Warp frames to 84x84 as done in the Nature paper and later work. If the environment uses dictionary observations, `dict_space_key` can be specified which indicates which observation should be warped.
(self, env, width=84, height=84, grayscale=True, dict_space_key=None)
| 169 | |
| 170 | class WarpFrame(gym.ObservationWrapper): |
| 171 | def __init__(self, env, width=84, height=84, grayscale=True, dict_space_key=None): |
| 172 | """ |
| 173 | Warp frames to 84x84 as done in the Nature paper and later work. |
| 174 | If the environment uses dictionary observations, `dict_space_key` can be specified which indicates which |
| 175 | observation should be warped. |
| 176 | """ |
| 177 | super().__init__(env) |
| 178 | self._width = width |
| 179 | self._height = height |
| 180 | self._grayscale = grayscale |
| 181 | self._key = dict_space_key |
| 182 | if self._grayscale: |
| 183 | num_colors = 1 |
| 184 | else: |
| 185 | num_colors = 3 |
| 186 | |
| 187 | new_space = gym.spaces.Box( |
| 188 | low=0, |
| 189 | high=255, |
| 190 | shape=(self._height, self._width, num_colors), |
| 191 | dtype=np.uint8, |
| 192 | ) |
| 193 | if self._key is None: |
| 194 | original_space = self.observation_space |
| 195 | self.observation_space = new_space |
| 196 | else: |
| 197 | original_space = self.observation_space.spaces[self._key] |
| 198 | self.observation_space.spaces[self._key] = new_space |
| 199 | assert original_space.dtype == np.uint8 and len(original_space.shape) == 3 |
| 200 | |
| 201 | def observation(self, obs): |
| 202 | if self._key is None: |