Skip to content

lit_sb3.core.preprocessing

source module lit_sb3.core.preprocessing

Functions

source is_image_space_channels_first(observation_space: spaces.Box) → bool

Check if an image observation space (see is_image_space).

is channels-first (CxHxW, True) or channels-last (HxWxC, False).

Use a heuristic that channel dimension is the smallest of the three. If second dimension is smallest, raise an exception (no support).

Parameters

  • observation_space : spaces.Box

  • observation_space : spaces.Box — The observation space parameter.

Returns

  • bool — True if observation space is channels-first image, False if channels-last.

source is_image_space(observation_space: spaces.Space, check_channels: bool = False, normalized_image: bool = False) → bool

Check if a observation space has the shape, limits and dtype.

of a valid image. The check is conservative, so that it returns False if there is a doubt.

Valid images: RGB, RGBD, GrayScale with values in [0, 255]

Parameters

  • observation_space : spaces.Space

  • check_channels : bool — Whether to do or not the check for the number of channels. e.g., with frame-stacking, the observation space may have more channels than expected.

  • normalized_image : bool — Whether to assume that the image is already normalized or not (this disables dtype and bounds checks): when True, it only checks that the space is a Box and has 3 dimensions. Otherwise, it checks that it has expected dtype (uint8) and bounds (values in [0, 255]).

  • observation_space : spaces.Space — The observation space parameter.

source maybe_transpose(observation: np.ndarray, observation_space: spaces.Space) → np.ndarray

Handle the different cases for images as PyTorch use channel first format.

Parameters

  • observation : np.ndarray — The input observation to process.

  • observation_space : spaces.Space — The observation space parameter.

Returns

  • np.ndarray — channel first observation if observation is an image

source preprocess_obs(obs: torch.Tensor | dict[str, torch.Tensor], observation_space: spaces.Space, normalize_images: bool = True) → torch.Tensor | dict[str, torch.Tensor]

Preprocess observation to be to a neural network.

For images, it normalizes the values by dividing them by 255 (to have values in [0, 1]) For discrete observations, it create a one hot vector.

Parameters

  • obs : torch.Tensor | dict[str, torch.Tensor] — Observation

  • observation_space : spaces.Space

  • normalize_images : bool — Whether to normalize images or not (True by default)

  • observation_space : spaces.Space — The observation space parameter.

Raises

  • NotImplementedError

source get_obs_shape(observation_space: spaces.Space) → tuple[int, ...] | dict[str, tuple[int, ...]]

Get the shape of the observation (useful for the buffers).

Parameters

  • observation_space : spaces.Space

  • observation_space : spaces.Space — The observation space parameter.

Raises

  • NotImplementedError

source get_flattened_obs_dim(observation_space: spaces.Space) → int

Get the dimension of the observation space when flattened.

It does not apply to image observation space.

Used by the FlattenExtractor to compute the input shape.

Parameters

  • observation_space : spaces.Space

  • observation_space : spaces.Space — The observation space parameter.

source get_action_dim(action_space: spaces.Space) → int

Get the dimension of the action space.

Parameters

  • action_space : spaces.Space

  • action_space : spaces.Space — The action space parameter.

Raises

  • NotImplementedError

source check_for_nested_spaces(obs_space: spaces.Space) → None

Make sure the observation space does not have nested spaces (Dicts/Tuples inside Dicts/Tuples).

If so, raise an Exception informing that there is no support for this.

Parameters

  • obs_space : spaces.Space — an observation space

Raises

  • NotImplementedError