# Causal State Distillation for Explainable Reinforcement Learning

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

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 28 of its 55 references that are in the map.
- First posted on arXiv 30 Dec 2023.

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

- [Learning Transferable Visual Models From Natural Language Supervision](https://mlsanity.com/paper/2103.00020/) (2021)
- [Representation Learning with Contrastive Predictive Coding](https://mlsanity.com/paper/1807.03748/) (2018)
- ["Why Should I Trust You?": Explaining the Predictions of Any Classifier](https://mlsanity.com/paper/1602.04938/) (2016)
- [Axiomatic Attribution for Deep Networks](https://mlsanity.com/paper/1703.01365/) (2017)
- [Invariant Risk Minimization](https://mlsanity.com/paper/1907.02893/) (2019)
- [Deep Reinforcement Learning with Double Q-learning](https://mlsanity.com/paper/1509.06461/) (2015)
- [Learning Important Features Through Propagating Activation Differences](https://mlsanity.com/paper/1704.02685/) (2017)
- [MINE: Mutual Information Neural Estimation](https://mlsanity.com/paper/1801.04062/) (2018)
- [Agent57: Outperforming the Atari Human Benchmark](https://mlsanity.com/paper/2003.13350/) (2020)
- [CLUB: A Contrastive Log-ratio Upper Bound of Mutual Information](https://mlsanity.com/paper/2006.12013/) (2020)
- [Towards Causal Representation Learning](https://mlsanity.com/paper/2102.11107/) (2021)
- [Explainable Reinforcement Learning Through a Causal Lens](https://mlsanity.com/paper/1905.10958/) (2019)
- [Visualizing and Understanding Atari Agents](https://mlsanity.com/paper/1711.00138/) (2017)
- [Causality Inspired Representation Learning for Domain Generalization](https://mlsanity.com/paper/2203.14237/) (2022)
- [Unsupervised Learning of Object Keypoints for Perception and Control](https://mlsanity.com/paper/1906.11883/) (2019)
- [Causal Interpretability for Machine Learning -- Problems, Methods and Evaluation](https://mlsanity.com/paper/2003.03934/) (2020)
- [A Survey on Causal Reinforcement Learning](https://mlsanity.com/paper/2302.05209/) (2023)
- [Generative causal explanations of black-box classifiers](https://mlsanity.com/paper/2006.13913/) (2020)
- [Transparency and Explanation in Deep Reinforcement Learning Neural Networks](https://mlsanity.com/paper/1809.06061/) (2018)
- [Counterfactual State Explanations for Reinforcement Learning Agents via Generative Deep Learning](https://mlsanity.com/paper/2101.12446/) (2021)
- [OrphicX: A Causality-Inspired Latent Variable Model for Interpreting Graph Neural Networks](https://mlsanity.com/paper/2203.15209/) (2022)
- [A Survey of Explainable Reinforcement Learning](https://mlsanity.com/paper/2202.08434/) (2022)
- [A Relational Intervention Approach for Unsupervised Dynamics Generalization in Model-Based Reinforcement Learning](https://mlsanity.com/paper/2206.04551/) (2022)
- [Internally Rewarded Reinforcement Learning](https://mlsanity.com/paper/2302.00270/) (2023)
- [Integrating Policy Summaries with Reward Decomposition for Explaining Reinforcement Learning Agents](https://mlsanity.com/paper/2210.11825/) (2022)
- [A Closer Look at Reward Decomposition for High-Level Robotic Explanations](https://mlsanity.com/paper/2304.12958/) (2023)
- [Learning "What-if" Explanations for Sequential Decision-Making](https://mlsanity.com/paper/2007.13531/) (2020)
- [Causal Explanation for Reinforcement Learning: Quantifying State and Temporal Importance](https://mlsanity.com/paper/2210.13507/) (2022)

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

- [Agentic AI framework for End-to-End Medical Data Inference](https://mlsanity.com/paper/2507.18115/) (2025)
- [Mental Modeling of Reinforcement Learning Agents by Language Models](https://mlsanity.com/paper/2406.18505/) (2024)

## Abstract (arXiv)

Reinforcement learning (RL) is a powerful technique for training intelligent agents, but understanding why these agents make specific decisions can be quite challenging. This lack of transparency in RL models has been a long-standing problem, making it difficult for users to grasp the reasons behind an agent's behaviour. Various approaches have been explored to address this problem, with one promising avenue being reward decomposition (RD). RD is appealing as it sidesteps some of the concerns associated with other methods that attempt to rationalize an agent's behaviour in a post-hoc manner. RD works by exposing various facets of the rewards that contribute to the agent's objectives during training. However, RD alone has limitations as it primarily offers insights based on sub-rewards and does not delve into the intricate cause-and-effect relationships that occur within an RL agent's neural model. In this paper, we present an extension of RD that goes beyond sub-rewards to provide more informative explanations. Our approach is centred on a causal learning framework that leverages information-theoretic measures for explanation objectives that encourage three crucial properties of causal factors: causal sufficiency, sparseness, and orthogonality. These properties help us distill the cause-and-effect relationships between the agent's states and actions or rewards, allowing for a deeper understanding of its decision-making processes. Our framework is designed to generate local explanations and can be applied to a wide range of RL tasks with multiple reward channels. Through a series of experiments, we demonstrate that our approach offers more meaningful and insightful explanations for the agent's action selections.