Saved in:
Bibliographic Details
Main Authors: Fang, Qiaoyi, Du, Weiyu, Wang, Hang, Zhang, Junshan
Format: Preprint
Published: 2024
Subjects:
Online Access:https://arxiv.org/abs/2501.00195
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866929653757771776
author Fang, Qiaoyi
Du, Weiyu
Wang, Hang
Zhang, Junshan
author_facet Fang, Qiaoyi
Du, Weiyu
Wang, Hang
Zhang, Junshan
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2501_00195
institution arXiv
publishDate 2024
record_format arxiv
spellingShingle Towards Unraveling and Improving Generalization in World Models
Fang, Qiaoyi
Du, Weiyu
Wang, Hang
Zhang, Junshan
Machine Learning
Artificial Intelligence
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.
title Towards Unraveling and Improving Generalization in World Models
topic Machine Learning
Artificial Intelligence
url https://arxiv.org/abs/2501.00195