# Setting Up Environment Variation Parameters for Robustness Evaluation of World Models

> Explore robust world models by setting up environment variation parameters. Customize visual and physical factors for domain randomization and OOD testing with Stable WorldModel.

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

---

**Stable WorldModel exposes a hierarchical variation space through `World.reset()` options, allowing you to resample visual and physical factors or fix them to specific values for domain randomization and out-of-distribution testing.**

The `stable-worldmodel` library provides a systematic approach to robustness evaluation by equipping each registered environment with a controllable **variation space**. This architecture enables researchers to domain-randomize training data and create reproducible out-of-distribution (OOD) test sets by manipulating visual colors, object scales, physics constants, and start positions at reset time.

## Understanding the Variation Space Architecture

Every environment in the Stable WorldModel registry contains a hierarchical `Dict` space—implemented in [[`stable_worldmodel/spaces.py`](https://github.com/galilai-group/stable-worldmodel/blob/main/stable_worldmodel/spaces.py)](https://github.com/galilai-group/stable-worldmodel/blob/main/stable_worldmodel/spaces.py)—that describes every controllable property using standard `gymnasium` space types such as `Box`, `Discrete`, and `MultiDiscrete`. When wrapped by the default preprocessing pipeline (`MegaWrapper` → `EverythingToInfoWrapper`), this space is accessible via `env.unwrapped.variation_space`.

The variation system is **stateless across steps**; parameters reside in the variation space rather than the environment dynamics, ensuring that identical random seeds produce identical trajectories when the same configuration options are supplied.

## Configuring Parameters at Reset Time

The `World.reset` method forwards a user-provided `options` dictionary to the `reset_variation_space` helper, enabling fine-grained control over which factors to randomize or fix.

### Resampling Specific Factors

To randomize only a subset of properties, pass a list of dotted keys via `options['variation']`. The helper resamples these keys from their bounded distributions while preserving previous values for all other parameters.

```python
import stable_worldmodel as swm

world = swm.World('swm/PushT-v1', num_envs=4, image_shape=(84, 84))

# Resample only agent and block colors

world.reset(seed=42, options={'variation': ['agent.color', 'block.color']})

```

### Randomizing All Variations

For comprehensive domain randomization, use the special value `['all']` to resample every key in the variation space.

```python

# Randomize all controllable factors

world.reset(seed=0, options={'variation': ['all']})

```

### Fixing Explicit Values for OOD Testing

To create deterministic, reproducible test configurations—essential for out-of-distribution evaluation—supply exact values through `options['variation_values']`. This mapping overwrites any sampled values, guaranteeing that specific visual or physical configurations are exactly as specified.

```python
import numpy as np

# Freeze specific visual factors for OOD evaluation

world.reset(
    seed=0,
    options={
        'variation': ['agent.color', 'background.color'],
        'variation_values': {
            'agent.color': np.array([255, 0, 0], dtype=np.uint8),      # Red agent

            'background.color': np.array([0, 0, 0], dtype=np.uint8)  # Black background

        }
    }
)

```

## Accessing Variation Values During Rollouts

The `EverythingToInfoWrapper`—defined in [[`stable_worldmodel/wrapper/default.py`](https://github.com/galilai-group/stable-worldmodel/blob/main/stable_worldmodel/wrapper/default.py)](https://github.com/galilai-group/stable-worldmodel/blob/main/stable_worldmodel/wrapper/default.py) at lines 254-276—reads the `variation` option to determine which keys to monitor. During `reset`, it injects current values into the rollout `info` dictionary under the namespace `variation.<key>`.

```python

# Step through environment and access variation values

obs, info = world.step(action)
current_agent_color = info['variation.agent.color']  # e.g., [128, 45, 200]

```

This propagation makes sampled factors visible to downstream components such as world-model training loops and evaluation metrics collectors.

## Core Implementation Flow

The following execution chain governs how setting up environment variation parameters translates into configured environment states:

1. `World.reset(seed, options)` forwards the options dict to `reset_variation_space`.
2. The helper seeds the space, resets it, updates (resamples) selected keys, and applies explicit values from `variation_values`.
3. `EverythingToInfoWrapper` extracts the sampled values and populates the `info` dictionary.

```text
World.reset(seed, options)
   └─> stable_worldmodel.spaces.reset_variation_space(
          env.unwrapped.variation_space,
          seed,
          options,
          default_variations)
          ├─ variation_space.seed(...)
          ├─ variation_space.reset()
          ├─ variation_space.update(keys)   # Resample selected keys

          └─ variation_space.set_value(...) # Apply explicit values

   └─> EverythingToInfoWrapper (reset)
          └─ for each watched key:
                 info[f'variation.{key}'] = variation_space.<key>.value

```

## Summary

- The variation space is a hierarchical `Dict` of `gymnasium` spaces exposed via `env.unwrapped.variation_space` and defined in [`spaces.py`](https://github.com/galilai-group/stable-worldmodel/blob/main/spaces.py).
- Use `World.reset(options={'variation': [...]})` to resample specific keys or `['all']` for complete randomization.
- Supply `options['variation_values']` to override sampling with explicit, reproducible configurations for OOD testing.
- Sampled values appear in `info['variation.<key>']` via `EverythingToInfoWrapper` in [`wrapper/default.py`](https://github.com/galilai-group/stable-worldmodel/blob/main/wrapper/default.py), enabling robustness evaluation during training and inference.

## Frequently Asked Questions

### How do I view all available variation keys for an environment?

Call `world.envs.single_variation_space.to_str()` after initializing the `World`. This prints a human-readable representation of the nested keys, such as `agent.color`, `block.scale`, and `background.color`, along with their space types and bounds.

### What is the difference between `variation` and `variation_values` in reset options?

The `variation` option accepts a list of dotted keys (or `['all']`) specifying which parameters to **resample** from their distributions. In contrast, `variation_values` accepts a dictionary mapping specific keys to exact values, which **overwrite** any sampled results to guarantee reproducibility for OOD experiments.

### How are variation parameters stored between environment steps?

Variation parameters are **stateless** across steps. They reside within the variation space object (accessed via `env.unwrapped.variation_space`) rather than in the environment's internal dynamics. This design ensures that resetting with the same seed and options reproduces identical configurations regardless of previous steps.

### Where in the codebase are variation values injected into the info dictionary?

The `EverythingToInfoWrapper` class in [[`stable_worldmodel/wrapper/default.py`](https://github.com/galilai-group/stable-worldmodel/blob/main/stable_worldmodel/wrapper/default.py)](https://github.com/galilai-group/stable-worldmodel/blob/main/stable_worldmodel/wrapper/default.py) (lines 254-276) handles this injection. During the reset phase, it reads the active variation keys and writes their current values into the `info` dict under the `variation.` namespace.