stable-worldmodel: A Platform for Reproducible World Modeling Research and Evaluation

The evolution of world models has transitioned from classical control theories relying on analytical models to deep learning approaches that learn predictive representations directly from raw data. Early methods utilized hand-crafted simulators or explicit dynamics equations, which limited scalability and adaptability. Recent advances leverage neural networks, such as Variational Autoencoders, Transformers, and JEPA frameworks, to encode high-dimensional sensory inputs into latent states, enabling complex environment understanding. Notable works include PlaNet, Dreamer, and VQ-VAE-based models, which demonstrated success in control tasks. However, fragmentation in codebases, data handling inefficiencies, and inconsistent evaluation protocols hinder progress. The lack of standardized benchmarks and reproducibility remains a critical bottleneck, especially when comparing models across different environments and perturbations.

Current world modeling research faces significant challenges: disparate implementations lead to inconsistent results; data I/O bottlenecks slow down training and evaluation; and the absence of comprehensive, standardized benchmarks hampers fair comparison and progress assessment. These issues are compounded by the difficulty of evaluating models’ robustness and generalization, particularly under out-of-distribution conditions. As models are increasingly deployed in safety-critical applications like robotics and autonomous driving, ensuring reliability and reproducibility becomes essential. Addressing these bottlenecks requires a unified framework that streamlines data handling, standardizes evaluation, and supports diverse environments with controllable factors, enabling systematic analysis of model performance and limitations.

The platform introduces Lance as a high-speed, scalable data storage format, drastically reducing I/O bottlenecks. It integrates multiple cutting-edge world models and planning algorithms within a unified interface, supporting diverse environments with adjustable visual, geometric, and physical factors. The modular architecture separates data collection, model training, and evaluation, facilitating reproducibility and extensibility. It also standardizes benchmarking procedures, enabling fair comparison across models and environments. The environment suite spans control, robotics, and gaming, with controllable perturbations for robustness testing. These innovations collectively address fragmentation, data bottlenecks, and evaluation inconsistencies, setting a new standard for systematic world model research.

• �� Data collection: Using environment wrappers, generate trajectories via random, expert, or model-based policies, with controllable visual and physical variations. • Data storage: Convert collected data into Lance format for high throughput, supporting cloud streaming and local access. • Model training: Implement state-of-the-art models like DINO-WM (ViT-based), JEPA variants, and decoder-free models like TD-MPC2, with flexible training objectives. • Planning: Integrate solvers such as CEM and MPPI, which optimize action sequences by rolling out predicted trajectories, interfaced via a common get_cost method. • Evaluation: Perform systematic testing across multiple environments, introducing environment shifts and perturbations, measuring success rates and prediction errors. • Interaction: Environment wrappers support real-time interventions, enabling robustness and generalization assessments.

Experiments span control tasks (Push-T, MuJoCo), arcade games (Atari), and robotics environments, with datasets collected under various policies. Models are trained on identical datasets, then evaluated under standard and perturbed conditions, including visual and physical variations. Success metrics include success rate and prediction error, with hyperparameters tuned for fairness. Ablation studies compare data formats, model architectures, and planning algorithms, revealing the impact of data throughput, model complexity, and environment perturbations. Results demonstrate the platform’s ability to reproduce state-of-the-art performance and expose robustness gaps, guiding future improvements.

DINO-WM achieved 94% success in Push-T, outperforming LeWM at 72%. Data loading with Lance exceeded 4000 samples/sec, reducing I/O delays. Under environment shifts, success rates dropped sharply, e.g., from 94% to 30%, highlighting robustness issues. The platform's standardized evaluation revealed that current models struggle with out-of-distribution generalization, especially under visual and physical perturbations. These findings emphasize the importance of robustness testing and the need for models that can adapt to environment variations, guiding future research directions.

This platform supports robotics, autonomous navigation, and game AI development, providing a standardized testing environment for model robustness and generalization. It enables researchers and industry practitioners to evaluate models systematically, accelerating deployment in real-world scenarios. The ability to simulate environment variations and disturbances prepares models for unpredictable real-world conditions, making it valuable for safety-critical applications. Long-term, it fosters the development of more reliable, adaptable AI agents capable of operating across diverse and changing environments.

The current focus on simulation environments limits direct transferability to real-world scenarios due to sim-to-real gaps. Handling large-scale real robot data remains challenging, requiring further optimization. Models exhibit significant performance drops under severe environment shifts, indicating robustness limitations. Future work must address these issues by integrating real-world datasets, enhancing robustness techniques, and reducing computational costs for large-scale deployment.