| 168 | |
| 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: |
| 203 | frame = obs |
| 204 | else: |
| 205 | frame = obs[self._key] |
| 206 | |
| 207 | if self._grayscale: |
| 208 | frame = cv2.cvtColor(frame, cv2.COLOR_RGB2GRAY) |
| 209 | frame = cv2.resize( |
| 210 | frame, (self._width, self._height), interpolation=cv2.INTER_AREA |
| 211 | ) |
| 212 | if self._grayscale: |
| 213 | frame = np.expand_dims(frame, -1) |
| 214 | |
| 215 | if self._key is None: |
| 216 | obs = frame |
| 217 | else: |
| 218 | obs = obs.copy() |
| 219 | obs[self._key] = frame |
| 220 | return obs |
| 221 | |
| 222 | |
| 223 | def make_atari(env_id, skip=4, max_episode_steps=None): |