# Rethinking State Disentanglement in Causal Reinforcement Learning

arXiv: https://arxiv.org/abs/2408.13498 · HTML: https://mlsanity.com/paper/2408.13498/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Deep RL algorithms and representations, one of 15 directions in Deep reinforcement learning methods.
- Builds on 22 of its 33 references that are in the map.
- First posted on arXiv 24 Aug 2024.

## Placed in

- Deep reinforcement learning methods / [Deep RL algorithms and representations](https://mlsanity.com/deeprl/deep-rl-algorithms-and-representations/)

## Builds on (most cited first, 22 of 22)

- [Dream to Control: Learning Behaviors by Latent Imagination](https://mlsanity.com/paper/1912.01603/) (2019)
- [Mastering Diverse Domains through World Models](https://mlsanity.com/paper/2301.04104/) (2023)
- [Mastering Atari with Discrete World Models](https://mlsanity.com/paper/2010.02193/) (2020)
- [CURL: Contrastive Unsupervised Representations for Reinforcement Learning](https://mlsanity.com/paper/2004.04136/) (2020)
- [Variational Autoencoders and Nonlinear ICA: A Unifying Framework](https://mlsanity.com/paper/1907.04809/) (2019)
- [dm_control: Software and Tasks for Continuous Control](https://mlsanity.com/paper/2006.12983/) (2020)
- [Planning to Explore via Self-Supervised World Models](https://mlsanity.com/paper/2005.05960/) (2020)
- [Embed to Control: A Locally Linear Latent Dynamics Model for Control from Raw Images](https://mlsanity.com/paper/1506.07365/) (2015)
- [Provably efficient RL with Rich Observations via Latent State Decoding](https://mlsanity.com/paper/1901.09018/) (2019)
- [Kinematic State Abstraction and Provably Efficient Rich-Observation Reinforcement Learning](https://mlsanity.com/paper/1911.05815/) (2019)
- [Invariant Causal Prediction for Block MDPs](https://mlsanity.com/paper/2003.06016/) (2020)
- [Identifiability Guarantees for Causal Disentanglement from Soft Interventions](https://mlsanity.com/paper/2307.06250/) (2023)
- [DreamerPro: Reconstruction-Free Model-Based Reinforcement Learning with Prototypical Representations](https://mlsanity.com/paper/2110.14565/) (2021)
- [On Oracle-Efficient PAC RL with Rich Observations](https://mlsanity.com/paper/1803.00606/) (2018)
- [Denoised MDPs: Learning World Models Better Than the World Itself](https://mlsanity.com/paper/2206.15477/) (2022)
- [Learning Task Informed Abstractions](https://mlsanity.com/paper/2106.15612/) (2021)
- [From Pixels to Torques: Policy Learning with Deep Dynamical Models](https://mlsanity.com/paper/1502.02251/) (2015)
- [Identifying Weight-Variant Latent Causal Models](https://mlsanity.com/paper/2208.14153/) (2022)
- [Action-Sufficient State Representation Learning for Control with Structural Constraints](https://mlsanity.com/paper/2110.05721/) (2021)
- [Robust Predictable Control](https://mlsanity.com/paper/2109.03214/) (2021)
- [Provable RL with Exogenous Distractors via Multistep Inverse Dynamics](https://mlsanity.com/paper/2110.08847/) (2021)
- [Identifiable Latent Polynomial Causal Models Through the Lens of Change](https://mlsanity.com/paper/2310.15580/) (2023)

## Abstract (arXiv)

One of the significant challenges in reinforcement learning (RL) when dealing with noise is estimating latent states from observations. Causality provides rigorous theoretical support for ensuring that the underlying states can be uniquely recovered through identifiability. Consequently, some existing work focuses on establishing identifiability from a causal perspective to aid in the design of algorithms. However, these results are often derived from a purely causal viewpoint, which may overlook the specific RL context. We revisit this research line and find that incorporating RL-specific context can reduce unnecessary assumptions in previous identifiability analyses for latent states. More importantly, removing these assumptions allows algorithm design to go beyond the earlier boundaries constrained by them. Leveraging these insights, we propose a novel approach for general partially observable Markov Decision Processes (POMDPs) by replacing the complicated structural constraints in previous methods with two simple constraints for transition and reward preservation. With the two constraints, the proposed algorithm is guaranteed to disentangle state and noise that is faithful to the underlying dynamics. Empirical evidence from extensive benchmark control tasks demonstrates the superiority of our approach over existing counterparts in effectively disentangling state belief from noise.