stable-worldmodel

World model research made simple. From data collection to training and evaluation.

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

Quick Example

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), as well as gradient-based approaches.

Efficiency

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.

Supported Environments

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.

Environment ID # FoV swm/PushT-v1 16 swm/TwoRoom-v1 17 swm/OGBCube-v0 11 swm/OGBScene-v0 12 swm/HumanoidDMControl-v0 7 swm/CheetahDMControl-v0 7 swm/HopperDMControl-v0 7 swm/ReacherDMControl-v0 8 swm/WalkerDMControl-v0 8 swm/AcrobotDMControl-v0 8 swm/PendulumDMControl-v0 6 swm/CartpoleDMControl-v0 6 swm/BallInCupDMControl-v0 9 swm/FingerDMControl-v0 10 swm/ManipulatorDMControl-v0 8 swm/QuadrupedDMControl-v0 7

Solver Type Cross-Entropy Method (CEM) Sampling Model Predictive Path Integral (MPPI) Sampling Gradient Descent (SGD, Adam) Gradient Projected Gradient Descent (PGD) Gradient Baselines Type DINO-WM JEPA GCBC Behaviour Cloning HILP RL GCIVL RL GCIQL RL

Installing stable-worldmodel

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

pip install stable-worldmodel

Note: The library is still in active development.

Install from Source

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.

Questions

If you have a question, please file an issue.

Citation

@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},
}