# How to Implement Action Repeat and Frame Stacking Wrappers for World Models

> Learn to implement action repeat and frame stacking wrappers for world models using the galilai-group/stable-worldmodel library. Streamline your environment modifications efficiently.

- Repository: [GalilAI-group/stable-worldmodel](https://github.com/galilai-group/stable-worldmodel)
- Tags: how-to-guide
- Published: 2026-05-30

---

**You implement action repeat and frame stacking in stable-worldmodel by passing custom Gymnasium wrappers to the `pre_wrappers` and `extra_wrappers` arguments of the `World` class, which insert environment modifications before and after the canonical observation processing pipeline.**

The **stable-worldmodel** framework provides a flexible wrapper system that lets you modify raw environment interactions without altering core world-model logic. By leveraging the `pre_wrappers` and `extra_wrappers` parameters in [`stable_worldmodel/world/world.py`](https://github.com/galilai-group/stable-worldmodel/blob/main/stable_worldmodel/world/world.py), you can implement action repetition and frame stacking while preserving compatibility with the `MegaWrapper` and `EverythingToInfoWrapper` classes that handle observation lifting.

## Understanding the Wrapper Pipeline Architecture

The `World` class constructs the environment pipeline using two distinct wrapper insertion points defined in lines 31-44 of [`stable_worldmodel/world/world.py`](https://github.com/galilai-group/stable-worldmodel/blob/main/stable_worldmodel/world/world.py):

- **`pre_wrappers`**: Applied **before** the `MegaWrapper`, closest to the raw environment. Use this for action repeat, reward shaping, or observation injection.
- **`extra_wrappers`**: Applied **after** the `MegaWrapper`, operating on the canonical observation dictionary. Use this for frame stacking, normalization, or pixel augmentations.

Because the `EverythingToInfoWrapper` lifts every observation into the **info** dictionary, wrappers that modify step mechanics or stack historical frames can operate cleanly without touching the world-model's core logic.

## Implementing Action Repeat Wrappers (pre-wrapper)

Action repetition aggregates multiple environment steps into a single world-model step, summing rewards and respecting early termination. The repository demonstrates this pattern for MuJoCo environments in [`stable_worldmodel/envs/dmcontrol/dmcontrol.py`](https://github.com/galilai-group/stable-worldmodel/blob/main/stable_worldmodel/envs/dmcontrol/dmcontrol.py) (lines 38-45 and 111-124).

### The ActionRepeatWrapper Class

Create a wrapper that repeats actions and accumulates rewards:

```python
import gymnasium as gym
import numpy as np
from typing import Any

class ActionRepeatWrapper(gym.Wrapper):
    """Repeat each action N times, summing rewards and stopping early on termination."""
    def __init__(self, env: gym.Env, repeat: int = 1):
        super().__init__(env)
        if repeat < 1:
            raise ValueError("repeat must be >= 1")
        self.repeat = repeat

    def step(self, action: Any) -> tuple[Any, float, bool, bool, dict]:
        total_reward = 0.0
        terminated = truncated = False
        info: dict = {}
        
        for _ in range(self.repeat):
            obs, reward, terminated, truncated, info = self.env.step(action)
            total_reward += float(reward)
            if terminated or truncated:
                break
                
        return obs, total_reward, terminated, truncated, info

    def reset(self, **kwargs: Any):
        return self.env.reset(**kwargs)

```

### Integration with World.pre_wrappers

Pass the wrapper via `functools.partial` to set the repetition count:

```python
import functools as ft
from stable_worldmodel import World

world = World(
    env_name="swm/PushT-v1",
    num_envs=8,
    image_shape=(84, 84),
    pre_wrappers=[ft.partial(ActionRepeatWrapper, repeat=4)],
    extra_wrappers=[],
)

```

The wrapper executes inside the `EnvPool` before `MegaWrapper` processes observations, ensuring downstream code receives the summed reward and final observation state.

## Implementing Frame Stacking Wrappers (extra-wrapper)

Frame stacking provides temporal context by concatenating consecutive observations along a time dimension. Because `MegaWrapper` lifts observations into the info dictionary, you must stack frames inside the `info` dict rather than modifying the raw observation space.

### The FrameStackInfoWrapper Class

Implement a wrapper that maintains deques of historical frames per environment:

```python
import collections
import gymnasium as gym
import numpy as np
from typing import Any

class FrameStackInfoWrapper(gym.Wrapper):
    """
    Stacks the last ``num_stack`` frames of a chosen key (default: 'pixels')
    inside the ``info`` dict. The stacked array has shape
    (num_stack, H, W, C) and replaces the original entry.
    """
    def __init__(self, env: gym.Env, num_stack: int = 4, key: str = "pixels"):
        super().__init__(env)
        self.num_stack = num_stack
        self.key = key
        self._buffers: list[collections.deque] = [
            collections.deque(maxlen=num_stack) for _ in range(self.env.unwrapped.num_envs)
        ]

    def reset(self, **kwargs: Any):
        obs, info = self.env.reset(**kwargs)
        
        for buf in self._buffers:
            buf.clear()
            
        for i in range(self.env.unwrapped.num_envs):
            frame = info[self.key][i] if self.key in info else None
            for _ in range(self.num_stack):
                self._buffers[i].append(frame)
                
        self._update_info(info)
        return obs, info

    def step(self, action: Any):
        obs, reward, terminated, truncated, info = self.env.step(action)
        
        for i in range(self.env.unwrapped.num_envs):
            frame = info[self.key][i] if self.key in info else None
            self._buffers[i].append(frame)
            
        self._update_info(info)
        return obs, reward, terminated, truncated, info

    def _update_info(self, info: dict):
        """Replace the key with a stacked array (num_env, num_stack, …)."""
        stacked = np.stack([np.stack(buf, axis=0) for buf in self._buffers], axis=0)
        info[self.key] = stacked

```

### Integration with World.extra_wrappers

Insert the wrapper in the `extra_wrappers` list to operate after `MegaWrapper` has populated the `pixels` entry:

```python
import functools as ft
from stable_worldmodel import World

world = World(
    env_name="swm/PushT-v1",
    num_envs=8,
    image_shape=(84, 84),
    pre_wrappers=[],
    extra_wrappers=[
        ft.partial(FrameStackInfoWrapper, num_stack=3, key="pixels")
    ],
)

```

The world-model policy now receives a `pixels` tensor of shape `(num_envs, 3, 84, 84, 3)` representing three stacked frames, accessible via `info["pixels"]` in `WorldModelPolicy` or custom solvers.

## Complete Implementation Example

Combine both wrappers for a full preprocessing pipeline that repeats actions and stacks frames:

```python
import functools as ft
import stable_worldmodel as swm

# Configure world with action repeat and frame stacking

world = swm.World(
    env_name="swm/PushT-v1",
    num_envs=4,
    image_shape=(84, 84),
    pre_wrappers=[
        ft.partial(swm.wrapper.ActionRepeatWrapper, repeat=2)
    ],
    extra_wrappers=[
        ft.partial(swm.wrapper.FrameStackInfoWrapper, num_stack=4, key="pixels")
    ],
)

# Use for data collection

from stable_worldmodel.policy import RandomPolicy
world.set_policy(RandomPolicy(seed=123))
world.collect(path="data/stacked_repeat.lance", episodes=100)

# Use for evaluation with a world model

policy = swm.policy.WorldModelPolicy(
    solver=swm.solver.TDMPCSolver(),
    config=swm.policy.PlanConfig(
        horizon=15,
        receding_horizon=5,
        action_block=2,
        history_len=4,  # Matches frame_stack size

    ),
)
world.set_policy(policy)
results = world.evaluate(episodes=20, seed=42)

```

## Summary

- **Action repeat wrappers** belong in `pre_wrappers` to aggregate rewards and terminal conditions before the `MegaWrapper` processes observations.
- **Frame stacking wrappers** belong in `extra_wrappers` to stack the `pixels` entry inside the info dictionary after `MegaWrapper` lifts observations.
- The `World` class in [`stable_worldmodel/world/world.py`](https://github.com/galilai-group/stable-worldmodel/blob/main/stable_worldmodel/world/world.py) applies wrappers in the order: `pre_wrappers` → `MegaWrapper` → `extra_wrappers`.
- Use `functools.partial` to configure wrapper parameters when passing them to the `World` constructor.
- Custom wrappers should inherit from `gymnasium.Wrapper` and maintain compatibility with vectorized environments via `self.env.unwrapped.num_envs`.

## Frequently Asked Questions

### What is the difference between pre_wrappers and extra_wrappers in stable-worldmodel?

**`pre_wrappers`** are applied before the `MegaWrapper` and operate on raw environment outputs, making them suitable for action repeat logic that must sum rewards across multiple physical steps. **`extra_wrappers`** are applied after the `MegaWrapper` and can access the canonical observation dictionary (including the `pixels` key lifted into `info`), making them appropriate for frame stacking and pixel augmentations.

### How does frame stacking interact with the info dictionary?

The `EverythingToInfoWrapper` (used inside `MegaWrapper`) moves observations into the `info` dictionary under keys like `pixels`. The `FrameStackInfoWrapper` must read from and write to this `info` key rather than the observation space, maintaining the dictionary structure that downstream world-model policies expect.

### Can I use multiple action repeat and frame stacking wrappers together?

Yes. The `World` class accepts lists for both arguments, applying them in sequence. You can chain multiple preprocessing steps—such as an action repeat wrapper followed by a reward shaping wrapper in `pre_wrappers`, or a frame stacker followed by a normalization wrapper in `extra_wrappers`—without conflicts.

### Where should I save custom wrapper classes for reuse?

Store reusable wrapper classes in a new module such as [`stable_worldmodel/wrapper/custom.py`](https://github.com/galilai-group/stable-worldmodel/blob/main/stable_worldmodel/wrapper/custom.py) or a local [`wrappers.py`](https://github.com/galilai-group/stable-worldmodel/blob/main/wrappers.py) file in your project. Import them and reference them via `functools.partial` when constructing `World` instances, keeping your experiment configuration separate from the wrapper implementations.