References & Citations
Computer Science > Machine Learning
Title: Learning Latent Dynamic Robust Representations for World Models
(Submitted on 10 May 2024 (v1), last revised 30 May 2024 (this version, v2))
Abstract: Visual Model-Based Reinforcement Learning (MBRL) promises to encapsulate agent's knowledge about the underlying dynamics of the environment, enabling learning a world model as a useful planner. However, top MBRL agents such as Dreamer often struggle with visual pixel-based inputs in the presence of exogenous or irrelevant noise in the observation space, due to failure to capture task-specific features while filtering out irrelevant spatio-temporal details. To tackle this problem, we apply a spatio-temporal masking strategy, a bisimulation principle, combined with latent reconstruction, to capture endogenous task-specific aspects of the environment for world models, effectively eliminating non-essential information. Joint training of representations, dynamics, and policy often leads to instabilities. To further address this issue, we develop a Hybrid Recurrent State-Space Model (HRSSM) structure, enhancing state representation robustness for effective policy learning. Our empirical evaluation demonstrates significant performance improvements over existing methods in a range of visually complex control tasks such as Maniskill \cite{gu2023maniskill2} with exogenous distractors from the Matterport environment. Our code is avaliable at this https URL
Submission history
From: Hongyu Zang [view email][v1] Fri, 10 May 2024 06:28:42 GMT (19471kb,D)
[v2] Thu, 30 May 2024 09:40:02 GMT (19965kb,D)
Link back to: arXiv, form interface, contact.