wehub-resource-sync
150 lines
5.8 KiB
Python
150 lines
5.8 KiB
Python
from typing import Any, Dict, List, Optional
|
|
|
|
import gymnasium as gym
|
|
import numpy as np
|
|
import tree # pip install dm_tree
|
|
|
|
from ray.rllib.connectors.connector_v2 import ConnectorV2
|
|
from ray.rllib.core.columns import Columns
|
|
from ray.rllib.core.rl_module.rl_module import RLModule
|
|
from ray.rllib.utils.annotations import override
|
|
from ray.rllib.utils.typing import EpisodeType
|
|
from ray.util.annotations import PublicAPI
|
|
|
|
|
|
@PublicAPI(stability="alpha")
|
|
class FrameStacking(ConnectorV2):
|
|
"""A connector piece that stacks the previous n observations into one."""
|
|
|
|
@override(ConnectorV2)
|
|
def recompute_output_observation_space(
|
|
self,
|
|
input_observation_space: gym.Space,
|
|
input_action_space: gym.Space,
|
|
) -> gym.Space:
|
|
# Change our observation space according to the given stacking settings.
|
|
if self._multi_agent:
|
|
ret = {}
|
|
for agent_id, obs_space in input_observation_space.spaces.items():
|
|
ret[agent_id] = self._convert_individual_space(obs_space)
|
|
return gym.spaces.Dict(ret)
|
|
else:
|
|
return self._convert_individual_space(input_observation_space)
|
|
|
|
def __init__(
|
|
self,
|
|
input_observation_space: Optional[gym.Space] = None,
|
|
input_action_space: Optional[gym.Space] = None,
|
|
*,
|
|
num_frames: int = 1,
|
|
multi_agent: bool = False,
|
|
as_learner_connector: bool = False,
|
|
**kwargs,
|
|
):
|
|
"""Initializes a FrameStackingConnector instance.
|
|
|
|
Args:
|
|
num_frames: The number of observation frames to stack up (into a single
|
|
observation) for the RLModule's forward pass.
|
|
multi_agent: Whether this is a connector operating on a multi-agent
|
|
observation space mapping AgentIDs to individual agents' observations.
|
|
as_learner_connector: Whether this connector is part of a Learner connector
|
|
pipeline, as opposed to an env-to-module pipeline.
|
|
"""
|
|
super().__init__(
|
|
input_observation_space=input_observation_space,
|
|
input_action_space=input_action_space,
|
|
**kwargs,
|
|
)
|
|
|
|
self._multi_agent = multi_agent
|
|
self.num_frames = num_frames
|
|
self._as_learner_connector = as_learner_connector
|
|
|
|
@override(ConnectorV2)
|
|
def __call__(
|
|
self,
|
|
*,
|
|
rl_module: RLModule,
|
|
batch: Dict[str, Any],
|
|
episodes: List[EpisodeType],
|
|
explore: Optional[bool] = None,
|
|
shared_data: Optional[dict] = None,
|
|
**kwargs,
|
|
) -> Any:
|
|
# Learner connector pipeline. Episodes have been numpy'ized.
|
|
if self._as_learner_connector:
|
|
for sa_episode in self.single_agent_episode_iterator(
|
|
episodes, agents_that_stepped_only=False
|
|
):
|
|
|
|
def _map_fn(s, _sa_episode=sa_episode):
|
|
# Squeeze out last dim.
|
|
s = np.squeeze(s, axis=-1)
|
|
# Calculate new shape and strides
|
|
new_shape = (len(_sa_episode), self.num_frames) + s.shape[1:]
|
|
new_strides = (s.strides[0],) + s.strides
|
|
# Create a strided view of the array.
|
|
# But return a copy to avoid non-contiguous memory in the object
|
|
# store (which is very expensive to deserialize).
|
|
return np.transpose(
|
|
np.lib.stride_tricks.as_strided(
|
|
s, shape=new_shape, strides=new_strides
|
|
),
|
|
axes=[0, 2, 3, 1],
|
|
).copy()
|
|
|
|
# Get all observations from the episode in one np array (except for
|
|
# the very last one, which is the final observation not needed for
|
|
# learning).
|
|
self.add_n_batch_items(
|
|
batch=batch,
|
|
column=Columns.OBS,
|
|
items_to_add=tree.map_structure(
|
|
_map_fn,
|
|
sa_episode.get_observations(
|
|
indices=slice(-self.num_frames + 1, len(sa_episode)),
|
|
neg_index_as_lookback=True,
|
|
fill=0.0,
|
|
),
|
|
),
|
|
num_items=len(sa_episode),
|
|
single_agent_episode=sa_episode,
|
|
)
|
|
|
|
# Env-to-module pipeline. Episodes still operate on lists.
|
|
else:
|
|
for sa_episode in self.single_agent_episode_iterator(episodes):
|
|
assert not sa_episode.is_numpy
|
|
# Get the list of observations to stack.
|
|
obs_stack = sa_episode.get_observations(
|
|
indices=slice(-self.num_frames, None),
|
|
fill=0.0,
|
|
)
|
|
# Observation components are (w, h, 1)
|
|
# -> concatenate along axis=-1 to (w, h, [num_frames]).
|
|
stacked_obs = tree.map_structure(
|
|
lambda *s: np.concatenate(s, axis=2),
|
|
*obs_stack,
|
|
)
|
|
self.add_batch_item(
|
|
batch=batch,
|
|
column=Columns.OBS,
|
|
item_to_add=stacked_obs,
|
|
single_agent_episode=sa_episode,
|
|
)
|
|
|
|
return batch
|
|
|
|
def _convert_individual_space(self, obs_space):
|
|
# Some assumptions: Space is box AND last dim (the stacking one) is 1.
|
|
assert isinstance(obs_space, gym.spaces.Box), obs_space
|
|
assert obs_space.shape[-1] == 1, obs_space
|
|
|
|
return gym.spaces.Box(
|
|
low=np.repeat(obs_space.low, repeats=self.num_frames, axis=-1),
|
|
high=np.repeat(obs_space.high, repeats=self.num_frames, axis=-1),
|
|
shape=list(obs_space.shape)[:-1] + [self.num_frames],
|
|
dtype=obs_space.dtype,
|
|
)
|