"""evaluate_prediction_planning_stack.py --load_from <wandb ID> --seed <seed>
--num_episodes <num_episodes> --risk_level <a list of risk-levels>
--num_samples <a list of numbers of prediction samples>
This script loads a trained predictor from <wand ID>, runs a batch of open-loop MPC evaluations
(i.e., without replanning) using <num_episodes> episodes while varying risk-levels and numbers of
prediction samples. Results are stored in scripts/logs/planner_eval/run-<wandb ID>_<seed> as a
collection of pickle files.
"""
import argparse
import os
import pickle
from time import time_ns
from typing import List, Tuple
import sys
import torch
from mmcv import Config
from pytorch_lightning import seed_everything
from tqdm import trange
from risk_biased.mpc_planner.planner import MPCPlanner, MPCPlannerParams
from risk_biased.predictors.biased_predictor import LitTrajectoryPredictor
from risk_biased.scene_dataset.loaders import SceneDataLoaders
from risk_biased.scene_dataset.scene import RandomScene, RandomSceneParams
from risk_biased.utils.callbacks import get_fast_slow_scenes
from risk_biased.utils.load_model import load_from_config, config_argparse
from risk_biased.utils.planner_utils import (
evaluate_control_sequence,
get_interaction_cost,
AbstractState,
to_state,
)
def evaluate_main(
load_from: str,
seed: int,
num_episodes: int,
risk_level_list: List[float],
num_prediction_samples_list: List[int],
):
print(f"Risk-sensitivity level(s) to test: {risk_level_list}")
print(f"Number(s) of prediction samples to test: {num_prediction_samples_list} ")
save_dir = os.path.join(
os.path.dirname(os.path.realpath(__file__)),
"logs",
"planner_eval",
f"run-{load_from}_{seed}",
)
if not os.path.exists(save_dir):
os.makedirs(save_dir)
cfg, planner = get_cfg_and_planner(load_from=load_from)
if not planner.solver.params.mean_warm_start == False:
print(
"switching to mean_warm_start = False for open-loop evaluation (i.e. without re-planning)"
)
planner.solver.params.mean_warm_start = False
for scene_type in [
"safer_slow",
"safer_fast",
]:
print(f"\nTesting {scene_type} scenes")
seed_everything(seed)
(
scene,
ado_state_history_batch,
ado_state_future_batch,
) = get_scene_and_ado_trajectory(
cfg, scene_type=scene_type, num_episodes=num_episodes
)
(
ego_state_history,
ego_state_target_trajectory,
) = get_ego_state_history_and_target_trajectory(cfg, scene)
for stack_risk_level in risk_level_list:
print(f" Risk_level: {stack_risk_level}")
file_name = f"{scene_type}_no_policy_opt_risk_level_{stack_risk_level}"
print(f" {file_name}")
stats_dict_no_policy_opt = evaluate_prediction_planning_stack(
planner,
ado_state_history_batch,
ado_state_future_batch,
ego_state_history,
ego_state_target_trajectory,
optimize_policy=False,
stack_risk_level=stack_risk_level,
risk_in_predictor=False,
)
with open(os.path.join(save_dir, file_name + ".pkl"), "wb") as f:
pickle.dump(stats_dict_no_policy_opt, f)
for num_prediction_samples in num_prediction_samples_list:
file_name = f"{scene_type}_{num_prediction_samples}_samples_risk_level_{stack_risk_level}"
if stack_risk_level == 0.0:
print(f" {file_name}")
stats_dict_risk_neutral = evaluate_prediction_planning_stack(
planner,
ado_state_history_batch,
ado_state_future_batch,
ego_state_history,
ego_state_target_trajectory,
optimize_policy=True,
stack_risk_level=stack_risk_level,
risk_in_predictor=False,
num_prediction_samples=num_prediction_samples,
)
with open(os.path.join(save_dir, file_name + ".pkl"), "wb") as f:
pickle.dump(stats_dict_risk_neutral, f)
else:
file_name_in_predictor = file_name + "_in_predictor"
print(f" {file_name_in_predictor}")
stats_dict_risk_in_predictor = evaluate_prediction_planning_stack(
planner,
ado_state_history_batch,
ado_state_future_batch,
ego_state_history,
ego_state_target_trajectory,
optimize_policy=True,
stack_risk_level=stack_risk_level,
risk_in_predictor=True,
num_prediction_samples=num_prediction_samples,
)
with open(
os.path.join(save_dir, file_name_in_predictor + ".pkl"), "wb"
) as f:
pickle.dump(stats_dict_risk_in_predictor, f)
file_name_in_planner = file_name + "_in_planner"
print(f" {file_name_in_planner}")
stats_dict_risk_in_planner = evaluate_prediction_planning_stack(
planner,
ado_state_history_batch,
ado_state_future_batch,
ego_state_history,
ego_state_target_trajectory,
optimize_policy=True,
stack_risk_level=stack_risk_level,
risk_in_predictor=False,
num_prediction_samples=num_prediction_samples,
)
with open(
os.path.join(save_dir, file_name_in_planner + ".pkl"), "wb"
) as f:
pickle.dump(stats_dict_risk_in_planner, f)
def evaluate_prediction_planning_stack(
planner: MPCPlanner,
ado_state_history_batch: AbstractState,
ado_state_future_batch: AbstractState,
ego_state_history: AbstractState,
ego_state_target_trajectory: AbstractState,
optimize_policy: bool = True,
stack_risk_level: float = 0.0,
risk_in_predictor: bool = False,
num_prediction_samples: int = 128,
num_prediction_samples_for_policy_eval: int = 4096,
) -> dict:
assert planner.solver.params.mean_warm_start == False
if risk_in_predictor:
predictor_risk_level, planner_risk_level = stack_risk_level, 0.0
else:
predictor_risk_level, planner_risk_level = 0.0, stack_risk_level
stats_dict = {
"stack_risk_level": stack_risk_level,
"predictor_risk_level": predictor_risk_level,
"planner_risk_level": planner_risk_level,
}
num_episodes = ado_state_history_batch.shape[0]
assert num_episodes == ado_state_future_batch.shape[0]
for episode_id in trange(num_episodes, desc="episodes", leave=False):
ado_state_history = ado_state_history_batch[episode_id]
ado_state_future = ado_state_future_batch[episode_id]
(
ado_state_future_samples_for_policy_eval,
sample_weights,
) = planner.solver.sample_prediction(
planner.predictor,
ado_state_history,
planner.normalizer,
ego_state_history=ego_state_history,
ego_state_future=ego_state_target_trajectory,
num_prediction_samples=num_prediction_samples_for_policy_eval,
risk_level=0.0,
)
if optimize_policy:
start = time_ns()
solver_info = planner.solver.solve(
planner.predictor,
ego_state_history,
ego_state_target_trajectory,
ado_state_history,
planner.normalizer,
num_prediction_samples=num_prediction_samples,
verbose=False,
risk_level=stack_risk_level,
resample_prediction=False,
risk_in_predictor=risk_in_predictor,
)
end = time_ns()
computation_time_ms = (end - start) * 1e-6
else:
planner.solver.reset()
computation_time_ms = 0.0
solver_info = None
interaction_cost_gt = get_ground_truth_interaction_cost(
planner, ado_state_future, ego_state_history
)
interaction_risk, tracking_cost = evaluate_control_sequence(
planner.solver.control_sequence,
planner.solver.dynamics_model,
ego_state_history,
ego_state_target_trajectory,
ado_state_future_samples_for_policy_eval,
sample_weights,
planner.solver.interaction_cost,
planner.solver.tracking_cost,
risk_level=stack_risk_level,
risk_estimator=planner.solver.risk_estimator,
)
stats_dict_this_run = {
"computation_time_ms": computation_time_ms,
"interaction_cost_ground_truth": interaction_cost_gt,
"interaction_risk": interaction_risk,
"tracking_cost": tracking_cost,
"control_sequence": planner.solver.control_sequence,
"solver_info": solver_info,
"ado_unbiased_predictions": ado_state_future_samples_for_policy_eval.position.detach()
.cpu()
.numpy(),
"sample_weights": sample_weights.detach().cpu().numpy(),
"ado_position_future": ado_state_future.position.detach().cpu().numpy(),
"ado_position_history": ado_state_history.position.detach().cpu().numpy(),
}
stats_dict[episode_id] = stats_dict_this_run
return stats_dict
def get_cfg_and_predictor() -> Tuple[Config, LitTrajectoryPredictor]:
config_path = os.path.join(
os.path.dirname(os.path.realpath(__file__)),
"..",
"..",
"risk_biased",
"config",
"learning_config.py",
)
cfg = config_argparse(config_path)
predictor, _, cfg = load_from_config(cfg)
return cfg, predictor
def get_cfg_and_planner(load_from: str) -> Tuple[Config, MPCPlanner]:
planner_config_path = os.path.join(
os.path.dirname(os.path.realpath(__file__)),
"..",
"..",
"risk_biased",
"config",
"planning_config.py",
)
planner_cfg = Config.fromfile(planner_config_path)
cfg, predictor = get_cfg_and_predictor()
# joint_dict = {**dict(cfg), **dict(planner_cfg)}
# assert joint_dict == {
# **dict(planner_cfg),
# **dict(cfg),
# }, f"some of the entries conflict between {cfg.filename} and {planner_cfg.filename}"
# joint_cfg = Config(joint_dict)
cfg.update(planner_cfg)
planner_params = MPCPlannerParams.from_config(cfg)
normalizer = SceneDataLoaders.normalize_trajectory
planner = MPCPlanner(planner_params, predictor, normalizer)
return cfg, planner
def get_scene_and_ado_trajectory(
cfg: Config, scene_type: str, num_episodes: int
) -> Tuple[RandomScene, torch.Tensor, torch.Tensor]:
scene_params = RandomSceneParams.from_config(cfg)
safer_fast_scene, safer_slow_scene = get_fast_slow_scenes(
scene_params, num_episodes
)
assert scene_type in ["safer_fast", "safer_slow"]
if scene_type == "safer_fast":
scene = safer_fast_scene
elif scene_type == "safer_slow":
scene = safer_slow_scene
else:
raise ValueError(f"unknown scene type {scene_type}")
ado_trajectory = torch.from_numpy(
scene.get_pedestrians_trajectories().astype("float32")
)
ado_position_history = to_state(ado_trajectory[..., : cfg.num_steps, :], cfg.dt)
ado_position_future = to_state(ado_trajectory[..., cfg.num_steps :, :], cfg.dt)
return scene, ado_position_history, ado_position_future
def get_ego_state_history_and_target_trajectory(
cfg: Config, scene: RandomScene
) -> Tuple[torch.Tensor, torch.Tensor]:
ego_state_traj = to_state(
torch.from_numpy(
scene.get_ego_ref_trajectory(cfg.sample_times).astype("float32")
),
cfg.dt,
)
ego_state_history = ego_state_traj[0, :, : cfg.num_steps]
ego_state_target_trajectory = ego_state_traj[0, :, cfg.num_steps :]
return ego_state_history, ego_state_target_trajectory
def get_ground_truth_interaction_cost(
planner: MPCPlanner,
ado_state_future: AbstractState, # (num_agents, num_steps_future)
ego_state_history: AbstractState, # (1, 1, num_steps)
) -> float:
ego_state_future = planner.solver.dynamics_model.simulate(
ego_state_history[..., -1], planner.solver.control_sequence.unsqueeze(0)
)
interaction_cost = get_interaction_cost(
ego_state_future,
ado_state_future,
planner.solver.interaction_cost,
)
return interaction_cost.item()
if __name__ == "__main__":
parser = argparse.ArgumentParser(
description="evaluate prediciton-planning stack using safer_fast and safer_slow scenes"
)
parser.add_argument(
"--load_from",
type=str,
required=True,
help="WandB ID to load trained predictor from",
)
parser.add_argument(
"--seed",
type=int,
default=0,
help="Random seed for evaluation",
)
parser.add_argument(
"--num_episodes",
type=int,
default=100,
help="Number of episodes",
)
parser.add_argument(
"--risk_level",
type=float,
nargs="+",
help="Risk-sensitivity level(s) to test",
default=[0.95, 1],
)
parser.add_argument(
"--num_samples",
type=int,
nargs="+",
help="Number(s) of prediction samples to test",
default=[1, 4, 16, 64, 256, 1024],
)
parser.add_argument(
"--force_config",
action="store_true",
help="""Use this flag to force the use of the local config file
when loading a model from a checkpoint. Otherwise the checkpoint config file is used.
In any case the parameters can be overwritten with an argparse argument.""",
)
parser.add_argument(
"--load_last",
action="store_true",
help="""Use this flag to force the use of the last checkpoint instead of the best one
when loading a model.""",
)
args = parser.parse_args()
# Args will be re-parsed, this keeps only the arguments that are compatible with the second parser.
keep_list = ["--load_from", "--seed", "--load_last", "--force_config"]
sys.argv = [ss for s in sys.argv for ss in s.split("=")]
sys.argv = [
sys.argv[i]
for i in range(len(sys.argv))
if sys.argv[i] in keep_list or sys.argv[i - 1] in keep_list or i == 0
]
evaluate_main(
args.load_from, args.seed, args.num_episodes, args.risk_level, args.num_samples
)