| 168 | |
| 169 | |
| 170 | class GrayScaleObservation(gym.ObservationWrapper): |
| 171 | def __init__(self, env): |
| 172 | super().__init__(env) |
| 173 | obs_shape = self.observation_space.shape[:2] |
| 174 | self.observation_space = Box(low=0, high=255, shape=obs_shape, dtype=np.uint8) |
| 175 | |
| 176 | def permute_orientation(self, observation): |
| 177 | # permute [H, W, C] array to [C, H, W] tensor |
| 178 | observation = np.transpose(observation, (2, 0, 1)) |
| 179 | observation = torch.tensor(observation.copy(), dtype=torch.float) |
| 180 | return observation |
| 181 | |
| 182 | def observation(self, observation): |
| 183 | observation = self.permute_orientation(observation) |
| 184 | transform = T.Grayscale() |
| 185 | observation = transform(observation) |
| 186 | return observation |
| 187 | |
| 188 | |
| 189 | class ResizeObservation(gym.ObservationWrapper): |