| 59 | |
| 60 | |
| 61 | class NoopResetEnv(gym.Wrapper): |
| 62 | def __init__(self, env, noop_max=30): |
| 63 | """Sample initial states by taking random number of no-ops on reset. |
| 64 | No-op is assumed to be action 0. |
| 65 | """ |
| 66 | gym.Wrapper.__init__(self, env) |
| 67 | self.noop_max = noop_max |
| 68 | self.override_num_noops = None |
| 69 | self.noop_action = 0 |
| 70 | assert env.unwrapped.get_action_meanings()[0] == 'NOOP' |
| 71 | |
| 72 | def reset(self, **kwargs): |
| 73 | """ Do no-op action for a number of steps in [1, noop_max].""" |
| 74 | self.env.reset(**kwargs) |
| 75 | if self.override_num_noops is not None: |
| 76 | noops = self.override_num_noops |
| 77 | else: |
| 78 | noops = self.unwrapped.np_random.randint(1, self.noop_max + 1) #pylint: disable=E1101 |
| 79 | assert noops > 0 |
| 80 | obs = None |
| 81 | for _ in range(noops): |
| 82 | obs, _, done, _ = self.env.step(self.noop_action) |
| 83 | if done: |
| 84 | obs = self.env.reset(**kwargs) |
| 85 | return obs |
| 86 | |
| 87 | def step(self, ac): |
| 88 | return self.env.step(ac) |
| 89 | |
| 90 | |
| 91 | class EpisodicLifeEnv(gym.Wrapper): |