Learning Latent Dynamic Robust Representations for World Models

TL;DR

This paper introduces a novel approach combining spatio-temporal masking, bisimulation, and latent reconstruction within a hybrid recurrent state-space model to improve visual model-based reinforcement learning in noisy environments.

Contribution

It proposes a new latent dynamic robust representation method and a hybrid recurrent state-space model to enhance world model robustness and performance in complex visual tasks.

Findings

Significant performance improvements in visually complex control tasks.

Effective filtering of irrelevant noise in environment observations.

Enhanced stability in joint training of representations, dynamics, and policy.

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 https://github.com/bit1029public/HRSSM.