| 53 | |
| 54 | |
| 55 | class _FireResetEnv(gym.Wrapper): |
| 56 | def __init__(self, env): |
| 57 | """Take action on reset for environments that are fixed until firing.""" |
| 58 | gym.Wrapper.__init__(self, env) |
| 59 | assert env.unwrapped.get_action_meanings()[1] == 'FIRE' |
| 60 | assert len(env.unwrapped.get_action_meanings()) >= 3 |
| 61 | |
| 62 | def reset(self): |
| 63 | self.env.reset() |
| 64 | obs, _, done, _ = self.env.step(1) |
| 65 | if done: |
| 66 | self.env.reset() |
| 67 | obs, _, done, _ = self.env.step(2) |
| 68 | if done: |
| 69 | self.env.reset() |
| 70 | return obs |
| 71 | |
| 72 | def step(self, action): |
| 73 | return self.env.step(action) |
| 74 | |
| 75 | |
| 76 | def FireResetEnv(env): |