Quick Example | Environments | Installation | Documentation | Contributing | Citation

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

# collect a dataset
world = swm.World('swm/PushT-v1', num_envs=8)
world.set_policy(your_expert_policy)
world.record_dataset(dataset_name='pusht_demo', episodes=100)

# load dataset and train your world model
dataset = HDF5Dataset(name='pusht_demo', num_steps=16)
world_model = ...  # your world-model

# 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}%")

stable-worldmodel eases reproducibility by already implementing several baselines: scripts/train/prejepa.py reproduces results from the DINO-WM paper and scripts/train/gcivl.py implements several goal-conditioned RL algorithms. To foster research in MPC for world models, several planning solvers are already implemented, including zeroth-order (CEM, MPPI), gradient-based (GradientSolver, PGD), and constrained gradient approaches (LagrangianSolver).

We support multiple dataset formats to optimize efficiency: MP4 enables fast and convenient visualization, while HDF5 ensures high-performance data loading, reduces CPU bottlenecks, and improves overall GPU utilization.

GPU utilization for DINO-WM trained on Push-T with a DINOv2-Small backbone.

See the full documentation here.

stable-worldmodel supports a large collection of environments from the DeepMind Control Suite, OGBench, and classical world model benchmarks such as Two-Room and PushT.

Each environment includes visual and physical factor variations to evaluate robustness and generalization. New environments can easily be added to stable-worldmodel as they only need to follow the Gymnasium interface.

After installation, the swm command-line tool is available to inspect your datasets, environments, and checkpoints without writing any code:

# list cached datasets
swm datasets

# inspect a specific dataset
swm inspect pusht_expert_train

# list all registered environments
swm envs

# show factors of variation for an environment
swm fovs PushT-v1

# list available model checkpoints
swm checkpoints

stable-worldmodel is available on PyPI and can be installed with:

pip install stable-worldmodel

Note: The library is still in active development.

To set up a development environment from source:

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

Note: All datasets and models will be saved in the $STABLEWM_HOME environment variable. By default this is ~/.stable-wm/. Adapt this directory according to your storage needs.

If you have a question, please file an issue.

The following research projects have been built on top of stable-worldmodel:

• C-JEPA
• LeWM

@misc{maes_lelidec2026swm-1,
      title={stable-worldmodel-v1: Reproducible World Modeling Research and Evaluation},
      author = {Lucas Maes and Quentin Le Lidec and Dan Haramati and
                Nassim Massaudi and Damien Scieur and Yann LeCun and
                Randall Balestriero},
      year={2026},
      eprint={2602.08968},
      archivePrefix={arXiv},
      primaryClass={cs.AI},
      url={https://arxiv.org/abs/2602.08968},
}