Learning World Models with Identifiable Factorization
This work addresses the challenge of disentangling information in reinforcement learning environments, which is incremental as it builds on existing world model methods by adding identifiability guarantees.
The paper tackled the problem of extracting stable and compact representations in high-dimensional, noisy, and non-stationary environments for reinforcement learning by proposing IFactor, a framework that models four distinct categories of latent state variables with block-wise identifiability, demonstrating accurate identification of ground-truth variables in synthetic worlds and superior performance in DeepMind Control Suite and RoboDesk variants.
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 these 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 redundants 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.