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

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, 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:

• 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 (lines 38-45 and 111-124).

The ActionRepeatWrapper Class

Create a wrapper that repeats actions and accumulates rewards:

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:

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:

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:

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:

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 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 or a local 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.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}