def test_e2e_no_models(ego_cat_dataset: EgoDataset,
simulation_evaluator: ClosedLoopEvaluator) -> None:
sim_cfg = SimulationConfig(use_ego_gt=True,
use_agents_gt=True,
disable_new_agents=False,
distance_th_far=1,
distance_th_close=0)
sim_loop = ClosedLoopSimulator(sim_cfg, ego_cat_dataset,
torch.device("cpu"))
sim_out = sim_loop.unroll(list(range(len(ego_cat_dataset.dataset.scenes))))
simulation_evaluator.evaluate(sim_out)
agg = ValidationCountingAggregator().aggregate(
simulation_evaluator.validation_results())
assert agg["displacement_error_l2_validator"].item() == 0
assert agg["distance_ref_trajectory_validator"].item() == 0
assert agg["collision_front_validator"].item() == 0
assert agg["collision_rear_validator"].item() == 0
assert agg["collision_side_validator"].item() == 0
def test_e2e_both(ego_cat_dataset: EgoDataset,
simulation_evaluator: ClosedLoopEvaluator,
advance_x_ego: float, advance_x_agent: float) -> None:
sim_cfg = SimulationConfig(
use_ego_gt=False,
use_agents_gt=False,
disable_new_agents=False,
distance_th_far=60,
distance_th_close=30,
num_simulation_steps=20,
start_frame_index=10) # start from 0 as leading is rotated at 0
ego_model = MockModel(advance_x=advance_x_ego)
agents_model = MockModel(advance_x=advance_x_agent)
sim_loop = ClosedLoopSimulator(sim_cfg,
ego_cat_dataset,
torch.device("cpu"),
model_ego=ego_model,
model_agents=agents_model)
sim_out = sim_loop.unroll(list(range(len(ego_cat_dataset.dataset.scenes))))
simulation_evaluator.evaluate(sim_out)
validation_results = simulation_evaluator.validation_results()
agg = ValidationCountingAggregator().aggregate(validation_results)
validators = simulation_evaluator.evaluation_plan.validators_dict().keys()
if advance_x_ego == 2.0 and advance_x_agent == 2.0: # distance ref trajectory
validator_to_trigger = "distance_ref_trajectory_validator"
else:
raise ValueError(
f"advance_x_ego {advance_x_ego} and advance_x_agents {advance_x_agent} not valid"
)
for validator in validators:
if validator == validator_to_trigger:
assert agg[validator].item() == 4
else:
assert agg[validator].item() == 0
def test_e2e_ego(ego_cat_dataset: EgoDataset,
simulation_evaluator: ClosedLoopEvaluator,
advance_x: float) -> None:
sim_cfg = SimulationConfig(use_ego_gt=False,
use_agents_gt=True,
disable_new_agents=False,
distance_th_far=1,
distance_th_close=0,
num_simulation_steps=20)
ego_model = MockModel(advance_x=advance_x)
sim_loop = ClosedLoopSimulator(sim_cfg,
ego_cat_dataset,
torch.device("cpu"),
model_ego=ego_model)
sim_out = sim_loop.unroll(list(range(len(ego_cat_dataset.dataset.scenes))))
simulation_evaluator.evaluate(sim_out)
validation_results = simulation_evaluator.validation_results()
agg = ValidationCountingAggregator().aggregate(validation_results)
validators = simulation_evaluator.evaluation_plan.validators_dict().keys()
if advance_x == 0.0: # bumps by rear car
validator_to_trigger = "collision_rear_validator"
elif advance_x == 2.0: # too fast
validator_to_trigger = "distance_ref_trajectory_validator"
elif advance_x == 5.0: # too fast bumps into leading car
validator_to_trigger = "collision_front_validator"
else:
raise ValueError(f"advance_x {advance_x} not valid")
for validator in validators:
if validator == validator_to_trigger:
assert agg[validator].item() == 4
else:
assert agg[validator].item() == 0
def test_unroll(zarr_cat_dataset: ChunkedDataset, dmg: LocalDataManager,
cfg: dict) -> None:
rasterizer = build_rasterizer(cfg, dmg)
# change the first yaw of scene 1
# this will break if some broadcasting happens
zarr_cat_dataset.frames = np.asarray(zarr_cat_dataset.frames)
slice_frames = get_frames_slice_from_scenes(zarr_cat_dataset.scenes[1])
rot = zarr_cat_dataset.frames[slice_frames.start]["ego_rotation"].copy()
zarr_cat_dataset.frames[
slice_frames.start]["ego_rotation"] = yaw_as_rotation33(
rotation33_as_yaw(rot + 0.75))
scene_indices = list(range(len(zarr_cat_dataset.scenes)))
ego_dataset = EgoDataset(cfg, zarr_cat_dataset, rasterizer)
# control only agents at T0, control them forever
sim_cfg = SimulationConfig(use_ego_gt=False,
use_agents_gt=False,
disable_new_agents=True,
distance_th_close=1000,
distance_th_far=1000,
num_simulation_steps=10)
# ego will move by 1 each time
ego_model = MockModel(advance_x=1.0)
# agents will move by 0.5 each time
agents_model = MockModel(advance_x=0.5)
sim = ClosedLoopSimulator(sim_cfg, ego_dataset, torch.device("cpu"),
ego_model, agents_model)
sim_outputs = sim.unroll(scene_indices)
# check ego movement
for sim_output in sim_outputs:
ego_tr = sim_output.simulated_ego[
"ego_translation"][:sim_cfg.num_simulation_steps, :2]
ego_dist = np.linalg.norm(np.diff(ego_tr, axis=0), axis=-1)
assert np.allclose(ego_dist, 1.0)
ego_tr = sim_output.simulated_ego_states[:sim_cfg.num_simulation_steps,
TrajectoryStateIndices.
X:TrajectoryStateIndices.Y +
1]
ego_dist = np.linalg.norm(np.diff(ego_tr.numpy(), axis=0), axis=-1)
assert np.allclose(ego_dist, 1.0, atol=1e-3)
# all rotations should be the same as the first one as the MockModel outputs 0 for that
rots_sim = sim_output.simulated_ego[
"ego_rotation"][:sim_cfg.num_simulation_steps]
r_rep = sim_output.recorded_ego["ego_rotation"][0]
for r_sim in rots_sim:
assert np.allclose(rotation33_as_yaw(r_sim),
rotation33_as_yaw(r_rep),
atol=1e-2)
# all rotations should be the same as the first one as the MockModel outputs 0 for that
rots_sim = sim_output.simulated_ego_states[:sim_cfg.
num_simulation_steps,
TrajectoryStateIndices.
THETA]
r_rep = sim_output.recorded_ego_states[0, TrajectoryStateIndices.THETA]
for r_sim in rots_sim:
assert np.allclose(r_sim, r_rep, atol=1e-2)
# check agents movements
for sim_output in sim_outputs:
# we need to know which agents were controlled during simulation
# TODO: this is not ideal, we should keep track of them through the simulation
sim_dataset = SimulationDataset.from_dataset_indices(
ego_dataset, [sim_output.scene_id], sim_cfg)
sim_dataset.rasterise_agents_frame_batch(
0) # this will fill agents_tracked
agents_tracks = [el[1] for el in sim_dataset._agents_tracked]
for track_id in agents_tracks:
states = sim_output.simulated_agents
agents = filter_agents_by_track_id(
states, track_id)[:sim_cfg.num_simulation_steps]
agent_dist = np.linalg.norm(np.diff(agents["centroid"], axis=0),
axis=-1)
assert np.allclose(agent_dist, 0.5)
def test_unroll_subset(zarr_cat_dataset: ChunkedDataset, dmg: LocalDataManager,
cfg: dict, frame_range: Tuple[int, int]) -> None:
rasterizer = build_rasterizer(cfg, dmg)
scene_indices = list(range(len(zarr_cat_dataset.scenes)))
ego_dataset = EgoDataset(cfg, zarr_cat_dataset, rasterizer)
# control only agents at T0, control them forever
sim_cfg = SimulationConfig(use_ego_gt=False,
use_agents_gt=False,
disable_new_agents=True,
distance_th_close=1000,
distance_th_far=1000,
num_simulation_steps=frame_range[1],
start_frame_index=frame_range[0])
# ego will move by 1 each time
ego_model = MockModel(advance_x=1.0)
# agents will move by 0.5 each time
agents_model = MockModel(advance_x=0.5)
sim = ClosedLoopSimulator(sim_cfg, ego_dataset, torch.device("cpu"),
ego_model, agents_model)
sim_outputs = sim.unroll(scene_indices)
for sim_out in sim_outputs:
assert zarr_cat_dataset.frames[
0] != sim_out.recorded_dataset.dataset.frames[0]
assert zarr_cat_dataset.frames[
0] != sim_out.simulated_dataset.dataset.frames[0]
for ego_in_out in sim_out.ego_ins_outs:
assert "positions" in ego_in_out.outputs and "yaws" in ego_in_out.outputs
assert np.allclose(ego_in_out.outputs["positions"][:, 0], 1.0)
assert np.allclose(ego_in_out.outputs["positions"][:, 1], 0.0)
for agents_in_out in sim_out.agents_ins_outs:
for agent_in_out in agents_in_out:
assert "positions" in agent_in_out.outputs and "yaws" in agent_in_out.outputs
assert np.allclose(agent_in_out.outputs["positions"][:, 0],
0.5)
assert np.allclose(agent_in_out.outputs["positions"][:, 1],
0.0)
if None not in frame_range:
assert len(
sim_out.recorded_dataset.dataset.frames) == frame_range[1]
assert len(
sim_out.simulated_dataset.dataset.frames) == frame_range[1]
assert len(sim_out.simulated_ego_states) == frame_range[1]
assert len(sim_out.recorded_ego_states) == frame_range[1]
assert len(sim_out.recorded_ego) == frame_range[1]
assert len(sim_out.simulated_ego) == frame_range[1]
assert len(sim_out.ego_ins_outs) == len(
sim_out.agents_ins_outs) == frame_range[1]
ego_tr = sim_out.simulated_ego[
"ego_translation"][:sim_cfg.num_simulation_steps, :2]
ego_dist = np.linalg.norm(np.diff(ego_tr, axis=0), axis=-1)
assert np.allclose(ego_dist, 1.0)