| 144 | # frame = cv2.resize(frame, (self.width, self.height), interpolation=cv2.INTER_AREA) |
| 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: |
| 162 | frame = cv2.cvtColor(frame, cv2.COLOR_RGB2GRAY) |
| 163 | frame = cv2.resize(frame, (self.width, self.height), interpolation=cv2.INTER_AREA) |
| 164 | if self.grayscale: |
| 165 | frame = np.expand_dims(frame, -1) |
| 166 | return frame |
| 167 | |
| 168 | |
| 169 | class FrameStack(gym.Wrapper): |