Causal Discovery with Reinforcement Learning

Shengyu Zhu; Ignavier Ng; Zhitang Chen

Source-linked AI summary

Causal Discovery with Reinforcement Learning

Shengyu Zhu, Ignavier Ng, Zhitang Chen

arXiv:1906.04477v4Jun 11, 2019cs.LGstat.ML

TL;DR

Causal discovery from observational data is difficult because score-based DAG search is combinatorial and randomized experiments may be impractical. The paper uses reinforcement learning with an encoder-decoder graph generator and acyclicity-aware rewards, reporting improved search ability and favorable comparisons with existing methods.

Problem
Causal discovery from observational data is needed when randomized experiments are expensive or impossible, but score-based DAG search is generally NP-hard under the combinatorial acyclicity constraint.
Method
The method uses reinforcement learning as a search strategy, generating graph adjacency matrices from observed data and rewarding them with a predefined score plus two acyclicity penalties.
Results
The approach shows improved search ability and flexible score selection, with BIC outperforming GES on LiNGAM and linear-Gaussian datasets and recent gradient-based methods for nonlinear relationships.
Takeaways & Limitations
RL can search for high-scoring causal DAGs while retaining flexibility in the predefined score function under an acyclicity constraint.
Takeaways & Limitations
Graphs with more than 50 nodes remain challenging, although the method has been demonstrated on graphs with up to 30 nodes.

Abstract

from arXiv · show

Discovering causal structure among a set of variables is a fundamental problem in many empirical sciences. Traditional score-based casual discovery methods rely on various local heuristics to search for a Directed Acyclic Graph (DAG) according to a predefined score function. While these methods, e.g., greedy equivalence search, may have attractive results with infinite samples and certain model assumptions, they are usually less satisfactory in practice due to finite data and possible violation of assumptions. Motivated by recent advances in neural combinatorial optimization, we propose to use Reinforcement Learning (RL) to search for the DAG with the best scoring. Our encoder-decoder model takes observable data as input and generates graph adjacency matrices that are used to compute rewards. The reward incorporates both the predefined score function and two penalty terms for enforcing acyclicity. In contrast with typical RL applications where the goal is to learn a policy, we use RL as a search strategy and our final output would be the graph, among all graphs generated during training, that achieves the best reward. We conduct experiments on both synthetic and real datasets, and show that the proposed approach not only has an improved search ability but also allows a flexible score function under the acyclicity constraint.

1 INTRODUCTION

Causal discovery from observational data is valuable when randomized experiments are costly or impossible, but score-based DAG search is computationally hard and existing heuristics offer limited flexibility. The paper proposes reinforcement learning as a search strategy that combines a predefined score with acyclicity penalties.

Observational causal discovery is attractive because controlled randomized experiments can be expensive or impossible in fields such as social sciences and bioinformatics.
Score-based methods assign scores to directed graphs and search over all DAGs for the best-scoring structure.
The optimization is generally NP-hard because the acyclicity constraint is combinatorial and the number of DAGs grows super-exponentially with graph size.
Existing approaches use local heuristics such as GES, which enforces acyclicity one edge at a time, while score functions and search strategies lack a principled problem-specific combination.
The proposed RL method generates graphs from observed data, rewards them using a predefined score and two acyclicity penalties, and returns the best graph generated during training.
The method reportedly improves search ability, preserves score-function flexibility, and with BIC outperforms GES on LiNGAM and linear-Gaussian data.

2 RELATED WORK

Related work spans constraint-based, functional-model, neural, and reinforcement-learning approaches to causal discovery and combinatorial optimization. The paper distinguishes its actor and reward designs by generating adjacency matrices and explicitly incorporating causal scores with acyclicity constraints.

Constraint-based methods use conditional independence tests to estimate causal skeletons and orient edges up to a Markov equivalence class.
Functional causal models can distinguish DAGs within an equivalence class by imposing additional assumptions on data distributions or functional classes.
Neural causal-discovery methods include models using prior skeletons or adversarial causal generative modeling, with the latter not guaranteeing acyclicity.
Neural sequence-to-sequence models and RL have been applied to combinatorial optimization, including traveling-salesman and graph problems.
Unlike typical RL applications that learn a policy, this paper uses RL for search and selects the best generated graph.
Its actor is an encoder-decoder that generates graph adjacency matrices, while its reward combines a causal score with an acyclicity constraint.

3 MODEL DEFINITION

The model represents each variable as a node in a DAG and generates observations from parent variables through functions plus independent additive noise. Causal discovery uses sampled observational data to infer the graph, with identifiability depending on model assumptions.

Each variable x_i corresponds to a node in a d-node DAG, and its observed value is generated from its parents through a function plus additive noise n_i.
The parent set contains variables with directed edges into x_i, and the noises are jointly independent.
Causal minimality requires each function f_i to be nonconstant in every argument.
Without further assumptions on functions or noises, the model is identifiable only up to a Markov equivalence class under standard Markov and faithfulness assumptions.
The experiments use fully identifiable synthetic models so estimated graphs can be evaluated against the true DAG.
The observed dataset consists of independently sampled vectors, and the objective is to infer the causal graph that best describes the data-generating procedure.

4 NEURAL NETWORK ARCHITECTURE FOR GRAPH GENERATION

The graph generator maps observed data to a binary adjacency matrix using an attention-based encoder and an element-wise decoder. This architecture is designed to capture interactions among variables while avoiding self-loops and supporting acyclic graph search.

A feed-forward network that independently outputs d^2 scalars for an adjacency matrix failed to produce promising results, possibly because it captured insufficient variable interactions.
The method samples variable-wise data representations and treats the variables analogously to cities in a sequence-to-sequence combinatorial optimization problem.
The target is a binary adjacency matrix A ∈ {0, 1}^{d×d} whose graph is acyclic and achieves the best score.
The encoder uses self-attention in a Transformer structure to represent the variables and capture causal relations together with the DAG constraint.
The decoder builds pairwise relationships between encoder outputs and generates adjacency entries element by element.
Each adjacency entry is transformed by a sigmoid and sampled from a Bernoulli distribution, while diagonal entries are masked to prevent self-loops.
The authors report that a single-layer decoder performs best empirically, possibly because it has fewer parameters and is easier to train.

5 REINFORCEMENT LEARNING FOR SEARCH

The method uses reinforcement learning as a search strategy for score-based causal discovery, combining flexible score functions with penalties that enforce acyclicity. An encoder-decoder agent generates graphs, while adaptive penalties and policy-gradient training guide exploration toward high-reward DAGs.

The approach uses RL to search for the highest-scoring DAG while allowing flexible predefined score functions under an acyclicity constraint.
Score Function: The BIC score combines a likelihood term based on residual sum of squares with a penalty proportional to the number of graph edges.
Acyclicity: Acyclicity is enforced with a smooth matrix-exponential penalty and an indicator penalty for non-DAGs, because the smooth term alone may be small on cyclic graphs.
Reward: The reward subtracts the score and both acyclicity penalties, with larger penalty weights making high-reward graphs more likely to be acyclic.
Penalty Selection: With properly chosen penalty parameters, the penalized optimization problem is equivalent to the original optimization with a hard acyclicity constraint.
RL Training: The algorithm normalizes score ranges, updates penalty weights during training, and uses actor-critic policy gradients with entropy regularization to encourage exploration.
Implementation: Computed rewards and decomposed BIC residual terms are recorded to reduce repeated computation during RL training.

6 EXPERIMENTAL RESULTS

Experiments evaluate the approach on synthetic linear and nonlinear models, larger graphs, and a real protein-signaling dataset, using standard causal-discovery metrics and comparisons with established methods.

Evaluation setup: The evaluation compares the proposed methods with traditional and gradient-based causal-discovery algorithms using FDR, TPR, and SHD.The benchmark includes synthetic and real datasets, with SHD measuring edge additions, deletions, and reversals.
Linear models: 683,784 different graphs were generated for a 12-node linear-Gaussian problem, far fewer than the approximately 5.22 × 10^26 possible DAGs, while the pruned DAG exactly matched the true graph.The experiment used 64 samples, 20,000 iterations, and a pruning threshold of 0.3.
Linear models: RL-BIC2 recovered all true causal graphs on both LiNGAM and linear-Gaussian data, whereas RL-BIC performed worse.With the same BIC score, RL-BIC nevertheless performed much better than GES on both datasets.
Larger graphs: On 30-node LiNGAM graphs, RL-BIC2 achieved FDR 0.14 ± 0.15, TPR 0.94 ± 0.07, and SHD 19.8 ± 23.0, comparable to NOTEARS.NOTEARS obtained FDR 0.13 ± 0.09, TPR 0.94 ± 0.04, and SHD 17.2 ± 13.12.
Nonlinear models: For quadratic causal relationships, RL-BIC2 achieved the best performance among the methods evaluated on 10-node graphs.The experiment used quadratic regression, BIC scoring, and threshold-based pruning of first- and second-order coefficients.
Nonlinear models: For Gaussian-process causal relationships, RL-BIC outperformed all other methods, while both proposed methods performed reasonably well.Fixed kernel bandwidth caused overfitting and often produced a non-DAG with the highest reward; normalization and median bandwidth heuristics were therefore applied.
Real data: On the Sachs protein-signaling dataset, both RL-BIC and RL-BIC2 achieved promising results compared with other methods.The analysis used 853 observational samples, Gaussian-process regression, median bandwidth heuristics, and CAM pruning.

7 CONCLUDING REMARKS AND FUTURE WORKS

The paper uses reinforcement learning to search for optimally scored DAGs and evaluates the method on synthetic and real data. It identifies scalability and score-computation efficiency as remaining challenges.

Contributions: The proposed RL approach searches for DAGs with optimal scores using an actor-critic algorithm, encoder-decoder actor, and acyclicity-penalized rewards.The reward combines a predefined score function with two penalty terms enforcing acyclicity.
Contributions: Experiments on synthetic and real datasets report advantages over other causal discovery methods.
Limitations: The method is effective on 30-node graphs, but graphs with more than 50 nodes remain challenging.The authors note that decomposition, prior knowledge, or constraint-based methods may reduce larger search spaces.
Future work: Score computation is more time-consuming than neural-network training, motivating more efficient score functions and early stopping criteria.The authors also identify alternative RL algorithms as a possible direction.

A MORE DETAILS ABOUT DECODERS

The decoder section describes several neural architectures for generating binary adjacency matrices from encoder outputs. On 12-node linear-Gaussian graphs, the single-layer decoder performs best, possibly because it is simpler to train.

Decoder architectures: The decoders generate binary adjacency matrices from neural-network representations of variables.The described choices include single-layer, bilinear, neural tensor network, and Transformer decoders.
Decoder architectures: The single-layer decoder uses trainable parameters W1, W2, and u, with dh denoting the decoder hidden dimension.
Decoder architectures: The bilinear decoder and neural tensor network decoder model pairwise relationships between encoder outputs using trainable bilinear parameters or a tensor.The neural tensor network uses K tensor slices and additional trainable parameters.
Decoder architectures: The Transformer decoder applies multi-head attention, shared feed-forward weights, and elementwise sigmoid sampling to produce an adjacency matrix.Its outputs are treated as the rows of a d × d matrix before binary sampling.
Empirical comparison: On linear-Gaussian data with 12-node graphs, the single-layer decoder performs best, possibly because fewer parameters make training easier.The encoder is reported to provide sufficient interactions among variables.

C PENALTY WEIGHT CHOICE

The penalty-weight analysis shows that an acyclicity indicator alone may not guide reinforcement learning toward DAGs. The method therefore uses penalty design to provide graded guidance while preserving exploration.

Penalty design: Using only the acyclicity indicator can be formally equivalent to the hard constraint but may fail to guide the RL agent toward DAGs.
Penalty design: Two cyclic graphs can receive identical indicator penalties even when one is much closer to acyclicity by edge edits or smooth acyclicity measures.In a dense linear model, the denser cyclic graph may nevertheless obtain a better BIC score and higher reward.
Initialization: Initializing the agent near DAGs also performs poorly because small penalties permit cyclic high-scoring graphs, whereas large penalties restrict exploration and produce suboptimal DAGs.

D IMPLEMENTATION DETAILS

The implementation compares RL-based causal discovery with established causal discovery methods using existing implementations. It builds on a neural combinatorial optimization implementation while modifying the reward and decoder.

Compared methods: The comparison includes ICA-LiNGAM, GES, PC, CAM, NOTEARS, DAG-GNN, and GraN-DAG.
Compared methods: ICA-LiNGAM uses independent component analysis for linear non-Gaussian additive models, while CAM separates causal-order search from feature or edge selection.
Compared methods: NOTEARS, DAG-GNN, and GraN-DAG optimize weighted adjacency representations under smooth acyclicity constraints using least-squares, evidence-lower-bound, or log-likelihood objectives.
Implementation: The implementation extends a TensorFlow neural combinatorial optimizer with entropy regularization and the paper’s modified reward and decoder.The authors state that their code was made publicly available.

E.1 EXPERIMENT 1 IN SECTION 6.1

The experiments compare causal-discovery results across LiNGAM and linear-Gaussian models while describing implementation choices for graph generation and optimization. They also examine sparse-graph initialization and preprocessing intended to support causal discovery.

0.05 ± 0.04, 0.93 ± 0.06, and 3.2 ± 2.93 are GraN-DAG's FDR, TPR, and SHD on LiNGAM data, respectively.
0.05 ± 0.04, 0.95 ± 0.03, and 2.40 ± 1.85 are GraN-DAG's FDR, TPR, and SHD on linear-Gaussian data, respectively.
Additional linear-model experiments sample noise variances uniformly from Unif([0.5, 2]).The results are reported in Table 6.2.
A trainable decoder bias is added when sparse true causal graphs are known a priori for 30-node experiments.The bias is initialized to −10 to encourage a favorable starting point for generating adjacency matrices.
For NOTEARS with quadratic features, an equivalent weighted adjacency matrix aggregates coefficient magnitudes so zero entries indicate no variable effect.The resulting optimization uses an acyclicity constraint and augmented Lagrangian optimization with Adam.