228 lines
8.7 KiB
Python
228 lines
8.7 KiB
Python
import warnings
|
|
from typing import Dict, Tuple, 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 (int(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 int(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)."
|
|
)
|