def _step_wait(
self
) -> Tuple[Dict[str, np.ndarray], np.ndarray, np.ndarray, Iterable[Dict[
Any, Any]]]:
"""
Wait for the step taken with step_async().
:return: ([int] or [float], [float], [bool], dict) observation, reward, done, information
"""
results = [remote.recv() for remote in self.remotes]
self.waiting = False
obs, rews, env_dones, infos, actor_dones, actor_ids, episode_stats, env_times = zip(
*results)
self._env_times = np.stack(env_times)
self._actor_dones = np.stack(actor_dones)
self._actor_ids = actor_ids
# collect episode statistics
for stat in episode_stats:
if stat is not None:
self.epoch_stats.receive(stat)
return stack_numpy_dict_list(obs), np.stack(rews), np.stack(
env_dones), infos
def stack(cls, records: List['SpacesRecord']) -> 'SpacesRecord':
"""Stack multiple records into a single spaces record. Useful for processing multiple records in a batch.
All the records should be in numpy and have the same structure of the spaces (i.e. come from the same
environment etc.).
:param records: Records to stack.
:return: Single stacked record, containing all the given records, and having the corresponding batch shape.
"""
assert len(set([
r.substep_key for r in records
])) == 1, "Cannot batch records for different sub-step keys."
assert len(set([
r.agent_id for r in records
])) == 1, "Cannot batch records for different agent ids."
stacked_record = SpacesRecord(
actor_id=records[0].actor_id,
observation=stack_numpy_dict_list([r.observation
for r in records]),
action=stack_numpy_dict_list([r.action for r in records]),
reward=np.stack([r.reward for r in records]),
done=np.stack([r.done for r in records]))
if records[0].next_observation:
stacked_record.next_observation = stack_numpy_dict_list(
[r.next_observation for r in records])
if records[0].logits:
stacked_record.logits = stack_torch_dict_list(
[r.logits for r in records])
stacked_record.batch_shape = [len(records)]
if records[0].batch_shape:
stacked_record.batch_shape += records[0].batch_shape
return stacked_record
def reset(self) -> Dict[str, np.ndarray]:
"""VectorEnv implementation"""
observations = []
for env in self.envs:
observations.append(env.reset())
# send the episode statistics of the environment collected before the reset()
self.epoch_stats.receive(
env.get_stats(LogStatsLevel.EPISODE).last_stats)
self._env_times = np.array([env.get_env_time() for env in self.envs])
self._actor_ids = [env.actor_id() for env in self.envs]
self._actor_dones = np.hstack(
[env.is_actor_done() for env in self.envs])
return stack_numpy_dict_list(observations)
def compute_action(self,
observation: ObservationType,
maze_state: Optional[MazeStateType],
env: Optional[BaseEnv] = None,
actor_id: Optional[ActorID] = None,
deterministic: bool = False) -> ActionType:
"""Sample multiple actions together."""
if actor_id:
action_space = self.action_spaces_dict[actor_id.step_key]
else:
assert len(self.action_spaces_dict) == 1, "action spaces for multiple sub-steps are available, please " \
"specify actor ID explicitly"
action_space = list(self.action_spaces_dict.values())[0]
return stack_numpy_dict_list(
[action_space.sample() for _ in range(self.concurrency)])
def reset(self) -> Dict[str, np.ndarray]:
"""VectorEnv implementation"""
for remote in self.remotes:
remote.send(('reset', None))
results = [remote.recv() for remote in self.remotes]
obs, actor_dones, actor_ids, episode_stats, env_times = zip(*results)
self._env_times = np.stack(env_times)
self._actor_dones = np.stack(actor_dones)
self._actor_ids = actor_ids
# collect episode statistics
for stat in episode_stats:
if stat is not None:
self.epoch_stats.receive(stat)
return stack_numpy_dict_list(obs)
def step(
self, actions: ActionType
) -> Tuple[ObservationType, np.ndarray, np.ndarray, Iterable[Dict[Any,
Any]]]:
"""Step the environments with the given actions.
:param actions: the list of actions for the respective envs.
:return: observations, rewards, dones, information-dicts all in env-aggregated form.
"""
actions = unstack_numpy_list_dict(actions)
observations, rewards, env_dones, infos, actor_dones, actor_ids = [], [], [], [], [], []
for i, env in enumerate(self.envs):
o, r, env_done, i = env.step(actions[i])
actor_dones.append(env.is_actor_done())
actor_ids.append(env.actor_id())
if env_done:
o = env.reset()
# collect the episode statistics for finished environments
self.epoch_stats.receive(
env.get_stats(LogStatsLevel.EPISODE).last_stats)
observations.append(o)
rewards.append(r)
env_dones.append(env_done)
infos.append(i)
obs = stack_numpy_dict_list(observations)
rewards = np.hstack(rewards).astype(np.float32)
env_dones = np.hstack(env_dones)
self._env_times = np.array([env.get_env_time() for env in self.envs])
self._actor_dones = np.hstack(actor_dones)
self._actor_ids = actor_ids
return obs, rewards, env_dones, infos
def test_torch_conversion():
stacked = convert_to_torch(stack_numpy_dict_list(unstacked_example()), in_place=True, cast=None, device="cpu")
for k, v in convert_to_torch(stacked_example(), in_place=True, cast=None, device="cpu").items():
assert torch.all(v == stacked[k])
def test_numpy_dict_list_stacking():
stacked = stack_numpy_dict_list(unstacked_example())
for k, v in stacked_example().items():
assert np.all(v == stacked[k])