import warnings from typing import Union import numpy as np import torch as th from gymnasium import spaces from torch.nn import functional as F def 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). :param observation_space: :return: True if observation space is channels-first image, False if channels-last. """ smallest_dimension = np.argmin(observation_space.shape).item() if smallest_dimension == 1: warnings.warn("Treating image space as channels-last, while second dimension was smallest of the three.") return smallest_dimension == 0 def 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] :param observation_space: :param check_channels: 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. :param normalized_image: 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]). :return: """ check_dtype = check_bounds = not normalized_image if isinstance(observation_space, spaces.Box) and len(observation_space.shape) == 3: # Check the type if check_dtype and observation_space.dtype != np.uint8: return False # Check the value range incorrect_bounds = np.any(observation_space.low != 0) or np.any(observation_space.high != 255) if check_bounds and incorrect_bounds: return False # Skip channels check if not check_channels: return True # Check the number of channels if is_image_space_channels_first(observation_space): n_channels = observation_space.shape[0] else: n_channels = observation_space.shape[-1] # GrayScale, RGB, RGBD return n_channels in [1, 3, 4] return False def maybe_transpose(observation: np.ndarray, observation_space: spaces.Space) -> np.ndarray: """ Handle the different cases for images as PyTorch use channel first format. :param observation: :param observation_space: :return: channel first observation if observation is an image """ # Avoid circular import from stable_baselines3.common.vec_env import VecTransposeImage if is_image_space(observation_space): if not (observation.shape == observation_space.shape or observation.shape[1:] == observation_space.shape): # Try to re-order the channels transpose_obs = VecTransposeImage.transpose_image(observation) if transpose_obs.shape == observation_space.shape or transpose_obs.shape[1:] == observation_space.shape: observation = transpose_obs return observation def preprocess_obs( obs: Union[th.Tensor, dict[str, th.Tensor]], observation_space: spaces.Space, normalize_images: bool = True, ) -> Union[th.Tensor, dict[str, th.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. :param obs: Observation :param observation_space: :param normalize_images: Whether to normalize images or not (True by default) :return: """ if isinstance(observation_space, spaces.Dict): # Do not modify by reference the original observation assert isinstance(obs, dict), f"Expected dict, got {type(obs)}" preprocessed_obs = {} for key, _obs in obs.items(): preprocessed_obs[key] = preprocess_obs(_obs, observation_space[key], normalize_images=normalize_images) return preprocessed_obs # type: ignore[return-value] assert isinstance(obs, th.Tensor), f"Expecting a torch Tensor, but got {type(obs)}" if isinstance(observation_space, spaces.Box): if normalize_images and is_image_space(observation_space): return obs.float() / 255.0 return obs.float() elif isinstance(observation_space, spaces.Discrete): # One hot encoding and convert to float to avoid errors return F.one_hot(obs.long(), num_classes=int(observation_space.n)).float() elif isinstance(observation_space, spaces.MultiDiscrete): # Tensor concatenation of one hot encodings of each Categorical sub-space return th.cat( [ F.one_hot(obs_.long(), num_classes=int(observation_space.nvec[idx])).float() for idx, obs_ in enumerate(th.split(obs.long(), 1, dim=1)) ], dim=-1, ).view(obs.shape[0], sum(observation_space.nvec)) elif isinstance(observation_space, spaces.MultiBinary): return obs.float() else: raise NotImplementedError(f"Preprocessing not implemented for {observation_space}") def get_obs_shape( observation_space: spaces.Space, ) -> Union[tuple[int, ...], dict[str, tuple[int, ...]]]: """ Get the shape of the observation (useful for the buffers). :param observation_space: :return: """ if isinstance(observation_space, spaces.Box): return observation_space.shape elif isinstance(observation_space, spaces.Discrete): # Observation is an int return (1,) elif isinstance(observation_space, spaces.MultiDiscrete): # Number of discrete features return (len(observation_space.nvec),) elif isinstance(observation_space, spaces.MultiBinary): # Number of binary features return observation_space.shape elif isinstance(observation_space, spaces.Dict): return {key: get_obs_shape(subspace) for (key, subspace) in observation_space.spaces.items()} # type: ignore[misc] else: raise NotImplementedError(f"{observation_space} observation space is not supported") def 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. :param observation_space: :return: """ # See issue https://github.com/openai/gym/issues/1915 # it may be a problem for Dict/Tuple spaces too... if isinstance(observation_space, spaces.MultiDiscrete): return sum(observation_space.nvec) else: # Use Gym internal method return spaces.utils.flatdim(observation_space) def get_action_dim(action_space: spaces.Space) -> int: """ Get the dimension of the action space. :param action_space: :return: """ if isinstance(action_space, spaces.Box): return int(np.prod(action_space.shape)) elif isinstance(action_space, spaces.Discrete): # Action is an int return 1 elif isinstance(action_space, spaces.MultiDiscrete): # Number of discrete actions return len(action_space.nvec) elif isinstance(action_space, spaces.MultiBinary): # Number of binary actions assert isinstance( action_space.n, int ), f"Multi-dimensional MultiBinary({action_space.n}) action space is not supported. You can flatten it instead." return int(action_space.n) else: raise NotImplementedError(f"{action_space} action space is not supported") def 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. :param obs_space: an observation space """ if isinstance(obs_space, (spaces.Dict, spaces.Tuple)): sub_spaces = obs_space.spaces.values() if isinstance(obs_space, spaces.Dict) else obs_space.spaces for sub_space in sub_spaces: if isinstance(sub_space, (spaces.Dict, spaces.Tuple)): raise NotImplementedError( "Nested observation spaces are not supported (Tuple/Dict space inside Tuple/Dict space)." )