Learning World Models with Identifiable Factorization

TL;DR

This paper introduces IFactor, a framework for extracting identifiable, disentangled latent representations in reinforcement learning, leading to more stable, compact, and effective world models for policy learning.

Contribution

We propose a novel framework that models four categories of latent variables with block-wise identifiability, improving representation stability and policy performance.

Findings

Accurately identifies ground-truth latent variables in synthetic environments.

Demonstrates superior performance over baselines in DeepMind Control Suite variants.

Ensures minimal and sufficient information retention for policy optimization.

Abstract

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.