Towards Unraveling And Improving Generalization In World Models

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 drif

Stats

• citations0
• S2 citations—
• github stars0
• HF likes0
• heat score0.00
• arxiv keyfang2024towards

Related papers

• Dreaming Of Many Worlds: Learning Contextual World Models Aids Zero-shot Generalization (2024)2.83
• Learning To Predict Without Looking Ahead: World Models Without Forward Prediction (2019)0.00
• Prototypical Context-aware Dynamics Generalization For High-dimensional Model-based Reinforcement Learning (2022)0.00
• Procedural Generalization By Planning With Self-supervised World Models (2021)0.00
• Continual Visual Reinforcement Learning With A Life-long World Model (2023)2.26
• Language-conditioned World Model Improves Policy Generalization By Reading Environmental Descriptions (2025)0.00
• Understanding What Affects The Generalization Gap In Visual Reinforcement Learning: Theory And Empirical Evidence (2024)5.84
• The Effectiveness Of World Models For Continual Reinforcement Learning (2022)0.00