def test_simulation_dataset_build(zarr_cat_dataset: ChunkedDataset,
dmg: LocalDataManager, cfg: dict,
tmp_path: Path) -> None:
# modify one frame to ensure everything works also when scenes are different
zarr_cat_dataset.frames = np.asarray(zarr_cat_dataset.frames)
for scene_idx in range(len(zarr_cat_dataset.scenes)):
frame_slice = get_frames_slice_from_scenes(zarr_cat_dataset.scenes)
zarr_cat_dataset.frames[
frame_slice.start]["ego_translation"] += np.random.randn(3)
rasterizer = build_rasterizer(cfg, dmg)
ego_dataset = EgoDataset(cfg, zarr_cat_dataset, rasterizer)
sim_cfg = SimulationConfig(use_ego_gt=True,
use_agents_gt=True,
disable_new_agents=False,
distance_th_far=30,
distance_th_close=10)
# we should be able to create the same object by using both constructor and factory
scene_indices = list(range(len(zarr_cat_dataset.scenes)))
scene_dataset_batch: Dict[int, EgoDataset] = {}
for scene_idx in scene_indices:
scene_dataset = ego_dataset.get_scene_dataset(scene_idx)
scene_dataset_batch[scene_idx] = scene_dataset
sim_1 = SimulationDataset(scene_dataset_batch, sim_cfg)
sim_2 = SimulationDataset.from_dataset_indices(ego_dataset, scene_indices,
sim_cfg)
for (k_1, v_1), (k_2, v_2) in zip(sim_1.scene_dataset_batch.items(),
sim_2.scene_dataset_batch.items()):
assert k_1 == k_2
assert np.allclose(v_1.dataset.frames["ego_translation"],
v_2.dataset.frames["ego_translation"])
def update_agents(dataset: SimulationDataset, frame_idx: int,
input_dict: Dict[str, np.ndarray],
output_dict: Dict[str, np.ndarray]) -> None:
"""Update the agents in frame_idx (across scenes) using agents_output_dict
:param dataset: the simulation dataset
:param frame_idx: index of the frame to modify
:param input_dict: the input to the agent model
:param output_dict: the output of the agent model
:return:
"""
agents_update_dict: Dict[Tuple[int, int], np.ndarray] = {}
world_from_agent = input_dict["world_from_agent"]
yaw = input_dict["yaw"]
pred_trs = transform_points(output_dict["positions"][:, :1],
world_from_agent)[:, 0]
pred_yaws = yaw + output_dict["yaws"][:, 0, 0]
next_agents = np.zeros(len(yaw), dtype=AGENT_DTYPE)
next_agents["centroid"] = pred_trs
next_agents["yaw"] = pred_yaws
next_agents["track_id"] = input_dict["track_id"]
next_agents["extent"] = input_dict["extent"]
next_agents["label_probabilities"][:, PERCEPTION_LABEL_TO_INDEX[
"PERCEPTION_LABEL_CAR"]] = 1
for scene_idx, next_agent in zip(input_dict["scene_index"],
next_agents):
agents_update_dict[(scene_idx,
next_agent["track_id"])] = np.expand_dims(
next_agent, 0)
dataset.set_agents(frame_idx, agents_update_dict)
def calculate_non_kinematic_rescale_params(sim_dataset: SimulationDataset) -> NonKinematicActionRescaleParams:
"""Calculate the action un-normalization parameters from the simulation dataset for non-kinematic model.
:param sim_dataset: the input dataset to calculate the action rescale parameters
:return: the unnormalized action
"""
x_component_frames = []
y_component_frames = []
yaw_component_frames = []
for index in range(1, len(sim_dataset) - 1):
ego_input = sim_dataset.rasterise_frame_batch(index)
x_component_frames.append([scene['target_positions'][0, 0] for scene in ego_input])
y_component_frames.append([scene['target_positions'][0, 1] for scene in ego_input])
yaw_component_frames.append([scene['target_yaws'][0, 0] for scene in ego_input])
x_components = np.concatenate(x_component_frames)
y_components = np.concatenate(y_component_frames)
yaw_components = np.concatenate(yaw_component_frames)
x_mu, x_std = np.mean(x_components), np.std(x_components)
y_mu, y_std = np.mean(y_components), np.std(y_components)
yaw_mu, yaw_std = np.mean(yaw_components), np.std(yaw_components)
# Keeping scale = 10 * std so that extreme values are not clipped
return NonKinematicActionRescaleParams(x_mu, 10 * x_std, y_mu, 10 * y_std, yaw_mu, 10 * yaw_std)
def update_ego(dataset: SimulationDataset, frame_idx: int, input_dict: Dict[str, np.ndarray],
output_dict: Dict[str, np.ndarray]) -> None:
"""Update ego across scenes for the given frame index.
:param dataset: The simulation dataset
:param frame_idx: index of the frame to modify
:param input_dict: the input to the ego model
:param output_dict: the output of the ego model
:return:
"""
world_from_agent = input_dict["world_from_agent"]
yaw = input_dict["yaw"]
pred_trs = transform_points(output_dict["positions"][:, :1], world_from_agent)
pred_yaws = np.expand_dims(yaw, -1) + output_dict["yaws"][:, :1, 0]
dataset.set_ego(frame_idx, 0, pred_trs, pred_yaws)
def test_invalid_simulation_dataset(zarr_cat_dataset: ChunkedDataset,
dmg: LocalDataManager, cfg: dict,
tmp_path: Path) -> None:
rasterizer = build_rasterizer(cfg, dmg)
scene_indices = [0, len(zarr_cat_dataset.scenes)]
ego_dataset = EgoDataset(cfg, zarr_cat_dataset, rasterizer)
sim_cfg = SimulationConfig(use_ego_gt=True,
use_agents_gt=True,
disable_new_agents=False,
distance_th_far=30,
distance_th_close=10)
with pytest.raises(ValueError):
SimulationDataset.from_dataset_indices(ego_dataset, scene_indices,
sim_cfg)
def test_simulation_agents_mock_insert(dmg: LocalDataManager, cfg: dict,
tmp_path: Path) -> None:
zarr_dataset = _mock_dataset()
rasterizer = build_rasterizer(cfg, dmg)
ego_dataset = EgoDataset(cfg, zarr_dataset, rasterizer)
sim_cfg = SimulationConfig(use_ego_gt=True,
use_agents_gt=True,
disable_new_agents=True,
distance_th_far=100,
distance_th_close=10)
dataset = SimulationDataset.from_dataset_indices(ego_dataset, [0], sim_cfg)
_ = dataset.rasterise_agents_frame_batch(0)
# insert (0, 1) in following frames
next_agent = np.zeros(1, dtype=AGENT_DTYPE)
next_agent["centroid"] = (-1, -1)
next_agent["yaw"] = -0.5
next_agent["track_id"] = 1
next_agent["extent"] = (1, 1, 1)
next_agent["label_probabilities"][:, 3] = 1
for frame_idx in [1, 2, 3]:
dataset.set_agents(frame_idx, {(0, 1): next_agent})
agents_dict = dataset.rasterise_agents_frame_batch(frame_idx)
assert len(agents_dict) == 1 and (0, 1) in agents_dict
assert np.allclose(agents_dict[(0, 1)]["centroid"], (-1, -1))
assert np.allclose(agents_dict[(0, 1)]["yaw"], -0.5)
def test_simulation_agents_mock_disable(dmg: LocalDataManager, cfg: dict,
tmp_path: Path) -> None:
zarr_dataset = _mock_dataset()
rasterizer = build_rasterizer(cfg, dmg)
ego_dataset = EgoDataset(cfg, zarr_dataset, rasterizer)
sim_cfg = SimulationConfig(use_ego_gt=True,
use_agents_gt=True,
disable_new_agents=True,
distance_th_far=100,
distance_th_close=10)
dataset = SimulationDataset.from_dataset_indices(ego_dataset, [0], sim_cfg)
# nothing should be tracked
assert len(dataset._agents_tracked) == 0
agents_dict = dataset.rasterise_agents_frame_batch(0)
# only (0, 1) should be in
assert len(agents_dict) == 1 and (0, 1) in agents_dict
assert len(dataset._agents_tracked) == 1
agents_dict = dataset.rasterise_agents_frame_batch(1)
# again, only (0, 1) should be in
assert len(agents_dict) == 1
assert (0, 1) in agents_dict
assert len(dataset._agents_tracked) == 1
agents_dict = dataset.rasterise_agents_frame_batch(2)
assert len(agents_dict) == 0
assert len(dataset._agents_tracked) == 0
def test_simulation_agents(zarr_cat_dataset: ChunkedDataset,
dmg: LocalDataManager, cfg: dict,
tmp_path: Path) -> None:
rasterizer = build_rasterizer(cfg, dmg)
scene_indices = list(range(len(zarr_cat_dataset.scenes)))
ego_dataset = EgoDataset(cfg, zarr_cat_dataset, rasterizer)
sim_cfg = SimulationConfig(use_ego_gt=True,
use_agents_gt=True,
disable_new_agents=False,
distance_th_far=100,
distance_th_close=30)
dataset = SimulationDataset.from_dataset_indices(ego_dataset,
scene_indices, sim_cfg)
# nothing should be tracked
assert len(dataset._agents_tracked) == 0
agents_dict = dataset.rasterise_agents_frame_batch(0)
# we should have the same agents in each scene
for k in agents_dict:
assert (0, k[1]) in agents_dict
# now everything should be tracked
assert len(dataset._agents_tracked) == len(agents_dict)
def unroll(self, scene_indices: List[int]) -> List[SimulationOutput]:
"""
Simulate the dataset for the given scene indices
:param scene_indices: the scene indices we want to simulate
:return: the simulated dataset
"""
sim_dataset = SimulationDataset.from_dataset_indices(self.dataset, scene_indices, self.sim_cfg)
agents_ins_outs: DefaultDict[int, List[List[UnrollInputOutput]]] = defaultdict(list)
ego_ins_outs: DefaultDict[int, List[UnrollInputOutput]] = defaultdict(list)
for frame_index in tqdm(range(len(sim_dataset)), disable=not self.sim_cfg.show_info):
next_frame_index = frame_index + 1
should_update = next_frame_index != len(sim_dataset)
# AGENTS
if not self.sim_cfg.use_agents_gt:
agents_input = sim_dataset.rasterise_agents_frame_batch(frame_index)
if len(agents_input): # agents may not be available
agents_input_dict = default_collate(list(agents_input.values()))
agents_output_dict = self.model_agents(move_to_device(agents_input_dict, self.device))
# for update we need everything as numpy
agents_input_dict = move_to_numpy(agents_input_dict)
agents_output_dict = move_to_numpy(agents_output_dict)
if should_update:
self.update_agents(sim_dataset, next_frame_index, agents_input_dict, agents_output_dict)
# update input and output buffers
agents_frame_in_out = self.get_agents_in_out(agents_input_dict, agents_output_dict,
self.keys_to_exclude)
for scene_idx in scene_indices:
agents_ins_outs[scene_idx].append(agents_frame_in_out.get(scene_idx, []))
# EGO
if not self.sim_cfg.use_ego_gt:
ego_input = sim_dataset.rasterise_frame_batch(frame_index)
ego_input_dict = default_collate(ego_input)
ego_output_dict = self.model_ego(move_to_device(ego_input_dict, self.device))
ego_input_dict = move_to_numpy(ego_input_dict)
ego_output_dict = move_to_numpy(ego_output_dict)
if should_update:
self.update_ego(sim_dataset, next_frame_index, ego_input_dict, ego_output_dict)
ego_frame_in_out = self.get_ego_in_out(ego_input_dict, ego_output_dict, self.keys_to_exclude)
for scene_idx in scene_indices:
ego_ins_outs[scene_idx].append(ego_frame_in_out[scene_idx])
simulated_outputs: List[SimulationOutput] = []
for scene_idx in scene_indices:
simulated_outputs.append(SimulationOutput(scene_idx, sim_dataset, ego_ins_outs, agents_ins_outs))
return simulated_outputs
def reset(self) -> Dict[str, np.ndarray]:
""" Resets the environment and outputs first frame of a new scene sample.
:return: the observation of first frame of sampled scene index
"""
# Define in / outs for new episode scene
self.agents_ins_outs: DefaultDict[
int, List[List[UnrollInputOutput]]] = defaultdict(list)
self.ego_ins_outs: DefaultDict[
int, List[UnrollInputOutput]] = defaultdict(list)
# Select Scene ID
self.scene_index = self.np_random.randint(0, self.max_scene_id)
if self.reset_scene_id is not None:
self.scene_index = min(self.reset_scene_id, self.max_scene_id - 1)
self.reset_scene_id += 1
# Select Frame ID (within bounds of the scene)
if self.randomize_start_frame:
scene_length = len(self.dataset.get_scene_indices(
self.scene_index))
self.eps_length = self.sim_cfg.num_simulation_steps or scene_length
end_frame = scene_length - self.eps_length
self.sim_cfg.start_frame_index = self.np_random.randint(
0, end_frame + 1)
# Prepare episode scene
self.sim_dataset = SimulationDataset.from_dataset_indices(
self.dataset, [self.scene_index], self.sim_cfg)
# Reset CLE evaluator
self.reward.reset()
# Output first observation
self.frame_index = 1 # Frame_index 1 has access to the true ego speed
ego_input = self.sim_dataset.rasterise_frame_batch(self.frame_index)
self.ego_input_dict = {
k: np.expand_dims(v, axis=0)
for k, v in ego_input[0].items()
}
# Reset Kinematic model
if self.use_kinematic:
init_kin_state = np.array(
[0.0, 0.0, 0.0, self.step_time * ego_input[0]['curr_speed']])
self.kin_model.reset(init_kin_state)
# Only output the image attribute
obs = {'image': ego_input[0]["image"]}
return obs
def test_visualise(ego_cat_dataset: EgoDataset) -> None:
mapAPI = ego_cat_dataset.rasterizer.sem_rast.mapAPI # type: ignore
sim_cfg = SimulationConfig(use_ego_gt=True,
use_agents_gt=True,
disable_new_agents=False,
distance_th_far=1,
distance_th_close=0,
num_simulation_steps=10)
sim_dataset = SimulationDataset.from_dataset_indices(
ego_cat_dataset, [0, 1], sim_cfg)
sim_out = SimulationOutput(0,
sim_dataset,
ego_ins_outs=defaultdict(list),
agents_ins_outs=defaultdict(list))
# ensure we can call the visualize
visualize(0, simulation_out_to_visualizer_scene(sim_out, mapAPI))
def _get_non_kin_rescale_params(self,
max_num_scenes: int = 10
) -> NonKinematicActionRescaleParams:
"""Determine the action un-normalization parameters for the non-kinematic model
from the current dataset in the L5Kit environment.
:param max_num_scenes: maximum number of scenes to consider to determine parameters
:return: Tuple of the action un-normalization parameters for non-kinematic model
"""
scene_ids = list(range(self.max_scene_id)) if not self.overfit else [
self.overfit_scene_id
]
if len(scene_ids) > max_num_scenes: # If too many scenes, CPU crashes
scene_ids = scene_ids[:max_num_scenes]
sim_dataset = SimulationDataset.from_dataset_indices(
self.dataset, scene_ids, self.sim_cfg)
return calculate_non_kinematic_rescale_params(sim_dataset)
def test_simulation_ego(zarr_cat_dataset: ChunkedDataset,
dmg: LocalDataManager, cfg: dict,
tmp_path: Path) -> None:
rasterizer = build_rasterizer(cfg, dmg)
scene_indices = list(range(len(zarr_cat_dataset.scenes)))
ego_dataset = EgoDataset(cfg, zarr_cat_dataset, rasterizer)
sim_cfg = SimulationConfig(use_ego_gt=True,
use_agents_gt=True,
disable_new_agents=False,
distance_th_far=30,
distance_th_close=10)
dataset = SimulationDataset.from_dataset_indices(ego_dataset,
scene_indices, sim_cfg)
# this also ensure order is checked
assert list(dataset.scene_dataset_batch.keys()) == scene_indices
# ensure we can call the aggregated get frame
out_0 = dataset.rasterise_frame_batch(0)
assert len(out_0) == len(scene_indices)
out_last = dataset.rasterise_frame_batch(len(dataset) - 1)
assert len(out_last) == len(scene_indices)
with pytest.raises(IndexError):
_ = dataset.rasterise_frame_batch(len(dataset))
# ensure we can set the ego in multiple frames for all scenes
frame_indices = np.random.randint(0, len(dataset), 10)
for frame_idx in frame_indices:
mock_tr = np.random.rand(len(scene_indices), 12, 2)
mock_yaw = np.random.rand(len(scene_indices), 12)
dataset.set_ego(frame_idx, 0, mock_tr, mock_yaw)
for scene_idx in scene_indices:
scene_zarr = dataset.scene_dataset_batch[scene_idx].dataset
ego_tr = scene_zarr.frames["ego_translation"][frame_idx]
ego_yaw = rotation33_as_yaw(
scene_zarr.frames["ego_rotation"][frame_idx])
assert np.allclose(mock_tr[scene_idx, 0], ego_tr[:2])
assert np.allclose(mock_yaw[scene_idx, 0], ego_yaw)