Warp frames to 84x84 as done in the Nature paper and later work.
(self, env, width=84, height=84, grayscale=True)
| 145 | # return frame[:, :, None] |
| 146 | class WarpFrame(gym.ObservationWrapper): |
| 147 | def __init__(self, env, width=84, height=84, grayscale=True): |
| 148 | """Warp frames to 84x84 as done in the Nature paper and later work.""" |
| 149 | gym.ObservationWrapper.__init__(self, env) |
| 150 | self.width = width |
| 151 | self.height = height |
| 152 | self.grayscale = grayscale |
| 153 | if self.grayscale: |
| 154 | self.observation_space = spaces.Box(low=0, high=255, |
| 155 | shape=(self.height, self.width, 1), dtype=np.uint8) |
| 156 | else: |
| 157 | self.observation_space = spaces.Box(low=0, high=255, |
| 158 | shape=(self.height, self.width, 3), dtype=np.uint8) |
| 159 | |
| 160 | def observation(self, frame): |
| 161 | if self.grayscale: |