lit_sb3.wrappers.atari.sticky_action
[docs]
module
lit_sb3.wrappers.atari.sticky_action
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 | import gymnasium as gym
import numpy as np
from lit_sb3.core.types import AtariResetReturn, AtariStepReturn
class StickyActionEnv(gym.Wrapper[np.ndarray, int, np.ndarray, int]):
"""Sticky action.
Paper: https://arxiv.org/abs/1709.06009
Official implementation: https://github.com/mgbellemare/Arcade-Learning-Environment
Args:
env: Environment to wrap
action_repeat_probability: Probability of repeating the last action
"""
def __init__(self, env: gym.Env, action_repeat_probability: float) -> None:
super().__init__(env)
self.action_repeat_probability = action_repeat_probability
assert env.unwrapped.get_action_meanings()[0] == "NOOP" # type: ignore[attr-defined]
def reset(self, **kwargs) -> AtariResetReturn:
self._sticky_action = 0 # NOOP
return self.env.reset(**kwargs)
def step(self, action: int) -> AtariStepReturn:
if self.np_random.random() >= self.action_repeat_probability:
self._sticky_action = action
return self.env.step(self._sticky_action)
|