Simple, Good, Fast: Self-Supervised World Models Free of Baggage

2025-06-03

International Conference on Learning Representations

AI Summary

Can world models remain effective without RNNs, Transformers, discrete representations, or explicit image reconstruction? This paper systematically investigates the essential components of lightweight world models, proposing a “simplicity-first” paradigm: it relies solely on self-supervised representation learning, frame- and action-stacking to model short-term dynamics, and reinforcement-learning-driven data augmentation to enhance robustness to prediction errors. All modules are fully differentiable end-to-end, eliminating the need for complex architectures or explicit pixel-level reconstruction. Evaluated on the Atari 100k benchmark, the approach achieves state-of-the-art performance. Ablation studies confirm the distinct contributions of each component. Training speed improves by 3×, and inference latency decreases by 60%, significantly reducing computational and implementation overhead. This work provides the first empirical validation that high-performance world modeling is feasible under an extremely minimal design.

Technology Category

• Computer Vision: Diffusion Models for Vision
• Machine Learning: Reinforcement Learning
• Search and Optimization: Learning to Search

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Responsible Web: Machine-in-the-loop, human agency and autonomy
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs

Abstract

What are the essential components of world models? How far do we get with world models that are not employing RNNs, transformers, discrete representations, and image reconstructions? This paper introduces SGF, a Simple, Good, and Fast world model that uses self-supervised representation learning, captures short-time dependencies through frame and action stacking, and enhances robustness against model errors through data augmentation. We extensively discuss SGF's connections to established world models, evaluate the building blocks in ablation studies, and demonstrate good performance through quantitative comparisons on the Atari 100k benchmark.

Problem

• Essential components of world models
• Performance of non-RNN, non-transformer world models
• Robustness enhancement through data augmentation

Innovation

• Self-supervised representation learning
• Frame and action stacking
• Data augmentation for robustness