Metadata-Version: 2.4
Name: stable-worldmodel
Version: 0.0.4
Summary: World Model Research Made Simple
Author: Lucas Maes, Quentin Le Lidec, Randall Balestriero
License-Expression: MIT
Classifier: Intended Audience :: Science/Research
Classifier: License :: OSI Approved :: MIT License
Classifier: Programming Language :: Python :: 3
Requires-Dist: torch
Requires-Dist: torchvision
Requires-Dist: numpy
Requires-Dist: loguru
Requires-Dist: tabulate
Requires-Dist: gymnasium
Requires-Dist: einops
Requires-Dist: h5py
Requires-Dist: hdf5plugin
Requires-Dist: tqdm
Requires-Dist: gdown
Requires-Dist: pygame ; extra == 'env'
Requires-Dist: pymunk ; extra == 'env'
Requires-Dist: shapely ; extra == 'env'
Requires-Dist: ogbench ; extra == 'env'
Requires-Dist: minigrid ; extra == 'env'
Requires-Dist: gymnasium[all] ; extra == 'env'
Requires-Dist: opencv-python ; extra == 'env'
Requires-Dist: stable-baselines3 ; extra == 'env'
Requires-Dist: transformers>=4.50.0 ; extra == 'train'
Requires-Dist: stable-pretraining ; extra == 'train'
Requires-Dist: hydra-core ; extra == 'train'
Requires-Dist: hydra-submitit-launcher ; extra == 'train'
Requires-Dist: wandb ; extra == 'train'
Maintainer: Lucas Maes
Maintainer-email: Lucas Maes <lucas_maes@brown.edu>
Requires-Python: >=3.10
Project-URL: Homepage, http://github.com/galilai-group/stable-worldmodel
Provides-Extra: env
Provides-Extra: train
Description-Content-Type: text/markdown

[![Documentation](https://img.shields.io/badge/Docs-blue.svg)](https://galilai-group.github.io/stable-worldmodel/)
![Tests](https://img.shields.io/github/actions/workflow/status/galilai-group/stable-worldmodel/tests.yaml?label=Tests)
[![PyPI](https://img.shields.io/pypi/v/stable-worldmodel.svg)](https://pypi.python.org/pypi/stable-worldmodel/#history)
[![PyTorch](https://img.shields.io/badge/PyTorch-ee4c2c?logo=pytorch&logoColor=white)](https://pytorch.org/get-started/locally/)
[![Ruff](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/ruff/main/assets/badge/v2.json)](https://github.com/astral-sh/ruff)

# stable-worldmodel

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

```bash
pip install stable-worldmodel
```

> **Note:** The library is still in active development.

See the full documentation at [here](https://galilai-group.github.io/stable-worldmodel/).


## Quick Example

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

## Supported Environments

<p align="center">
  <img src="docs/assets/envs_grid_1.gif" alt="Environments Grid 1">
  <br>
  <img src="docs/assets/envs_grid_2.gif" alt="Environments Grid 2">
</p>

## Contributing

Setup your codebase:

```bash
uv venv --python=3.10
source .venv/bin/activate
uv sync --all-extras --group dev
```

## Questions

If you have a question, please [file an issue](https://github.com/galilai-group/stable-worldmodel/issues).


## Citation

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