Skip to content

lit_sb3.wrappers.atari.fire_reset

[docs] module lit_sb3.wrappers.atari.fire_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
import gymnasium as gym
import numpy as np

from lit_sb3.core.types import AtariResetReturn


class FireResetEnv(gym.Wrapper[np.ndarray, int, np.ndarray, int]):
    """Take action on reset for environments that are fixed until firing.

    Args:
        env: Environment to wrap
    """

    def __init__(self, env: gym.Env) -> None:
        super().__init__(env)
        assert env.unwrapped.get_action_meanings()[1] == "FIRE"  # type: ignore[attr-defined]
        assert len(env.unwrapped.get_action_meanings()) >= 3  # type: ignore[attr-defined]

    def reset(self, **kwargs) -> AtariResetReturn:
        self.env.reset(**kwargs)
        obs, _, terminated, truncated, _ = self.env.step(1)
        if terminated or truncated:
            self.env.reset(**kwargs)
        obs, _, terminated, truncated, _ = self.env.step(2)
        if terminated or truncated:
            self.env.reset(**kwargs)
        return obs, {}