Causal discovery based on hierarchical reinforcement learning

Abstract

Conditional independence (CI) tests in causal discovery can determine a set of Markov equivalence classes w.r.t. the observed data by checking whether each pair of variables is d-separated under faithfulness and Markov assumptions. However, CI tests are intractable for high-dimensional conditional variables. Motivated by the advantages of reinforcement learning in exploring the solution space, firstly, we propose a causal discovery framework based on hierarchical reinforcement learning (CD-HRL). This framework trains both the discovery of the causal skeleton and the identification of direction using two interdependent high-level and low-level policies seperately. Dividing causal discovery into two distinct subtasks to high-level and low-level policies enhances exploration efficiency and minimises error accumulation. The high-level policy iteratively generates causal skeletons as subgoals for instructing the low-level policy, which then identifies causal directions of individual pairs of variables. Secondly, to avoid redundant exploration of familiar causal structures, we incorporate a memory module into the high-level agent and predefine an augmented reward that combines a causal score function and a curiosity item for exploring unknown causal structures. Lastly, experiments on both synthetic and real datasets show that the proposed approach outperforms the state-of-the-art methods under various data-generating procedures, which follow linear, nonlinear, and ordinary differential equations with additive Gaussian noise. The code for our CD-HRL method is available online in https://github.com/HITshenrj/CD-HRL.