| 7 | |
| 8 | |
| 9 | class NoopResetEnv(gym.Wrapper): |
| 10 | def __init__(self, env, noop_max=30): |
| 11 | """Sample initial states by taking random number of no-ops on reset. |
| 12 | No-op is assumed to be action 0. |
| 13 | """ |
| 14 | gym.Wrapper.__init__(self, env) |
| 15 | self.noop_max = noop_max |
| 16 | self.override_num_noops = None |
| 17 | self.noop_action = 0 |
| 18 | assert env.unwrapped.get_action_meanings()[0] == 'NOOP' |
| 19 | |
| 20 | def reset(self, **kwargs): |
| 21 | """ Do no-op action for a number of steps in [1, noop_max].""" |
| 22 | self.env.reset(**kwargs) |
| 23 | if self.override_num_noops is not None: |
| 24 | noops = self.override_num_noops |
| 25 | else: |
| 26 | noops = self.unwrapped.np_random.randint(1, self.noop_max + 1) # pylint: disable=E1101 |
| 27 | assert noops > 0 |
| 28 | obs = None |
| 29 | for _ in range(noops): |
| 30 | obs, _, done, _ = self.env.step(self.noop_action) |
| 31 | if done: |
| 32 | obs = self.env.reset(**kwargs) |
| 33 | return obs |
| 34 | |
| 35 | def step(self, ac): |
| 36 | return self.env.step(ac) |
| 37 | |
| 38 | |
| 39 | class FireResetEnv(gym.Wrapper): |