Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

From Pixels to Factors: Learning Independently Controllable State Variables for Reinforcement Learning

Rafael Rodriguez-Sanchez, Cameron Allen, George Konidaris

TL;DR

This work addresses the challenge of learning factored representations from high-dimensional observations when the ground-truth factors are not observed. It introduces Action Controllable Factorization (ACF), a contrastive, energy-based method that isolates independently controllable latent variables by contrasting action-driven transitions against natural dynamics, leveraging a sparse-action assumption and a no-op baseline. The approach combines an energy-parametrized forward model, an inverse dynamics objective, and ratio-based classifiers to align latent factors with controllable state components, with identifiability supported under sparsity and connectivity assumptions. Empirically, ACF recovers ground-truth controllable factors directly from pixels in Taxi, FourRooms, and MiniGrid-DoorKey and outperforms standard disentanglement baselines, indicating potential for improved sample efficiency in factored RL and world-model learning.

Abstract

Algorithms that exploit factored Markov decision processes are far more sample-efficient than factor-agnostic methods, yet they assume a factored representation is known a priori -- a requirement that breaks down when the agent sees only high-dimensional observations. Conversely, deep reinforcement learning handles such inputs but cannot benefit from factored structure. We address this representation problem with Action-Controllable Factorization (ACF), a contrastive learning approach that uncovers independently controllable latent variables -- state components each action can influence separately. ACF leverages sparsity: actions typically affect only a subset of variables, while the rest evolve under the environment's dynamics, yielding informative data for contrastive training. ACF recovers the ground truth controllable factors directly from pixel observations on three benchmarks with known factored structure -- Taxi, FourRooms, and MiniGrid-DoorKey -- consistently outperforming baseline disentanglement algorithms.

From Pixels to Factors: Learning Independently Controllable State Variables for Reinforcement Learning

TL;DR

This work addresses the challenge of learning factored representations from high-dimensional observations when the ground-truth factors are not observed. It introduces Action Controllable Factorization (ACF), a contrastive, energy-based method that isolates independently controllable latent variables by contrasting action-driven transitions against natural dynamics, leveraging a sparse-action assumption and a no-op baseline. The approach combines an energy-parametrized forward model, an inverse dynamics objective, and ratio-based classifiers to align latent factors with controllable state components, with identifiability supported under sparsity and connectivity assumptions. Empirically, ACF recovers ground-truth controllable factors directly from pixels in Taxi, FourRooms, and MiniGrid-DoorKey and outperforms standard disentanglement baselines, indicating potential for improved sample efficiency in factored RL and world-model learning.

Abstract

Algorithms that exploit factored Markov decision processes are far more sample-efficient than factor-agnostic methods, yet they assume a factored representation is known a priori -- a requirement that breaks down when the agent sees only high-dimensional observations. Conversely, deep reinforcement learning handles such inputs but cannot benefit from factored structure. We address this representation problem with Action-Controllable Factorization (ACF), a contrastive learning approach that uncovers independently controllable latent variables -- state components each action can influence separately. ACF leverages sparsity: actions typically affect only a subset of variables, while the rest evolve under the environment's dynamics, yielding informative data for contrastive training. ACF recovers the ground truth controllable factors directly from pixel observations on three benchmarks with known factored structure -- Taxi, FourRooms, and MiniGrid-DoorKey -- consistently outperforming baseline disentanglement algorithms.

Paper Structure

This paper contains 14 sections, 3 theorems, 16 equations, 8 figures, 6 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Background
  3. ACF: Action Controllable Factorization
  4. Problem Formulation
  5. Algorithm
  6. Evaluation
  7. Related Works
  8. Conclusion
  9. Theoretical Results
  10. Extended Empirical Results
  11. Network Architectures
  12. Domains
  13. Hyperparameters, Tuning and Computational Resources
  14. Detailed Empirical Results

Key Result

Theorem 3.1

Let the learned encoder $f: X\rightarrow Z$ be a diffeomorphism. If the following conditions hold then, there exists a factor-wise diffeomorphism $h : \mathcal{S} \rightarrow Z$ between the underlying ground truth factors of variation $\mathcal{S}$ and the learned encoding $Z$

Figures (8)

  • Figure 1: Factorization metrics. The left side bars show how much of the information is represented off the diagonal on average over all variables. The right side bars represent the mean diagonal value. Ideally, we would expect our $R^2$ matrices to be close to the identity: $0$ on the left bar, $1$ on the right bar. The error bars show the standard deviation over $5$ independent seeds.
  • Figure 2: Factorization matrices for DoorKey. Mean $R^2$ matrices over $5$ seeds.
  • Figure 3: Taxi latent traversals. In this Taxi rendering, the taxi is represented by a hollow square, the passengers are circles with colors matching their goal positions. When a passenger is in the taxi, the border of the frame is highlighted with stripes. By varying the value of a latent variable (columns), we can see its effect on the mean observation. Each row represents different latent variables.
  • Figure 4: DoorKey latent traversals. For this domain, we show a random sample from observations that have a particular value of the latent dimension. We only show the controllable elements in DoorKey, that includes the agent position and orientation, the key and the door state. Different rows correspond to different latent variables and different columns represent different values for the corresponding latent variable.
  • Figure 5: DoorKey Factorization Results
  • ...and 3 more figures

Theorems & Definitions (6)

  • Theorem 3.1: Identifiability of the Independently Controllable Factors
  • Definition A.1: Identifiability
  • Lemma A.2: Local Identifiability of the Independently Controllable Factors
  • proof
  • Proposition A.3: Global Identifiability of Independently Controllable Factors
  • proof