feat: implemented RepeatAction wrapper

gymnasium/wrappers/__init__.py

@@ -58,7 +58,7 @@
     TimeLimit,
 )
 from gymnasium.wrappers.rendering import HumanRendering, RecordVideo, RenderCollection
-from gymnasium.wrappers.stateful_action import StickyAction
+from gymnasium.wrappers.stateful_action import RepeatAction, StickyAction
 from gymnasium.wrappers.stateful_observation import (
     DelayObservation,
     FrameStackObservation,

gymnasium/wrappers/stateful_action.py

@@ -1,14 +1,18 @@
-"""``StickyAction`` wrapper - There is a probability that the action is taken again."""
+"""A collection of wrappers for modifying actions.
+
+* ``StickyAction`` wrapper - There is a probability that the action is taken again.
+* ``RepeatAction`` wrapper - Repeat a single action multiple times.
+"""
 from __future__ import annotations
 
-from typing import Any
+from typing import Any, SupportsFloat
 
 import gymnasium as gym
-from gymnasium.core import ActType, ObsType
+from gymnasium.core import ActType, ObsType, WrapperActType, WrapperObsType
 from gymnasium.error import InvalidProbability
 
 
-__all__ = ["StickyAction"]
+__all__ = ["StickyAction", "RepeatAction"]
 
 
 class StickyAction(
@@ -80,3 +84,71 @@ def action(self, action: ActType) -> ActType:
 
         self.last_action = action
         return action
+
+
+class RepeatAction(
+    gym.Wrapper[ObsType, ActType, ObsType, ActType], gym.utils.RecordConstructorArgs
+):
+    """Repeatedly executes a given action in the underlying environment.
+
+    Upon calling the `step` method of this wrapper, `num_repeats`-many steps will be taken
+    with the same action in the underlying environment.
+    The wrapper sums the rewards collected from the underlying environment and returns the last
+    environment state observed.
+    If a termination or truncation is encountered during these steps, the wrapper will stop prematurely.
+    The `info` will additionally contain a field `"num_action_repetitions"`, which specifies
+    how many steps were actually taken.
+
+    Example:
+        >>> import gymnasium as gym
+        >>> env = gym.make("CartPole-v1")
+        >>> wrapped = RepeatAction(env, num_repeats=2)
+        >>> env.reset(seed=123)
+        (array([ 0.01823519, -0.0446179 , -0.02796401, -0.03156282], dtype=float32), {})
+        >>> env.step(0)
+        (array([ 0.01734283, -0.23932791, -0.02859527,  0.25216764], dtype=float32), 1.0, False, False, {})
+        >>> env.step(0)         # Perform the same action again
+        (array([ 0.01255627, -0.43403012, -0.02355192,  0.5356957 ], dtype=float32), 1.0, False, False, {})
+        >>> wrapped.reset(seed=123)         # Now we do the same thing with the `RepeatAction` wrapper
+        (array([ 0.01823519, -0.0446179 , -0.02796401, -0.03156282], dtype=float32), {})
+        >>> wrapped.step(0)
+        (array([ 0.01255627, -0.43403012, -0.02355192,  0.5356957 ], dtype=float32), 2.0, False, False, {'num_action_repetitions': 2})
+    """
+
+    def __init__(self, env: gym.Env[ObsType, ActType], num_repeats: int):
+        """Initialize RepeatAction wrapper.
+
+        Args:
+            env (Env): the wrapped environment
+            num_repeats (int): the maximum number of times to repeat the action
+        """
+        if num_repeats <= 1:
+            raise ValueError(
+                f"Number of action repeats should be greater than 1, but got {num_repeats}"
+            )
+
+        gym.utils.RecordConstructorArgs.__init__(self, num_repeats=num_repeats)
+        gym.Wrapper.__init__(self, env)
+        self._num_repeats = num_repeats
+
+    def step(
+        self, action: WrapperActType
+    ) -> tuple[WrapperObsType, SupportsFloat, bool, bool, dict[str, Any]]:
+        """Repeat `action` several times.
+
+        This step method will execute `action` at most `num_repeats`-many times in `self.env`,
+        or until a termination or truncation is encountered. The reward returned
+        is the sum of rewards collected from `self.env`. The last observation from the
+        environment is returned.
+        """
+        num_steps = 0
+        total_reward = 0
+        assert self._num_repeats > 0
+        for _ in range(self._num_repeats):
+            observation, reward, terminated, truncated, info = self.env.step(action)
+            total_reward += reward
+            num_steps += 1
+            if terminated or truncated:
+                break
+        info["num_action_repetitions"] = num_steps
+        return observation, total_reward, terminated, truncated, info