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()