| 187 | |
| 188 | |
| 189 | class ResizeObservation(gym.ObservationWrapper): |
| 190 | def __init__(self, env, shape): |
| 191 | super().__init__(env) |
| 192 | if isinstance(shape, int): |
| 193 | self.shape = (shape, shape) |
| 194 | else: |
| 195 | self.shape = tuple(shape) |
| 196 | |
| 197 | obs_shape = self.shape + self.observation_space.shape[2:] |
| 198 | self.observation_space = Box(low=0, high=255, shape=obs_shape, dtype=np.uint8) |
| 199 | |
| 200 | def observation(self, observation): |
| 201 | transforms = T.Compose( |
| 202 | [T.Resize(self.shape, antialias=True), T.Normalize(0, 255)] |
| 203 | ) |
| 204 | observation = transforms(observation).squeeze(0) |
| 205 | return observation |
| 206 | |
| 207 | |
| 208 | # Apply Wrappers to environment |