Learning World Models with Identifiable Factorization

Extracting a stable and compact representation of the environment is crucial for efficient reinforcement learning in high-dimensional, noisy, and non-stationary environments. Different categories of information coexist in such environments -- how to effectively extract and disentangle the information remains a challenging problem. In this paper, we propose IFactor, a general framework to model four distinct categories of latent state variables that capture various aspects of information within the RL system, based on their interactions with actions and rewards. Our analysis establishes block-wise identifiability of these latent variables, which not only provides a stable and compact representation but also discloses that all reward-relevant factors are significant for policy learning. We further present a practical approach to learning the world model with identifiable blocks, ensuring the removal of redundancies but retaining minimal and sufficient information for policy optimization. Experiments in synthetic worlds demonstrate that our method accurately identifies the ground-truth latent variables, substantiating our theoretical findings. Moreover, experiments in variants of the DeepMind Control Suite and RoboDesk showcase the superior performance of our approach over baselines.

Authors

• Yuren Liu
• Biwei Huang University of California San Diego
• Zhengmao Zhu
• Honglong Tian
• Mingming Gong University of Melbourne Mohamed bin Zayed University of Artificial Intelligence
• Yang Yu School of Artificial Intelligence, Jilin University, Changchun Engineering Research Center of Knowledge-Driven Human-Machine Intelligence, Jilin University, Changchun
• Kun Zhang MBZUAI, Abu Dhabi, UAE, Carnegie Mellon University, Pittsburg, USA

Keywords

No keywords are indexed for this paper.