# Implementing Custom Environment Wrappers with Pre‑wrappers and Extra‑wrappers in stable‑worldmodel

> Learn to implement custom environment wrappers in stable-worldmodel using pre-wrappers for pre-processing and extra-wrappers for post-processing. Enhance your model's data pipeline effectively.

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

---

**Use `pre_wrappers` for modifications that need to run before the core `MegaWrapper` processes pixels, and `extra_wrappers` for augmentations that operate on the canonical info‑centric representation afterward.**

The `stable‑worldmodel` library provides a flexible hook system that lets you inject custom Gymnasium wrappers at two distinct points in the environment pipeline. This design allows you to modify raw environment behavior, inject custom observations, or apply visual augmentations without forking the core library.

## Understanding the Wrapper Pipeline Architecture

The `World` class orchestrates environment creation by assembling a sequential chain of wrappers defined in [`stable_worldmodel/world/world.py`](https://github.com/galilai-group/stable-worldmodel/blob/main/stable_worldmodel/world/world.py).

### Pre‑wrappers vs Extra‑wrappers

The pipeline distinguishes between two insertion points:

- **`pre_wrappers`** – Applied **before** `MegaWrapper`. These receive the raw Gymnasium environment and are ideal for low‑level modifications like action repetition, reward shaping, or observation injection that should be invisible to the pixel rendering logic.
- **`extra_wrappers`** – Applied **after** `MegaWrapper`. These operate on the standardized representation where pixels and goal images have been lifted into the `info` dictionary, making them suitable for frame stacking, normalization, or visual noise injection.

Lines 31‑43 of [`stable_worldmodel/world/world.py`](https://github.com/galilai-group/stable-worldmodel/blob/main/stable_worldmodel/world/world.py) construct the wrapper list sequentially:

```python
wrappers = [
    *(pre_wrappers or []),                           # 1️⃣ pre‑wrappers

    partial(
        MegaWrapper,
        image_shape=image_shape,
        pixels_transform=image_transform,
        goal_transform=goal_transform,
        separate_goal=goal_conditioned,
        image_resample=image_resample,
        add_pixels=add_pixels,
    ),                                               # 2️⃣ core MegaWrapper

    *(extra_wrappers or []),                         # 3️⃣ extra‑wrappers

]

```

### The MegaWrapper Core

`MegaWrapper` (defined in [`stable_worldmodel/wrapper/default.py`](https://github.com/galilai-group/stable-worldmodel/blob/main/stable_worldmodel/wrapper/default.py) lines 92‑106) serves as the central preprocessing hub. It handles pixel rendering, lifts all observations into the `info` dictionary under keys like `pixels`, validates required keys, and resizes goal images. Any wrapper placed after this point can assume the existence of a standardized `info['pixels']` entry if `add_pixels=True`.

## Creating a Custom Wrapper

All wrappers follow the standard Gymnasium pattern. They must inherit from `gym.Wrapper` (or `gym.ObservationWrapper`, `gym.ActionWrapper`) and implement the `reset` and/or `step` methods.

Here is a pre‑wrapper example that repeats actions:

```python
import gymnasium as gym
import numpy as np

class ActionRepeatWrapper(gym.Wrapper):
    """Repeat the same action `k` times per step."""
    def __init__(self, env: gym.Env, k: int = 4):
        super().__init__(env)
        self.k = int(k)

    def step(self, action):
        total_reward = 0.0
        for _ in range(self.k):
            obs, reward, terminated, truncated, info = self.env.step(action)
            total_reward += reward
            if terminated or truncated:
                break
        return obs, total_reward, terminated, truncated, info

```

Place this in any module (e.g., [`my_wrappers.py`](https://github.com/galilai-group/stable-worldmodel/blob/main/my_wrappers.py)) and pass it to the `World` constructor as a factory function.

## Configuring Pre‑wrappers and Extra‑wrappers in World

Instantiate `World` with lists of wrapper factories (callables that accept an `env` and return a wrapped `env`):

```python
import stable_worldmodel as swm
from stable_worldmodel.wrapper import BlurWrapper, NoiseWrapper
from my_wrappers import ActionRepeatWrapper

# 1️⃣ Pre‑wrappers – run before MegaWrapper

pre = [
    lambda env: ActionRepeatWrapper(env, k=3),
]

# 2️⃣ Extra‑wrappers – run after MegaWrapper

extra = [
    lambda env: BlurWrapper(env, kernel=7),
    lambda env: NoiseWrapper(env, std=5.0),
]

# Create a World with both hook points

world = swm.World(
    env_name="swm/PushT-v1",
    num_envs=8,
    image_shape=(84, 84),
    pre_wrappers=pre,
    extra_wrappers=extra,
    image_transform=None,
)

```

**Key implementation details:**

- `pre_wrappers` receive the **raw** environment before pixel rendering occurs.
- `extra_wrappers` see the **canonical** info dictionary produced by `MegaWrapper`, allowing them to safely access `info['pixels']`.
- Both parameters accept factory functions (e.g., `lambda env: ...`), enabling easy configuration from external config files or CLI arguments.

## Practical Examples

### Action Repetition as a Pre‑wrapper

Action repetition modifies the environment dynamics by holding an action constant across multiple physics steps. Since this affects the underlying simulation before any pixel observation is generated, it belongs in `pre_wrappers`:

```python
pre_wrappers=[
    lambda env: ActionRepeatWrapper(env, k=4)
]

```

This wrapper aggregates rewards over the repeated steps and correctly propagates termination flags, as shown in the `ActionRepeatWrapper` implementation above.

### Frame Stacking as an Extra‑wrapper

Frame stacking requires access to the rendered pixel buffer. Since `MegaWrapper` ensures `info['pixels']` exists, this belongs in `extra_wrappers`:

```python
import numpy as np

class FrameStackWrapper(gym.Wrapper):
    """Stack the last N pixel frames into a single array."""
    def __init__(self, env: gym.Env, n: int = 4):
        super().__init__(env)
        self.n = n
        self.buffer = np.zeros((n,) + env.observation_space.shape, dtype=np.uint8)

    def reset(self, *args, **kwargs):
        obs, info = self.env.reset(*args, **kwargs)
        self.buffer[:] = obs['pixels'] if isinstance(obs, dict) else obs
        info['pixels'] = self.buffer.copy()
        return obs, info

    def step(self, action):
        obs, reward, term, trunc, info = self.env.step(action)
        self.buffer = np.roll(self.buffer, shift=-1, axis=0)
        self.buffer[-1] = obs['pixels'] if isinstance(obs, dict) else obs
        info['pixels'] = self.buffer.copy()
        return obs, reward, term, trunc, info

```

Configure it as an extra wrapper:

```python
world = swm.World(
    ...,
    extra_wrappers=[lambda e: FrameStackWrapper(e, n=4)]
)

```

Now every entry in `info['pixels']` contains a stack of the four most recent frames, ready for temporal convolution networks.

## Summary

- **`pre_wrappers`** run before `MegaWrapper` and modify raw environment behavior (action repetition, reward shaping).
- **`extra_wrappers`** run after `MegaWrapper` and operate on the standardized pixel representation (frame stacking, visual augmentations).
- The pipeline is constructed in [`stable_worldmodel/world/world.py`](https://github.com/galilai-group/stable-worldmodel/blob/main/stable_worldmodel/world/world.py) lines 31‑43, applying wrappers in the order: pre → MegaWrapper → extra.
- Wrapper factories must be callable as `wrapper(env)` and can be provided as lambdas for easy parameterization.

## Frequently Asked Questions

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

`pre_wrappers` are applied to the raw Gymnasium environment before any pixel rendering or info lifting occurs, making them suitable for modifying physics or action spaces. `extra_wrappers` are applied after `MegaWrapper` processes the environment, meaning they operate on the canonical representation where observations are stored in `info['pixels']` and goal images are standardized.

### Can I use standard Gymnasium wrappers like NormalizeObservation as pre_wrappers?

Yes. Any callable that accepts an environment and returns a wrapped environment works as a pre‑wrapper. Standard Gymnasium wrappers like `gym.wrappers.NormalizeObservation` or custom wrappers inheriting from `gym.Wrapper` are fully compatible with the `pre_wrappers` list.

### Where does MegaWrapper store pixel observations?

`MegaWrapper` lifts rendered pixels into the `info` dictionary under the key `'pixels'` (and optionally separate goal keys if `separate_goal=True`). This happens in [`stable_worldmodel/wrapper/default.py`](https://github.com/galilai-group/stable-worldmodel/blob/main/stable_worldmodel/wrapper/default.py), ensuring that subsequent `extra_wrappers` can access visual data through `info['pixels']` regardless of the original observation space.

### How do I apply visual augmentations like blur or noise?

Import the ready‑made wrappers from `stable_worldmodel.wrapper.visual` (e.g., `BlurWrapper`, `NoiseWrapper`) and place them in `extra_wrappers`. These wrappers specifically target the `info['pixels']` entries and any other `info['pixels*']` keys, allowing them to augment both observations and goal images consistently.