Interpretable Imitation Learning with Dynamic Causal Relations

TL;DR

This work tackles the lack of interpretability in imitation learning by learning dynamic, causal graphs that relate state and action variables. It introduces CAIL, a three-component framework consisting of a dynamic causal discovery module (with DAG templates and Granger-causality-inspired learning), a causality-encoding module (edge-aware, multi-layer embeddings), and a prediction module trained end-to-end with adversarial imitation and an auxiliary regression objective. The approach demonstrates strong causal-discovery performance and competitive imitation accuracy on synthetic Kuramoto data and real-world MIMIC-IV and MiniGrid tasks, while producing meaningful DAGs that explain decision-making. By jointly optimizing causality discovery and policy learning, CAIL offers a scalable path to transparent, non-black-box imitation learning applicable to high-stakes domains.

Abstract

Imitation learning, which learns agent policy by mimicking expert demonstration, has shown promising results in many applications such as medical treatment regimes and self-driving vehicles. However, it remains a difficult task to interpret control policies learned by the agent. Difficulties mainly come from two aspects: 1) agents in imitation learning are usually implemented as deep neural networks, which are black-box models and lack interpretability; 2) the latent causal mechanism behind agents' decisions may vary along the trajectory, rather than staying static throughout time steps. To increase transparency and offer better interpretability of the neural agent, we propose to expose its captured knowledge in the form of a directed acyclic causal graph, with nodes being action and state variables and edges denoting the causal relations behind predictions. Furthermore, we design this causal discovery process to be state-dependent, enabling it to model the dynamics in latent causal graphs. Concretely, we conduct causal discovery from the perspective of Granger causality and propose a self-explainable imitation learning framework, {\method}. The proposed framework is composed of three parts: a dynamic causal discovery module, a causality encoding module, and a prediction module, and is trained in an end-to-end manner. After the model is learned, we can obtain causal relations among states and action variables behind its decisions, exposing policies learned by it. Experimental results on both synthetic and real-world datasets demonstrate the effectiveness of the proposed {\method} in learning the dynamic causal graphs for understanding the decision-making of imitation learning meanwhile maintaining high prediction accuracy.

Figures (8)

• Figure 1: Overview of CAIL for interpretable imitation learning from the causality perspective, which is composed of (a) causal discovery module, (b) causal encoding module, (c) prediction module. The arrow in the back shows the inference order of them.
• Figure 2: Overview of the dynamic causal discovery module.
• Figure 3: Sensitivity on weight of sparsity regularization, $\lambda_1$
• Figure 4: Influence of training instance amount.
• Figure 5: Sensitivity on $c$, i.e, acyclicity regularization $\mathcal{R}_{DAG}$