lit_sb3.vec_env.stacked_observations

source module lit_sb3.vec_env.stacked_observations

Classes

• StackedObservations — Frame stacking wrapper for data.

source class StackedObservations[TObs: (np.ndarray, dict[str, np.ndarray])](num_envs: int, n_stack: int, observation_space: spaces.Box | spaces.Dict, channels_order: str | Mapping[str, str | None] | None = None)

Frame stacking wrapper for data.

Dimension to stack over is either first (channels-first) or last (channels-last), which is detected automatically using common.preprocessing.is_image_space_channels_first if observation is an image space.

Parameters

• num_envs : int — Number of environments

• n_stack : int — Number of frames to stack

• observation_space : spaces.Box | spaces.Dict — Environment observation space

• channels_order : str | Mapping[str, str | None] | None — If "first", stack on first image dimension. If "last", stack on last dimension. If None, automatically detect channel to stack over in case of image observation or default to "last". For Dict space, channels_order can also be a dictionary.

Methods

• compute_stacking — Calculates the parameters in order to stack observations.

• reset — Reset the stacked_obs, add the reset observation to the stack, and return the stack.

• update — Add the observations to the stack and use the dones to update the infos.

source staticmethod StackedObservations.compute_stacking(n_stack: int, observation_space: spaces.Box, channels_order: str | None = None) → tuple[bool, int, tuple[int, ...], int]

Calculates the parameters in order to stack observations.

Parameters

• n_stack : int — Number of observations to stack

• observation_space : spaces.Box — Observation space

• channels_order : str | None — Order of the channels

Returns

• tuple[bool, int, tuple[int, ...], int] — Tuple of channels_first, stack_dimension, stackedobs, repeat_axis

source method StackedObservations.reset(observation: TObs) → TObs

Reset the stacked_obs, add the reset observation to the stack, and return the stack.

Parameters

• observation : TObs — Reset observation

Returns

• TObs — The stacked reset observation

source method StackedObservations.update(observations: TObs, dones: np.ndarray, infos: list[dict[str, Any]]) → tuple[TObs, list[dict[str, Any]]]

Add the observations to the stack and use the dones to update the infos.

Parameters

• observations : TObs — Observations

• dones : np.ndarray — Dones

• infos : list[dict[str, Any]] — Infos

Returns

• tuple[TObs, list[dict[str, Any]]] — Tuple of the stacked observations and the updated infos