import gymnasium as gym
import numpy as np
from lit_sb3.core.types import AtariStepReturn
class MaxAndSkipEnv(gym.Wrapper[np.ndarray, int, np.ndarray, int]):
"""Return only every ``skip``-th frame (frameskipping).
and return the max between the two last frames.
Args:
env: Environment to wrap
skip: Number of ``skip``-th frame
The same action will be taken ``skip`` times.
"""
def __init__(self, env: gym.Env, skip: int = 4) -> None:
super().__init__(env)
# most recent raw observations (for max pooling across time steps)
assert env.observation_space.dtype is not None, "No dtype specified for the observation space"
assert env.observation_space.shape is not None, "No shape defined for the observation space"
self._obs_buffer = np.zeros((2, *env.observation_space.shape), dtype=env.observation_space.dtype)
self._skip = skip
def step(self, action: int) -> AtariStepReturn:
"""Step the environment with the given action.
Repeat action, sum reward, and max over last observations.
Args:
action: the action
Returns:
observation, reward, terminated, truncated, information
"""
total_reward = 0.0
terminated = truncated = False
for i in range(self._skip):
obs, reward, terminated, truncated, info = self.env.step(action)
done = terminated or truncated
if i == self._skip - 2:
self._obs_buffer[0] = obs
if i == self._skip - 1:
self._obs_buffer[1] = obs
total_reward += float(reward)
if done:
break
# Note that the observation on the done=True frame
# doesn't matter
max_frame = self._obs_buffer.max(axis=0)
return max_frame, total_reward, terminated, truncated, info