Causal State Distillation for Explainable Reinforcement Learning

TL;DR

This work tackles explainability in reinforcement learning by introducing Causal State Distillation, a framework that learns latent causal factors $\alpha$ from states while separating non-causal components $\beta$, and ties them to actions $a$ and multi-channel rewards $r = \sum_i r^i$ via a structural causal model. It imposes causal sufficiency, sparsity, and orthogonality on the causal factors through an information-theoretic learning objective, enabling local explanations in the form of interpretable masks (R-Mask and Q-Mask) that reflect cause-driven attention to state components. The methodology includes interventions on $\beta$, metrics for causal intervention and sufficiency, sparsity and orthogonality constraints, and an optimization procedure that combines entropy, mutual information, and L1 penalties, with practical lite variants. Empirical results on Atari games (e.g., Gopher, MsPacman) and a Monster-Treasure toy demonstrate that causal factors yield meaningful, disentangled explanations with competitive performance, and qualitative masks illustrate the specific state components driving actions and sub-rewards. The work suggests broader applicability to counterfactual explanations and multi-domain RL, offering a principled alternative to post-hoc saliency by grounding explanations in the agent’s learning dynamics and causal structure.

Abstract

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.

Figures (21)

• Figure 1: The disentanglement of state representations and resulting sub-agents when uncovering the cause-effect relationships with causal state distillation (action omitted for brevity). Here, $s^0$ denotes the distilled non-causal components of state $s$, while $s^i$ captures the causal elements, each linked to a distinct reward aspect $r^i$. Sub-agents focus on a singular causal component for policy learning. The distillation process, consisting of multiple learning steps, is elaborated in Sec. \ref{['sec:learning framework']}.
• Figure 1: Evaluation results for RD, RD-pred, RD-pred-u.
• Figure 2: The causal graph for one-step RL explanations.
• Figure 3: Comparison of saliency maps (associated with ground and gopher rewards) of RD with RD-pred in a state where the agent filled the hole and attained reward 0.15. Q saliency refers to the generated saliency of Q-value; R saliency pertains to the generated saliency of reward.
• Figure 4: R-Mask attention masksfrom Gopher and their interpretation along with Q-value bars.