Leveraging Conditional Dependence for Efficient World Model Denoising

De-Chuan Zhan (Nanjing University) · Shaowei Zhang (Nanjing University) · Jiahan Cao (Nanjing University) · Dian Cheng (Independent Researcher) · Xunlan Zhou (Nanjing University) · Shenghua Wan (Nanjing University) · Le Gan (Nanjing University of Science and Technology)

collider-structure recurrent state-space modelconditional dependencecsrssmdecoupling regularizationdenoising experimental results inference process information extraction learning efficiency model-based reinforcement learningnoise interferencerolloutsstate space inferencetask-irrelevant componentstask-relevant components

Effective denoising is critical for managing complex visual inputs contaminated with noisy distractors in model-based reinforcement learning (RL). Current methods often oversimplify the decomposition of observations by neglecting the conditional dependence between task-relevant and task-irrelevant components given an observation. To address this limitation, we introduce CsDreamer, a model-based RL approach built upon the world model of Collider-structure Recurrent State-Space Model (CsRSSM). CsRSSM incorporates colliders to comprehensively model the denoising inference process and explicitly capture the conditional dependence. Furthermore, it employs a decoupling regularization to balance the influence of this conditional dependence. By accurately inferring a task-relevant state space, CsDreamer improves learning efficiency during rollouts. Experimental results demonstrate the effectiveness of CsRSSM in extracting task-relevant information, leading to CsDreamer outperforming existing approaches in environments characterized by complex noise interference.