1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242 | from typing import Union
import warnings
from gymnasium import spaces
import numpy as np
import torch
from torch.nn import functional as F # noqa: N812
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).
Args:
observation_space:
observation_space: The observation space parameter.
Returns:
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.", stacklevel=2
)
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]
Args:
observation_space:
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.
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]).
observation_space: The observation space parameter.
"""
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.
Args:
observation: The input observation to process.
observation_space: The observation space parameter.
Returns:
channel first observation if observation is an image
"""
# Avoid circular import
from lit_sb3.vec_env import VecTransposeImage
if is_image_space(observation_space) and 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: 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.
Args:
obs: Observation
observation_space:
normalize_images: Whether to normalize images or not
(True by default)
observation_space: The observation space parameter.
"""
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) # ty:ignore[invalid-argument-type]
return preprocessed_obs
assert isinstance(obs, torch.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 torch.cat(
[
F.one_hot(obs_.long(), num_classes=int(observation_space.nvec[idx])).float()
for idx, obs_ in enumerate(torch.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,
) -> tuple[int, ...] | dict[str, tuple[int, ...]]:
"""Get the shape of the observation (useful for the buffers).
Args:
observation_space:
observation_space: The observation space parameter.
"""
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.
Args:
observation_space:
observation_space: The observation space parameter.
"""
# 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.
Args:
action_space:
action_space: The action space parameter.
"""
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.
Args:
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)."
)
|