def test_unroll_none_input(ego_cat_dataset: EgoDataset) -> None: sim_cfg = SimulationConfig(use_ego_gt=True, use_agents_gt=False, disable_new_agents=True, distance_th_close=1000, distance_th_far=1000, num_simulation_steps=10) sim = ClosedLoopSimulator(sim_cfg, ego_cat_dataset, torch.device("cpu"), None, MockModel()) assert isinstance(sim.model_ego, torch.nn.Sequential) assert isinstance(sim.model_agents, MockModel) sim_cfg = SimulationConfig(use_ego_gt=False, use_agents_gt=True, disable_new_agents=True, distance_th_close=1000, distance_th_far=1000, num_simulation_steps=10) sim = ClosedLoopSimulator(sim_cfg, ego_cat_dataset, torch.device("cpu"), MockModel(), None) assert isinstance(sim.model_ego, MockModel) assert isinstance(sim.model_agents, torch.nn.Sequential) sim_cfg = SimulationConfig(use_ego_gt=True, use_agents_gt=True, disable_new_agents=True, distance_th_close=1000, distance_th_far=1000, num_simulation_steps=10) sim = ClosedLoopSimulator(sim_cfg, ego_cat_dataset, torch.device("cpu"), None, None) assert isinstance(sim.model_ego, torch.nn.Sequential) assert isinstance(sim.model_agents, torch.nn.Sequential)
def test_get_in_out_mock() -> None: with pytest.raises(ValueError): # repeated scene not allowed ClosedLoopSimulator.get_ego_in_out({"scene_index": np.ones(3)}, {"value": np.ones(3)}) input_ego_dict = { "scene_index": np.asarray([0, 1, 2]), "track_id": np.full(3, -1), "mock": np.zeros(3) } output_ego_dict = { "output_value": input_ego_dict["scene_index"], "mock_out": np.zeros(3) } out = ClosedLoopSimulator.get_ego_in_out(input_ego_dict, output_ego_dict) # we should have 3 elements (one per scene) assert len(out) == 3 for scene_index in input_ego_dict["scene_index"]: assert scene_index in out assert out[scene_index].track_id == -1 assert out[scene_index].outputs["output_value"] == scene_index # we can test with exclude_keys also no_keys = {"mock", "mock_out"} out = ClosedLoopSimulator.get_ego_in_out(input_ego_dict, output_ego_dict, keys_to_exclude=no_keys) assert len(out) == 3 for out_el in out.values(): assert len(no_keys.intersection(out_el.inputs.keys())) == 0 assert len(no_keys.intersection(out_el.outputs.keys())) == 0 # for agents repeated scenes are allowed input_agents_dict = { "scene_index": np.asarray([0, 0, 2]), "track_id": np.asarray([0, 1, 0]) } output_agents_dict = { "output_value_1": input_agents_dict["scene_index"], "output_value_2": input_agents_dict["track_id"] } out = ClosedLoopSimulator.get_agents_in_out(input_agents_dict, output_agents_dict) assert len(out[0]) == 2 assert len(out[2]) == 1 for scene_index, scene_out in out.items(): for agent_out in scene_out: assert agent_out.outputs["output_value_1"] == scene_index assert agent_out.outputs["output_value_2"] == agent_out.track_id
def test_unroll_invalid_input(ego_cat_dataset: EgoDataset) -> None: # try to use None models with wrong config 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) with pytest.raises(ValueError): ClosedLoopSimulator(sim_cfg, ego_cat_dataset, torch.device("cpu"), MockModel(), None) with pytest.raises(ValueError): ClosedLoopSimulator(sim_cfg, ego_cat_dataset, torch.device("cpu"), None, MockModel()) with pytest.raises(ValueError): ClosedLoopSimulator(sim_cfg, ego_cat_dataset, torch.device("cpu"), None, None)
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)
class L5Env(gym.Env): """Custom Environment of L5 Kit that can be registered in OpenAI Gym. :param env_config_path: path to the L5Kit environment configuration file :param dmg: local data manager object :param simulation_cfg: configuration of the L5Kit closed loop simulator :param train: flag to determine whether to use train or validation dataset :param reward: calculates the reward for the gym environment :param cle: flag to enable close loop environment updates :param rescale_action: flag to rescale the model action back to the un-normalized action space :param use_kinematic: flag to use the kinematic model :param kin_model: the kinematic model :param return_info: flag to return info when a episode ends :param randomize_start: flag to randomize the start frame of episode """ def __init__(self, env_config_path: Optional[str] = None, dmg: Optional[LocalDataManager] = None, sim_cfg: Optional[SimulationConfig] = None, train: bool = True, reward: Optional[Reward] = None, cle: bool = True, rescale_action: bool = True, use_kinematic: bool = False, kin_model: Optional[KinematicModel] = None, reset_scene_id: Optional[int] = None, return_info: bool = False, randomize_start: bool = True) -> None: """Constructor method """ super(L5Env, self).__init__() # Required to register environment if env_config_path is None: return # env config dm = dmg if dmg is not None else LocalDataManager(None) cfg = load_config_data(env_config_path) self.step_time = cfg["model_params"]["step_time"] # rasterisation rasterizer = build_rasterizer(cfg, dm) raster_size = cfg["raster_params"]["raster_size"][0] n_channels = rasterizer.num_channels() # load dataset of environment self.train = train self.overfit = cfg["gym_params"]["overfit"] self.randomize_start_frame = randomize_start if self.train or self.overfit: loader_key = cfg["train_data_loader"]["key"] else: loader_key = cfg["val_data_loader"]["key"] dataset_zarr = ChunkedDataset(dm.require(loader_key)).open() self.dataset = EgoDataset(cfg, dataset_zarr, rasterizer) # Define action and observation space # Continuous Action Space: gym.spaces.Box (X, Y, Yaw * number of future states) self.action_space = spaces.Box(low=-1, high=1, shape=(3, )) # Observation Space: gym.spaces.Dict (image: [n_channels, raster_size, raster_size]) obs_shape = (n_channels, raster_size, raster_size) self.observation_space = spaces.Dict({ 'image': spaces.Box(low=0, high=1, shape=obs_shape, dtype=np.float32) }) # Simulator Config within Gym self.sim_cfg = sim_cfg if sim_cfg is not None else SimulationConfigGym( ) self.simulator = ClosedLoopSimulator(self.sim_cfg, self.dataset, device=torch.device("cpu"), mode=ClosedLoopSimulatorModes.GYM) self.reward = reward if reward is not None else L2DisplacementYawReward( ) self.max_scene_id = cfg["gym_params"]["max_scene_id"] if not self.train: self.max_scene_id = cfg["gym_params"]["max_val_scene_id"] self.randomize_start_frame = False if self.overfit: self.overfit_scene_id = cfg["gym_params"]["overfit_id"] self.randomize_start_frame = False self.cle = cle self.rescale_action = rescale_action self.use_kinematic = use_kinematic if self.use_kinematic: self.kin_model = kin_model if kin_model is not None else UnicycleModel( ) self.kin_rescale = self._get_kin_rescale_params() else: self.non_kin_rescale = self._get_non_kin_rescale_params() # If not None, reset_scene_id is the scene_id that will be rolled out when reset is called self.reset_scene_id = reset_scene_id if self.overfit: self.reset_scene_id = self.overfit_scene_id # flag to decide whether to return any info at end of episode # helps to limit the IPC self.return_info = return_info self.seed() def seed(self, seed: int = None) -> List[int]: """Generate the random seed. :param seed: the seed integer :return: the output random seed """ self.np_random, seed = seeding.np_random(seed) # TODO : add a torch seed for future return [seed] def set_reset_id(self, reset_id: int = None) -> None: """Set the reset_id to unroll from specific scene_id. Useful during deterministic evaluation. :param reset_id: the scene_id to unroll """ self.reset_scene_id = reset_id 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 step(self, action: np.ndarray) -> GymStepOutput: """Inputs the action, updates the environment state and outputs the next frame. :param action: the action to perform on current state/frame :return: the namedTuple comprising the (next observation, reward, done, info) based on the current action """ frame_index = self.frame_index next_frame_index = frame_index + 1 episode_over = next_frame_index == (len(self.sim_dataset) - 1) # EGO if not self.sim_cfg.use_ego_gt: action = self._rescale_action(action) ego_output = self._convert_action_to_ego_output(action) self.ego_output_dict = ego_output if self.cle: # In closed loop training, the raster is updated according to predicted ego positions. self.simulator.update_ego(self.sim_dataset, next_frame_index, self.ego_input_dict, self.ego_output_dict) ego_frame_in_out = self.simulator.get_ego_in_out( self.ego_input_dict, self.ego_output_dict, self.simulator.keys_to_exclude) self.ego_ins_outs[self.scene_index].append( ego_frame_in_out[self.scene_index]) # generate simulated_outputs simulated_outputs = SimulationOutputCLE(self.scene_index, self.sim_dataset, self.ego_ins_outs, self.agents_ins_outs) # reward calculation reward = self.reward.get_reward(self.frame_index, [simulated_outputs]) # done is True when episode ends done = episode_over # Optionally we can pass additional info # We are using "info" to output rewards and simulated outputs (during evaluation) info: Dict[str, Any] info = { 'reward_tot': reward["total"], 'reward_dist': reward["distance"], 'reward_yaw': reward["yaw"] } if done and self.return_info: info = { "sim_outs": self.get_episode_outputs(), "reward_tot": reward["total"], "reward_dist": reward["distance"], "reward_yaw": reward["yaw"] } # Get next obs self.frame_index += 1 obs = self._get_obs(self.frame_index, episode_over) # return obs, reward, done, info return GymStepOutput(obs, reward["total"], done, info) def get_episode_outputs(self) -> List[EpisodeOutputGym]: """Generate and return the outputs at the end of the episode. :return: List of episode outputs """ episode_outputs = [ EpisodeOutputGym(self.scene_index, self.sim_dataset, self.ego_ins_outs, self.agents_ins_outs) ] return episode_outputs def render(self) -> None: """Render a frame during the simulation """ raise NotImplementedError def _get_obs(self, frame_index: int, episode_over: bool) -> Dict[str, np.ndarray]: """Get the observation corresponding to a given frame index in the scene. :param frame_index: the index of the frame which provides the observation :param episode_over: flag to determine if the episode is over :return: the observation corresponding to the frame index """ if episode_over: frame_index = 0 # Dummy final obs (when episode_over) ego_input = self.sim_dataset.rasterise_frame_batch(frame_index) self.ego_input_dict = { k: np.expand_dims(v, axis=0) for k, v in ego_input[0].items() } obs = {"image": ego_input[0]["image"]} return obs def _rescale_action(self, action: np.ndarray) -> np.ndarray: """Rescale the input action back to the un-normalized action space. PPO and related algorithms work well with normalized action spaces. The environment receives a normalized action and we un-normalize it back to the original action space for environment updates. :param action: the normalized action :return: the unnormalized action """ if self.rescale_action: if self.use_kinematic: action[0] = self.kin_rescale.steer_scale * action[0] action[1] = self.kin_rescale.acc_scale * action[1] else: action[ 0] = self.non_kin_rescale.x_mu + self.non_kin_rescale.x_scale * action[ 0] action[ 1] = self.non_kin_rescale.y_mu + self.non_kin_rescale.y_scale * action[ 1] action[ 2] = self.non_kin_rescale.yaw_mu + self.non_kin_rescale.yaw_scale * action[ 2] return action def _get_kin_rescale_params(self) -> KinematicActionRescaleParams: """Determine the action un-normalization parameters for the kinematic model from the current dataset in the L5Kit environment. :return: Tuple of the action un-normalization parameters for kinematic model """ global MAX_ACC, MAX_STEER return KinematicActionRescaleParams(MAX_STEER * self.step_time, MAX_ACC * self.step_time) 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 _convert_action_to_ego_output( self, action: np.ndarray) -> Dict[str, np.ndarray]: """Convert the input action into ego output format. :param action: the input action provided by policy :return: action in ego output format, a numpy dict with keys 'positions' and 'yaws' """ if self.use_kinematic: data_dict = self.kin_model.update(action[:2]) else: # [batch_size=1, num_steps, (X, Y, yaw)] data = action.reshape(1, 1, 3) pred_positions = data[:, :, :2] # [batch_size, num_steps, 1->(yaw)] pred_yaws = data[:, :, 2:3] data_dict = {"positions": pred_positions, "yaws": pred_yaws} return data_dict
def __init__(self, env_config_path: Optional[str] = None, dmg: Optional[LocalDataManager] = None, sim_cfg: Optional[SimulationConfig] = None, train: bool = True, reward: Optional[Reward] = None, cle: bool = True, rescale_action: bool = True, use_kinematic: bool = False, kin_model: Optional[KinematicModel] = None, reset_scene_id: Optional[int] = None, return_info: bool = False, randomize_start: bool = True) -> None: """Constructor method """ super(L5Env, self).__init__() # Required to register environment if env_config_path is None: return # env config dm = dmg if dmg is not None else LocalDataManager(None) cfg = load_config_data(env_config_path) self.step_time = cfg["model_params"]["step_time"] # rasterisation rasterizer = build_rasterizer(cfg, dm) raster_size = cfg["raster_params"]["raster_size"][0] n_channels = rasterizer.num_channels() # load dataset of environment self.train = train self.overfit = cfg["gym_params"]["overfit"] self.randomize_start_frame = randomize_start if self.train or self.overfit: loader_key = cfg["train_data_loader"]["key"] else: loader_key = cfg["val_data_loader"]["key"] dataset_zarr = ChunkedDataset(dm.require(loader_key)).open() self.dataset = EgoDataset(cfg, dataset_zarr, rasterizer) # Define action and observation space # Continuous Action Space: gym.spaces.Box (X, Y, Yaw * number of future states) self.action_space = spaces.Box(low=-1, high=1, shape=(3, )) # Observation Space: gym.spaces.Dict (image: [n_channels, raster_size, raster_size]) obs_shape = (n_channels, raster_size, raster_size) self.observation_space = spaces.Dict({ 'image': spaces.Box(low=0, high=1, shape=obs_shape, dtype=np.float32) }) # Simulator Config within Gym self.sim_cfg = sim_cfg if sim_cfg is not None else SimulationConfigGym( ) self.simulator = ClosedLoopSimulator(self.sim_cfg, self.dataset, device=torch.device("cpu"), mode=ClosedLoopSimulatorModes.GYM) self.reward = reward if reward is not None else L2DisplacementYawReward( ) self.max_scene_id = cfg["gym_params"]["max_scene_id"] if not self.train: self.max_scene_id = cfg["gym_params"]["max_val_scene_id"] self.randomize_start_frame = False if self.overfit: self.overfit_scene_id = cfg["gym_params"]["overfit_id"] self.randomize_start_frame = False self.cle = cle self.rescale_action = rescale_action self.use_kinematic = use_kinematic if self.use_kinematic: self.kin_model = kin_model if kin_model is not None else UnicycleModel( ) self.kin_rescale = self._get_kin_rescale_params() else: self.non_kin_rescale = self._get_non_kin_rescale_params() # If not None, reset_scene_id is the scene_id that will be rolled out when reset is called self.reset_scene_id = reset_scene_id if self.overfit: self.reset_scene_id = self.overfit_scene_id # flag to decide whether to return any info at end of episode # helps to limit the IPC self.return_info = return_info self.seed()