Towards Unraveling and Improving Generalization in World Models

TL;DR

This paper reframes world-model–based RL as a stochastic dynamical system to analyze how latent representation errors affect robustness and generalization. It shows that zero-drift latent errors can act as implicit regularization, potentially improving performance, while non-zero drift introduces bias that can destabilize learning, which can be mitigated via Jacobian regularization on the latent dynamics. The authors derive explicit theoretical expressions for regularization terms and provide bounds on error propagation during predictive rollouts, then validate the approach with extensive Mujoco experiments. Overall, Jacobian regularization enhances training stability, accelerates convergence, and improves long-horizon prediction accuracy, advancing the reliability of world models in perturbed or unseen environments.

Abstract

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.

Figures (9)

• Figure 1: Generalization against increasing degree of perturbation.
• Figure 2: Latent Representation Problem: The left and right denote the manifold $\mathcal{M}$ with lower dim $d_\mathcal{M}$ embedded in a larger Euclidean space, with latent space $Z$ a $d_\mathcal{M}$-dimensional ball in middle. Encoder and decoder as maps respectively pushing forward Q to P and P to Q.
• Figure 3: Walker. Eval (left) and train scores (right) under latent error process $\mu_t = 0, \sigma_t^2=5$
• Figure 4: Walker. Eval (left) and train scores (right) under latent error process $\mu_t \sim [0, 5], \sigma_t^2 \sim [0,5]$.
• Figure 5: Quad. Eval (left) and train scores (right) under latent error process $\mu_t = 0, \sigma_t^2=0.1$.

Theorems & Definitions (24)

• Theorem 3.3
• Corollary 3.4
• Theorem 4.1
• Corollary 4.2
• Definition A.1: $\mathcal{C}^{k,\alpha}$ manifold
• Definition A.2: Riemannian volume form
• Theorem A.6
• Proposition A.7: $\mathcal{C}^{k, \alpha}$, compact
• proof
• proof