Stable World Model provides a streamlined framework for conducting world model research with reproducible data collection, flexible model training, and comprehensive evaluation tools. Built on top of Gymnasium, it offers vectorized environments, domain randomization, and integrated support for multiple planning algorithms.
- Python >= 3.10
- CUDA-compatible GPU (recommended for training)
Using uv (recommended):
# Install uv
pip install uv
# Clone and install
git clone https://github.com/rbalestr-lab/stable-worldmodel.git
cd stable-worldmodel
uv pip install -e .Using pip:
git clone https://github.com/rbalestr-lab/stable-worldmodel.git
cd stable-worldmodel
pip install -e .For contributors and researchers developing new features:
uv pip install -e ".[dev,docs]"This includes testing tools (pytest, coverage) and documentation generators (sphinx).
stable_worldmodel/
├── envs/ # Gymnasium environments
│ ├── pusht.py
│ ├── simple_point_maze.py
│ ├── two_room.py
│ └── ogbench_cube.py
├── solver/ # Planning algorithms
│ ├── cem.py # Cross-Entropy Method
│ ├── mppi.py # Model Predictive Path Integral
│ ├── gd.py # Gradient Descent
│ └── nevergrad.py # Nevergrad
├── wm/ # World model architectures
│ ├── dinowm.py # DINO World Model
│ ├── dreamer.py # Dreamer
│ └── tdmpc.py # Temporal Difference MPC
├── policy.py
├── spaces.py # Extended Gymnasium spaces
├── world.py
├── data.py
└── utils.py
We maintain high test coverage to ensure reliability:
# Run all tests
pytest
# Run with coverage report
pytest --cov=stable_worldmodel --cov-report=term-missingThis project is licensed under the MIT License - see the LICENSE file for details.