Setting Up Environment Variation Parameters for Robustness Evaluation of World Models

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)—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.

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.

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

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

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

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.
• 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, 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) (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.

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