def step(self, model_states: StatesModel,
env_states: StatesEnv) -> StatesEnv:
"""
Set the environment to the target states by applying the specified \
actions an arbitrary number of time steps.
The state transitions will be calculated in parallel.
Args:
model_states: :class:`StatesModel` representing the data to be used \
to act on the environment.
env_states: :class:`StatesEnv` representing the data to be set in \
the environment.
Returns:
:class:`StatesEnv` containing the information that describes the \
new state of the Environment.
"""
split_env_states = [
env.step.remote(model_states=ms, env_states=es)
for env, ms, es in zip(
self.envs,
model_states.split_states(self.n_workers),
env_states.split_states(self.n_workers),
)
]
env_states = ray.get(split_env_states)
new_env_states: StatesEnv = StatesEnv.merge_states(env_states)
return new_env_states
def _make_transitions(self, model_states: StatesModel,
env_states: StatesEnv) -> List[StatesEnv]:
n_chunks = len(self._envs)
results = [
env.step(self._blocking, env_states=es, model_states=ms)
for env, es, ms in zip(self._envs, env_states.split_states(
n_chunks), model_states.split_states(n_chunks))
]
states = [result if self._blocking else result() for result in results]
return states