Towards Unraveling and Improving Generalization in World Models

TL;DR

This paper investigates the robustness and generalization of world models in reinforcement learning, introducing a stochastic differential equation framework and Jacobian regularization to improve stability and performance.

Contribution

It develops a novel stochastic dynamical system formulation for world models and proposes a Jacobian regularization method to enhance robustness and training stability.

Findings

Latent representation errors can act as implicit regularizers with zero drift.

Jacobian regularization improves training stability and accelerates convergence.

Regularization enhances long-horizon prediction accuracy.

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.