stable-worldmodel 0.0.6

World Model Research Made Simple

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), gradient-based (GradientSolver, PGD), and constrained gradient approaches (LagrangianSolver).

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 Improved CEM (iCEM) Sampling Model Predictive Path Integral (MPPI) Sampling Gradient Descent (SGD, Adam) Gradient Projected Gradient Descent (PGD) Gradient Augmented Lagrangian Constrained Opt Baselines Type DINO-WM JEPA GCBC Behaviour Cloning HILP RL GCIVL RL GCIQL RL

CLI

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

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