Towards Unraveling and Improving Generalization in World Models

Qiaoyi Fang, Weiyu Du, Hang Wang, Junshan Zhang

Keywords: world models, reinforcement learning, generalization

TL;DR: We analyzed the effects of latent representation errors on the generalization capacity of world models and improved it with Jacobian regularization.

Abstract: World model has recently emerged as a promising approach to reinforcement learning (RL), as evidenced by its great successes that world model based agents exhibit state-of-the-art performance on a wide range visual control tasks. In this study, we aim to first obtain a clear understanding of the generalization capability of world models by examining the impact of _latent representation error_, and then devise new methods to enhance its generalization. We hypothesize that latent representation errors may paradoxically bring generalization to the model. We develop a continuous-time stochastic dynamics framework to quantify the impact of these errors, by examining the regularization effects for both cases with zero-drift representation errors and non-zero-drift representation errors. We propose a Jacobian regularization scheme to mitigate the "destabilizing'' effects of non-zero drift errors, thereby enhancing training stability and model generalization. Our empirical results confirm that this regularization approach not only stabilizes training but also accelerates convergence and improves performance on long-horizon prediction.

Primary Area: Reinforcement learning

Submission Number: 13788