# Building Minimal and Reusable Causal State Abstractions for Reinforcement Learning

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

## Key facts

- Cited by 11 papers in the map (data through 14 Oct 2026).
- Placed in Deep RL algorithms and representations, one of 15 directions in Deep reinforcement learning methods.
- Builds on 23 of its 36 references that are in the map.
- First posted on arXiv 23 Jan 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, 23 of 23)

- [Representation Learning with Contrastive Predictive Coding](https://mlsanity.com/paper/1807.03748/) (2018)
- [Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor](https://mlsanity.com/paper/1801.01290/) (2018)
- [Categorical Reparameterization with Gumbel-Softmax](https://mlsanity.com/paper/1611.01144/) (2016)
- [Deep Reinforcement Learning in a Handful of Trials using Probabilistic Dynamics Models](https://mlsanity.com/paper/1805.12114/) (2018)
- [When to Trust Your Model: Model-Based Policy Optimization](https://mlsanity.com/paper/1906.08253/) (2019)
- [Reinforcement Learning with Deep Energy-Based Policies](https://mlsanity.com/paper/1702.08165/) (2017)
- [robosuite: A Modular Simulation Framework and Benchmark for Robot Learning](https://mlsanity.com/paper/2009.12293/) (2020)
- [Learning Invariant Representations for Reinforcement Learning without Reconstruction](https://mlsanity.com/paper/2006.10742/) (2020)
- [Sliced Score Matching: A Scalable Approach to Density and Score Estimation](https://mlsanity.com/paper/1905.07088/) (2019)
- [Neural Network Dynamics for Model-Based Deep Reinforcement Learning with Model-Free Fine-Tuning](https://mlsanity.com/paper/1708.02596/) (2017)
- [Model-Ensemble Trust-Region Policy Optimization](https://mlsanity.com/paper/1802.10592/) (2018)
- [Model-Based Value Estimation for Efficient Model-Free Reinforcement Learning](https://mlsanity.com/paper/1803.00101/) (2018)
- [Tianshou: a Highly Modularized Deep Reinforcement Learning Library](https://mlsanity.com/paper/2107.14171/) (2021)
- [Invariant Causal Prediction for Block MDPs](https://mlsanity.com/paper/2003.06016/) (2020)
- [Causal Dynamics Learning for Task-Independent State Abstraction](https://mlsanity.com/paper/2206.13452/) (2022)
- [Learning Task Informed Abstractions](https://mlsanity.com/paper/2106.15612/) (2021)
- [ContactNets: Learning Discontinuous Contact Dynamics with Smooth, Implicit Representations](https://mlsanity.com/paper/2009.11193/) (2020)
- [Deep Energy Estimator Networks](https://mlsanity.com/paper/1805.08306/) (2018)
- [Generalizing Goal-Conditioned Reinforcement Learning with Variational Causal Reasoning](https://mlsanity.com/paper/2207.09081/) (2022)
- [Action-Sufficient State Representation Learning for Control with Structural Constraints](https://mlsanity.com/paper/2110.05721/) (2021)
- [Learning Causal State Representations of Partially Observable Environments](https://mlsanity.com/paper/1906.10437/) (2019)
- [Decomposed Mutual Information Estimation for Contrastive Representation Learning](https://mlsanity.com/paper/2106.13401/) (2021)
- [Joint Energy-based Model Training for Better Calibrated Natural Language Understanding Models](https://mlsanity.com/paper/2101.06829/) (2021)

## Built on by (most cited first, 11 of 11)

- [Towards Empowerment Gain through Causal Structure Learning in Model-Based RL](https://mlsanity.com/paper/2502.10077/) (2025)
- [Curious Causality-Seeking Agents Learn Meta Causal World](https://mlsanity.com/paper/2506.23068/) (2025)
- [Why Online Reinforcement Learning is Causal](https://mlsanity.com/paper/2403.04221/) (2024)
- [On the Sample Efficiency of Abstractions and Potential-Based Reward Shaping in Reinforcement Learning](https://mlsanity.com/paper/2404.07826/) (2024)
- [Dynamic Sparsity: Challenging Common Sparsity Assumptions for Learning World Models in Robotic Reinforcement Learning Benchmarks](https://mlsanity.com/paper/2511.08086/) (2025)
- [Learning telic-controllable state representations](https://mlsanity.com/paper/2406.14476/) (2024)
- [Context-Sensitive Abstractions for Reinforcement Learning with Parameterized Actions](https://mlsanity.com/paper/2512.20831/) (2025)
- [From Causal Graphs to Markovian States: Multi-Order State Exposure for Deep RL](https://mlsanity.com/paper/2605.07057/) (2026)
- [SAT-RTS: A systematic framework for tactical knowledge extraction and visualization-based analysis in real-time strategy games](https://mlsanity.com/paper/2606.30090/) (2026)
- [Property-driven Causal Abstractions for Markov Decision Processes](https://mlsanity.com/paper/2607.26787/) (2026)
- [A Unifying Perspective on Causal World Models: From Observations to Representations to Structure](https://mlsanity.com/paper/2608.13456/) (2026)

## Abstract (arXiv)

Two desiderata of reinforcement learning (RL) algorithms are the ability to learn from relatively little experience and the ability to learn policies that generalize to a range of problem specifications. In factored state spaces, one approach towards achieving both goals is to learn state abstractions, which only keep the necessary variables for learning the tasks at hand. This paper introduces Causal Bisimulation Modeling (CBM), a method that learns the causal relationships in the dynamics and reward functions for each task to derive a minimal, task-specific abstraction. CBM leverages and improves implicit modeling to train a high-fidelity causal dynamics model that can be reused for all tasks in the same environment. Empirical validation on manipulation environments and Deepmind Control Suite reveals that CBM's learned implicit dynamics models identify the underlying causal relationships and state abstractions more accurately than explicit ones. Furthermore, the derived state abstractions allow a task learner to achieve near-oracle levels of sample efficiency and outperform baselines on all tasks.