galilai-group/stable-worldmodel

Source: https://github.com/galilai-group/stable-worldmodel

stable-worldmodel

A platform for reproducible world model research and evaluation.

Installation · Quick Start · Environments · Solvers & Baselines · Documentation · Paper · Citation

stable-worldmodel provides a single, unified interface for the three stages of world model research — collecting data, training, and evaluating with model-predictive control — across a large suite of standardized environments. It ships with reference implementations of common baselines and planning solvers so research code can stay focused on the contribution that matters: the model and the objective.

Installation

From PyPI:

pip install stable-worldmodel            # base only
pip install 'stable-worldmodel[all]'     # + training, environments, and data formats

LeRobot dataset support is a separate opt-in extra (requires Python 3.12+): pip install 'stable-worldmodel[lerobot]'.

From source (development):

git clone https://github.com/galilai-group/stable-worldmodel
cd stable-worldmodel
uv venv --python=3.10 && source .venv/bin/activate
uv sync --extra all --group dev

Datasets and checkpoints are stored under $STABLEWM_HOME (defaults to ~/.stable_worldmodel/). Override the variable to point at your preferred storage location.

The library is in active development. APIs may change between minor versions.

Quick Start

import stable_worldmodel as swm
from stable_worldmodel.policy import WorldModelPolicy, PlanConfig
from stable_worldmodel.solver import CEMSolver

# 1. Collect a dataset
world = swm.World("swm/PushT-v1", num_envs=8)
world.set_policy(your_expert_policy)
world.collect("data/pusht_demo.lance", episodes=100, seed=0)

# 2. Load it and train your world model (format is autodetected)
dataset = swm.data.load_dataset("data/pusht_demo.lance", num_steps=16)
world_model = ...  # your model

# 3. Evaluate with model-predictive control
solver = CEMSolver(model=world_model, num_samples=300)
policy = WorldModelPolicy(solver=solver, config=PlanConfig(horizon=10))

world.set_policy(policy)
results = world.evaluate(episodes=50)
print(f"Success Rate: {results['success_rate']:.1f}%")

Reference implementations are provided in scripts/train/: lewm.py implements LeWM, and prejepa.py reproduces DINO-WM.

GPU utilization for LeWM trained with Push-T LanceDB dataset on a H200 GPU.

Data Formats

Recording, loading, and conversion all go through a small format registry. Pick the backend that matches your trade-off, or register your own.

world.collect("data/pusht.lance", episodes=100)                  # default: lance
world.collect("data/pusht_video", episodes=100, format="video")  # mp4 episodes

ds = swm.data.load_dataset("data/pusht.lance", num_steps=16)     # autodetect
swm.data.convert("data/pusht.lance", "data/pusht_video",
                 dest_format="video", fps=30)                    # one-shot migration

Every writer accepts a mode kwarg ('append' (default), 'overwrite', 'error'); re-running world.collect extends the existing dataset rather than failing.

Environments

Top row: default appearance  ·  Bottom row: visual factor of variation

Environments are pulled from the DeepMind Control Suite, Gymnasium classic control, OGBench, Craftax, the Arcade Learning Environment (100+ Atari games), and classical world model benchmarks (Two-Room, PushT). Most environments ship with a set of factors of variation — independently controllable visual and physical parameters (lighting, textures, dynamics, morphology) — that make it straightforward to evaluate zero-shot generalization to distribution shifts without any additional setup. Adding a new environment only requires conforming to the Gymnasium interface.

Solvers and Baselines

Command-Line Interface

After installation, the swm command is available for inspecting/converting datasets, environments, and checkpoints without writing code:

swm datasets                                        # list cached datasets
swm inspect pusht_expert_train                      # inspect a specific dataset
swm envs                                            # list all registered environments
swm fovs PushT-v1                                   # show factors of variation for an environment
swm checkpoints                                     # list available model checkpoints
swm convert pusht_expert_train --dest-format video  # convert a dataset to another format

Documentation

The full documentation lives at galilai-group.github.io/stable-worldmodel, with API references, tutorials, and guides.

Built on stable-worldmodel

• C-JEPA
• LeWM

Citation

@misc{maes_lld2026swm,
  title  = {stable-worldmodel: A Platform for Reproducible World Modeling Research and Evaluation},
  author = {Lucas Maes and Quentin Le Lidec and Luiz Facury and Nassim Massaudi and
            Ayush Chaurasia and Francesco Capuano and Richard Gao and Taj Gillin and
            Dan Haramati and Damien Scieur and Yann LeCun and Randall Balestriero},
  year   = {2026},
  eprint = {2605.21800},
  archivePrefix = {arXiv},
  primaryClass = {cs.LG},
  url    = {https://arxiv.org/abs/2605.21800},
}

Questions

Open an issue — happy to help.