Overview

stable-worldmodel is a comprehensive platform engineered for reproducible world model research and evaluation. It streamlines the entire research workflow, from data collection and model training to robust evaluation using model-predictive control, all within a unified interface. This project provides a rich suite of standardized environments, including implementations of common baselines and planning solvers, allowing researchers to concentrate on their core contributions—the world model and its objective function.

Key Features

• Unified Research Pipeline: Seamlessly integrates data collection, world model training, and evaluation with Model-Predictive Control (MPC).
• Standardized Environments: Offers a vast collection of over 100 environments, including popular benchmarks from DeepMind Control Suite, Gymnasium, OGBench, Craftax, and the Arcade Learning Environment, many equipped with "factors of variation" for easy generalization testing.
• Flexible Data Handling: Supports multiple data formats (lance, hdf5, folder, video, lerobot) with efficient on-disk layouts and a pluggable registry for custom formats.
• Extensive Solvers & Baselines: Includes a variety of planning solvers (e.g., CEM, iCEM, MPPI) and reference implementations of leading baselines (e.g., LeWM, DINO-WM, GCBC).
• Reproducibility Focus: Built with reproducibility at its core, featuring reference implementations and a clear structure for research code.
• Command-Line Interface: Provides a convenient CLI for inspecting datasets, environments, and checkpoints without writing code.
• Active Development: The library is under active development with a clear roadmap and community support.

Typical Use Cases

• Accelerating World Model Research: Researchers can quickly set up experiments, leverage pre-built environments and baselines, and focus on developing novel world models and training objectives.
• Benchmarking and Comparison: Provides a standardized framework for comparing the performance of different world models and planning algorithms across a diverse set of tasks and visual variations.
• Evaluating Generalization: The integrated "factors of variation" in environments facilitate straightforward evaluation of zero-shot generalization capabilities to unseen visual or physical perturbations.
• Developing Robust Control Policies: Enables the training and evaluation of policies that utilize world models for planning and decision-making in complex environments.
• Educational Tool: Serves as an excellent resource for students and researchers looking to understand and contribute to the field of world modeling and model-based reinforcement learning.