lit_sb3.wrappers.atari.noop_reset
[docs]
module
lit_sb3.wrappers.atari.noop_reset
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37 | import gymnasium as gym
import numpy as np
from lit_sb3.core.types import AtariResetReturn
class NoopResetEnv(gym.Wrapper[np.ndarray, int, np.ndarray, int]):
"""Sample initial states by taking random number of no-ops on reset.
No-op is assumed to be action 0.
Args:
env: Environment to wrap
noop_max: Maximum value of no-ops to run
"""
def __init__(self, env: gym.Env, noop_max: int = 30) -> None:
super().__init__(env)
self.noop_max = noop_max
self.override_num_noops = None
self.noop_action = 0
assert env.unwrapped.get_action_meanings()[0] == "NOOP" # type: ignore[attr-defined]
def reset(self, **kwargs) -> AtariResetReturn:
self.env.reset(**kwargs)
if self.override_num_noops is not None:
noops = self.override_num_noops
else:
noops = self.unwrapped.np_random.integers(1, self.noop_max + 1)
assert noops > 0
obs = np.zeros(0)
info: dict = {}
for _ in range(noops):
obs, _, terminated, truncated, info = self.env.step(self.noop_action)
if terminated or truncated:
obs, info = self.env.reset(**kwargs)
return obs, info
|