# Towards Unraveling and Improving Generalization in World Models

arXiv: https://arxiv.org/abs/2501.00195 · HTML: https://mlsanity.com/paper/2501.00195/

## Key facts

- Cited by 4 papers in the map (data through 14 Oct 2026).
- Placed in Latent predictive world models, one of 11 directions in World models.
- Builds on 26 of its 45 references that are in the map.
- First posted on arXiv 31 Dec 2024.

## Placed in

- World models / [Latent predictive world models](https://mlsanity.com/wm/latent-predictive-world-models/)
- Deep reinforcement learning methods / [Deep RL algorithms and representations](https://mlsanity.com/deeprl/deep-rl-algorithms-and-representations/)

## Builds on (most cited first, 26 of 26)

- [Dream to Control: Learning Behaviors by Latent Imagination](https://mlsanity.com/paper/1912.01603/) (2019)
- [World Models](https://mlsanity.com/paper/1803.10122/) (2018)
- [On Large-Batch Training for Deep Learning: Generalization Gap and Sharp Minima](https://mlsanity.com/paper/1609.04836/) (2016)
- [Learning Latent Dynamics for Planning from Pixels](https://mlsanity.com/paper/1811.04551/) (2018)
- [Mastering Diverse Domains through World Models](https://mlsanity.com/paper/2301.04104/) (2023)
- [GAIA-1: A Generative World Model for Autonomous Driving](https://mlsanity.com/paper/2309.17080/) (2023)
- [Tutorial on Variational Autoencoders](https://mlsanity.com/paper/1606.05908/) (2016)
- [DayDreamer: World Models for Physical Robot Learning](https://mlsanity.com/paper/2206.14176/) (2022)
- [A Survey of Zero-shot Generalisation in Deep Reinforcement Learning](https://mlsanity.com/paper/2111.09794/) (2021)
- [Symplectic Recurrent Neural Networks](https://mlsanity.com/paper/1909.13334/) (2019)
- [Robust Learning with Jacobian Regularization](https://mlsanity.com/paper/1908.02729/) (2019)
- [AntisymmetricRNN: A Dynamical System View on Recurrent Neural Networks](https://mlsanity.com/paper/1902.09689/) (2019)
- [Robust Reinforcement Learning using Offline Data](https://mlsanity.com/paper/2208.05129/) (2022)
- [The Implicit and Explicit Regularization Effects of Dropout](https://mlsanity.com/paper/2002.12915/) (2020)
- [Improving Generalization in Meta-learning via Task Augmentation](https://mlsanity.com/paper/2007.13040/) (2020)
- [Noisy Recurrent Neural Networks](https://mlsanity.com/paper/2102.04877/) (2021)
- [On the Robustness of Safe Reinforcement Learning under Observational Perturbations](https://mlsanity.com/paper/2205.14691/) (2022)
- [Explicit Regularisation in Gaussian Noise Injections](https://mlsanity.com/paper/2007.07368/) (2020)
- [Natural Actor-Critic for Robust Reinforcement Learning with Function Approximation](https://mlsanity.com/paper/2307.08875/) (2023)
- [Theory of Deep Learning III: explaining the non-overfitting puzzle](https://mlsanity.com/paper/1801.00173/) (2017)
- [Active World Model Learning with Progress Curiosity](https://mlsanity.com/paper/2007.07853/) (2020)
- [The Effectiveness of World Models for Continual Reinforcement Learning](https://mlsanity.com/paper/2211.15944/) (2022)
- [Learning to Predict Without Looking Ahead: World Models Without Forward Prediction](https://mlsanity.com/paper/1910.13038/) (2019)
- [Combining Pessimism with Optimism for Robust and Efficient Model-Based Deep Reinforcement Learning](https://mlsanity.com/paper/2103.10369/) (2021)
- [Improving Generalization in Meta-RL with Imaginary Tasks from Latent Dynamics Mixture](https://mlsanity.com/paper/2105.13524/) (2021)
- [Exploring the Training Robustness of Distributional Reinforcement Learning against Noisy State Observations](https://mlsanity.com/paper/2109.08776/) (2021)

## Built on by (most cited first, 4 of 4)

- [SmallWorlds: Assessing Dynamics Understanding of World Models in Isolated Environments](https://mlsanity.com/paper/2511.23465/) (2025)
- [Cloning Deterministic Worlds: The Critical Role of Latent Geometry in Long-Horizon World Models](https://mlsanity.com/paper/2510.26782/) (2025)
- [From Digital Twins to World Models:Opportunities, Challenges, and Applications for Mobile Edge General Intelligence](https://mlsanity.com/paper/2603.17420/) (2026)
- [When Low Prediction Error Misleads Planning: Diagnosing Representation, Dynamics, and Decision Failures in Latent World Models](https://mlsanity.com/paper/2610.05550/) (2026)

## Abstract (arXiv)

World models have recently emerged as a promising approach to reinforcement learning (RL), achieving state-of-the-art performance across a wide range of visual control tasks. This work aims to obtain a deep understanding of the robustness and generalization capabilities of world models. Thus motivated, we develop a stochastic differential equation formulation by treating the world model learning as a stochastic dynamical system, and characterize the impact of latent representation errors on robustness and generalization, for both cases with zero-drift representation errors and with non-zero-drift representation errors. Our somewhat surprising findings, based on both theoretic and experimental studies, reveal that for the case with zero drift, modest latent representation errors can in fact function as implicit regularization and hence result in improved robustness. We further propose a Jacobian regularization scheme to mitigate the compounding error propagation effects of non-zero drift, thereby enhancing training stability and robustness. Our experimental studies corroborate that this regularization approach not only stabilizes training but also accelerates convergence and improves accuracy of long-horizon prediction.