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

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.

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 construct the wrapper list sequentially:

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

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

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:

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:

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:

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

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