| 4 | |
| 5 | |
| 6 | class AtariWrapper(Game): |
| 7 | def __init__(self, env, discount: float, cvt_string=True): |
| 8 | """Atari Wrapper |
| 9 | Parameters |
| 10 | ---------- |
| 11 | env: Any |
| 12 | another env wrapper |
| 13 | discount: float |
| 14 | discount of env |
| 15 | cvt_string: bool |
| 16 | True -> convert the observation into string in the replay buffer |
| 17 | """ |
| 18 | super().__init__(env, env.action_space.n, discount) |
| 19 | self.cvt_string = cvt_string |
| 20 | |
| 21 | def legal_actions(self): |
| 22 | return [_ for _ in range(self.env.action_space.n)] |
| 23 | |
| 24 | def get_max_episode_steps(self): |
| 25 | return self.env.get_max_episode_steps() |
| 26 | |
| 27 | def step(self, action): |
| 28 | observation, reward, done, info = self.env.step(action) |
| 29 | observation = observation.astype(np.uint8) |
| 30 | |
| 31 | if self.cvt_string: |
| 32 | observation = arr_to_str(observation) |
| 33 | |
| 34 | return observation, reward, done, info |
| 35 | |
| 36 | def reset(self, **kwargs): |
| 37 | observation = self.env.reset(**kwargs) |
| 38 | observation = observation.astype(np.uint8) |
| 39 | |
| 40 | if self.cvt_string: |
| 41 | observation = arr_to_str(observation) |
| 42 | |
| 43 | return observation |
| 44 | |
| 45 | def close(self): |
| 46 | self.env.close() |