Identifying Transition-Independent State Variables for Factorized World Models

Learning a compact world model representation with factorized state variables is crucial for efficient reinforcement learning. However, existing approaches of identifiable factorization for Markov decision processes (MDPs) are either constrained to block-wise identifiability or rely on the limiting condition of control-independent factors. To extend the capability of identifiable factorization in MDP, this paper proposes Transition-Independent Factorization (TIF), a technique that can further identify factors that have dependent control influence but independent transition dynamics. We first present a theoretical analysis to establish the identifiability of controllable and transition-independent factors, followed by a practical approach to learn factorized world models from high-dimensional observations. Experiments in synthetic worlds and pixel-based grid worlds validate that TIF can accurately identify the ground-truth factors with independent transition dynamics. Experiments in variants of the MiniGrid Empty and Four-Rooms demonstrate that the resulting factorized world model representations can bring substantial benefits to the planning and policy learning in reinforcement learning.