Reinforcement Learning of Causal Variables Using Mediation Analysis


  • Tue Herlau Technical University of Denmark
  • Rasmus Larsen Alexandra Institute



Machine Learning (ML), Reasoning Under Uncertainty (RU), Knowledge Representation And Reasoning (KRR)


We consider the problem of acquiring causal representations and concepts in a reinforcement learning setting. Our approach defines a causal variable as being both manipulable by a policy, and able to predict the outcome. We thereby obtain a parsimonious causal graph in which interventions occur at the level of policies. The approach avoids defining a generative model of the data, prior pre-processing, or learning the transition kernel of the Markov decision process. Instead, causal variables and policies are determined by maximizing a new optimization target inspired by mediation analysis, which differs from the expected return. The maximization is accomplished using a generalization of Bellman's equation which is shown to converge, and the method finds meaningful causal representations in a simulated environment.




How to Cite

Herlau, T., & Larsen, R. (2022). Reinforcement Learning of Causal Variables Using Mediation Analysis. Proceedings of the AAAI Conference on Artificial Intelligence, 36(6), 6910-6917.



AAAI Technical Track on Machine Learning I