wehub-resource-sync
145 lines
5.1 KiB
Python
145 lines
5.1 KiB
Python
from typing import List, Tuple
|
|
|
|
import numpy as np
|
|
|
|
from ray.rllib.env.single_agent_episode import SingleAgentEpisode
|
|
from ray.util.annotations import DeveloperAPI
|
|
|
|
|
|
@DeveloperAPI
|
|
def add_one_ts_to_episodes_and_truncate(episodes: List[SingleAgentEpisode]):
|
|
"""Adds an artificial timestep to an episode at the end.
|
|
|
|
In detail: The last observations, infos, actions, and all `extra_model_outputs`
|
|
will be duplicated and appended to each episode's data. An extra 0.0 reward
|
|
will be appended to the episode's rewards. The episode's timestep will be
|
|
increased by 1. Also, adds the truncated=True flag to each episode if the
|
|
episode is not already done (terminated or truncated).
|
|
|
|
Useful for value function bootstrapping, where it is required to compute a
|
|
forward pass for the very last timestep within the episode,
|
|
i.e. using the following input dict: {
|
|
obs=[final obs],
|
|
state=[final state output],
|
|
prev. reward=[final reward],
|
|
etc..
|
|
}
|
|
|
|
Args:
|
|
episodes: The list of SingleAgentEpisode objects to extend by one timestep
|
|
and add a truncation flag if necessary.
|
|
|
|
Returns:
|
|
A list of the original episodes' truncated values (so the episodes can be
|
|
properly restored later into their original states).
|
|
"""
|
|
orig_truncateds = []
|
|
for episode in episodes:
|
|
orig_truncateds.append(episode.is_truncated)
|
|
|
|
# Add timestep.
|
|
episode.t += 1
|
|
# Use the episode API that allows appending (possibly complex) structs
|
|
# to the data.
|
|
episode.observations.append(episode.observations[-1])
|
|
episode.infos.append(episode.infos[-1])
|
|
episode.actions.append(episode.actions[-1])
|
|
episode.rewards.append(0.0)
|
|
for v in episode.extra_model_outputs.values():
|
|
v.append(v[-1])
|
|
# Artificially make this episode truncated for the upcoming GAE
|
|
# computations.
|
|
if not episode.is_done:
|
|
episode.is_truncated = True
|
|
# Validate to make sure, everything is in order.
|
|
episode.validate()
|
|
|
|
return orig_truncateds
|
|
|
|
|
|
@DeveloperAPI
|
|
def remove_last_ts_from_data(
|
|
episode_lens: List[int],
|
|
*data: Tuple[np._typing.NDArray],
|
|
) -> Tuple[np._typing.NDArray]:
|
|
"""Removes the last timesteps from each given data item.
|
|
|
|
Each item in data is a concatenated sequence of episodes data.
|
|
For example if `episode_lens` is [2, 4], then data is a shape=(6,)
|
|
ndarray. The returned corresponding value will have shape (4,), meaning
|
|
both episodes have been shortened by exactly one timestep to 1 and 3.
|
|
|
|
..testcode::
|
|
|
|
from ray.rllib.algorithms.ppo.ppo_learner import PPOLearner
|
|
import numpy as np
|
|
|
|
unpadded = PPOLearner._remove_last_ts_from_data(
|
|
[5, 3],
|
|
np.array([0, 1, 2, 3, 4, 0, 1, 2]),
|
|
)
|
|
assert (unpadded[0] == [0, 1, 2, 3, 0, 1]).all()
|
|
|
|
unpadded = PPOLearner._remove_last_ts_from_data(
|
|
[4, 2, 3],
|
|
np.array([0, 1, 2, 3, 0, 1, 0, 1, 2]),
|
|
np.array([4, 5, 6, 7, 2, 3, 3, 4, 5]),
|
|
)
|
|
assert (unpadded[0] == [0, 1, 2, 0, 0, 1]).all()
|
|
assert (unpadded[1] == [4, 5, 6, 2, 3, 4]).all()
|
|
|
|
Args:
|
|
episode_lens: A list of current episode lengths. The returned
|
|
data will have the same lengths minus 1 timestep.
|
|
data: A tuple of data items (np.ndarrays) representing concatenated episodes
|
|
to be shortened by one timestep per episode.
|
|
Note that only arrays with `shape=(n,)` are supported! The
|
|
returned data will have `shape=(n-len(episode_lens),)` (each
|
|
episode gets shortened by one timestep).
|
|
|
|
Returns:
|
|
A tuple of new data items shortened by one timestep.
|
|
"""
|
|
# Figure out the new slices to apply to each data item based on
|
|
# the given episode_lens.
|
|
slices = []
|
|
sum = 0
|
|
for len_ in episode_lens:
|
|
slices.append(slice(sum, sum + len_ - 1))
|
|
sum += len_
|
|
# Compiling return data by slicing off one timestep at the end of
|
|
# each episode.
|
|
ret = []
|
|
for d in data:
|
|
ret.append(np.concatenate([d[s] for s in slices]))
|
|
return tuple(ret) if len(ret) > 1 else ret[0]
|
|
|
|
|
|
@DeveloperAPI
|
|
def remove_last_ts_from_episodes_and_restore_truncateds(
|
|
episodes: List[SingleAgentEpisode],
|
|
orig_truncateds: List[bool],
|
|
) -> None:
|
|
"""Reverts the effects of `_add_ts_to_episodes_and_truncate`.
|
|
|
|
Args:
|
|
episodes: The list of SingleAgentEpisode objects to extend by one timestep
|
|
and add a truncation flag if necessary.
|
|
orig_truncateds: A list of the original episodes' truncated values to be
|
|
applied to the `episodes`.
|
|
"""
|
|
|
|
# Fix all episodes.
|
|
for episode, orig_truncated in zip(episodes, orig_truncateds):
|
|
# Reduce timesteps by 1.
|
|
episode.t -= 1
|
|
# Remove all extra timestep data from the episode's buffers.
|
|
episode.observations.pop()
|
|
episode.infos.pop()
|
|
episode.actions.pop()
|
|
episode.rewards.pop()
|
|
for v in episode.extra_model_outputs.values():
|
|
v.pop()
|
|
# Fix the truncateds flag again.
|
|
episode.is_truncated = orig_truncated
|