Ordering-Based Causal Discovery with Reinforcement Learning
Authors: Xiaoqiang Wang
Yali Du
Shengyu Zhu
Liangjun Ke
Zhitang Chen
Jianye HAO
Jun Wang
Venue:  (2021)
Keywords: Causal Discovery
Reinforcement Learning
Ordering Search

Abstract: It is a long-standing question to discover causal relations among a set of variables in many empirical sciences. Recently, Reinforcement Learning (RL) has achieved promising results in causal discovery. However, searching the space of directed graphs directly and enforcing acyclicity by implicit penalties tend to be inefficient and restrict the method to the small problems. In this work, we alternatively consider searching an ordering by RL from the variable ordering space that is much smaller than that of directed graphs, which also helps avoid dealing with acyclicity. Specifically, we formulate the ordering search problem as a Markov decision process, and then use different reward designs to optimize the ordering generating model. A generated ordering is then processed using variable selection methods to obtain the final directed acyclic graph. In contrast to other causal discovery methods, our method can also utilize a pretrained model to accelerate training. We conduct experiments on both synthetic and real-world datasets, and show that the proposed method outperforms other baselines on important metrics even on large graph tasks.

Under review as a conference paper at ICLR 2021
ORDERING-BASED CAUSAL DISCOVERY WITH REINFORCEMENT LEARNING
Anonymous authors Paper under double-blind review
ABSTRACT
It is long-standing question to discover causal relations among a set of variables in many empirical sciences. Recently, Reinforcement Learning (RL) has achieved promising results in causal discovery. However, searching the space of directed graphs directly and enforcing acyclicity by implicit penalties tend to inefficient and restrict the method to small problems. In this work, we alternatively consider searching an ordering by RL from the variable ordering space that is much smaller than that of directed graphs, which also helps avoid dealing with acyclicity. Specif- ically, we formulate the ordering search problem as a Markov decision process, and then use different reward designs to optimize the ordering generating model. A generated ordering is then processed using variable selection methods to obtain the final directed acyclic graph. In contrast to other causal discovery methods, our method can also utilize pretrained model to accelerate training. We conduct experiments on both synthetic and real-world datasets, and show that the proposed method outperforms other basclines on important metries even on large graph tasks.
INTRODUCTION
Identifying causal structure from observational data is an important but also challenging task in many applications. This problem can be formulated as that of finding Directed Acyclic Graph (DAG) that minimizes some score function defined w.r.t. observed data. Though there exist well-studied score functions like Bayesian Information Criterion (BIC) or Minimum Description Length (MDL) (Schwarz et al. 1978 Chickering 2002), searching over the space of DAGS is known to be NP-hard, even each node has at most two parents (Chickering 1996). Consequently, traditional methods mostly rely on local heuristics to perform the search, including greedy hill climbing and Greedy Equivalence Search (GES) that explore the Markov equivalence classes (Chickering 1996).
Along with various search strategies, existing methods have also considered to reduce the search space while meeting the DAG constraint. A useful practice is to cast the causal structure learning problem as that of learning an optimal ordering of variables (Koller & Friedman, 2009). Because the ordering space is significantly smaller than the space of directed graphs and searching over ordering space can avoid the problem of dealing with acyclic constraints (Teyssier & Koller 2005). Many
algorithms are used to search for ordering such as genetic algorithm (Larranaga et al. 1996), Markoy chain Monte Carlo (Friedman & Koller 2003) and greedy local hill-climbing (Teyssier Koller 2005). However, these algorithms often cannot find the best ordering effectively.
Recently, with differentiable score functions, several gradient-based methods have been proposed based on a smooth characterization of acyclicity, including NOTEARS (Zheng et al. 2018) for linear Ju causal models and several subsequent works, e.g., Yu et al. (2019) Lachapelle et al. (2020) et al.)
(2019b)ar a). which use neural networks to model nonlinear causal relationships. As another attempt, Zhu et al. (2020) utilize Reinforcement Learning (RL) as search strategy to find the best DAG from the graph space and it can be incorporated with wide range of score functions. Unfortunately. its good performance is achieved only with around 30 variables, for at least two reasons: 1) the action space, consisting of directed graphs, is tremendous for large scale problems and hard to be explored efficiently; and 2) it has to compute scores for many non- non-DAGs generated during training but computing scores w.r.t. data is generally be time- conssunnninng It appears that the RL-based approach may not be able to achieve a close performance to other gradient- baade methods that directly optimize the same (differentiable) score function for large causal discovery problems, due its search nature.

---

Under review as a conference paper at ICLR 2021
In this work, we propose RL-based approach for causal discovery, named Causal discovery with Ordering-basse based Reinforcement Learning (CORL), which combines RL with the ordering based paradigm that we can exploit the powerful search ability of RL to search the best ordering efficiently. To achieve this, we formulate the ordering search problem as a Markov Decision Process (MIP), and then use different reward designs for RL to optimize the ordering generating model. In addition, we notice that pretrained model can be incorporated into our model to accelcrate training. For generated ordering, we prune it to the final DAG by variable selection. The proposed approach is evaluated on both synthetic and real datasets to validate its effectiveness. In particular, the proposed method can achieve a much improved performance than the previous RL-based method on both linear and non-linear data models, even outperforms NOTEARS, a gradient-based method, on 150-node linear data models, and competitive with Causal Additive Model (CAM) on non-linear data models.
RELATED WORKS
Existing causal discovery methods roughly fall into three classes. The first class, as described in the introduction, are the so-called score-based methods. Besides the mentioned BIC/MDL scores, other score functions include the Bayesian Gaussian equivalent score (Geiger & Heckerman, 1994).
the generalized score based on (conditional) independence relationship (Huang et al. 2018). and recently proposed meta-transfer score (Bengio et al. 2020). Another class of methods such as fast causal inference and PC (Spirtes et al. 2000 Zhang, 2008) which first find causal skeleton and then decide the orientations of the edges up to the Markov equivalence class, are viewed as constraint- Such methods usually involve multiple independent testing problems; the testing results may have conflicts and handling them is not easy. The last class of methods relies on properly defined functional causal models, includeing Linear Non-Gaussian Acyelic Model (LiNGAM), nonlinear Additive Noise Model (ANM) (Hoyer et al., 2009] Peters et al., 2014), and the post post-nonlinear causal model (Zhang & Hyvärinen 2009). By placing certain assumptions on the class of causal functions and/or noise distributions, these methods can distinguish different DAGS in the same Markov equivalence class.
Since particularly consider ordering- based approaches, here we present more detailed review of such methods. Most of the ordering-base methods belong to the class of score-based methods. Besides the mentioned heuristic search algorithms, Schmiitt al. (2007) proposed LIOBS conduct variable selection using & f1-regularization paths based on Teyssier & Koller (2005). Scanagatta et al.) (2015) further proposed an ordering exploration method on the basis of an approximated score function so as to scale to thousands of variables. The Causal Additive Model (CAM) was developed by BBühlmann et al. 4014) for nonlinear data models. Some recent ordering-basee -baed methods such as sparsest permutation (Raskutti& & Uhler 2018) and greedy sparest permutation (Solus et al. (017) can guarantee consistency of Markov equivalence class, based on some conditional independence relations and certain assumptions like faithfulness. A variant of greedy sparest permutation was also proposed in Bernstein et al. (2020) for the setting with latent variables. In the present work, we mainly work on identifiable cases which may have different assumptions from theirs.
In addition, several exact algorithms such as dynamic programming (Silander & Myllymäki 2006 Perrier et al. 2008) and integer or linear programming (Jakkkola et al. |010| Cussens] 2011] Bartlet & Cussens 2017) are used to solve causal discovery problem. However, these exact algorithms usually work on small graphs efficiently (De Campos & Ji. [001]. and in order to handle larger problems with hundreds of variables, they need to incorporate heuristics search (Xiang & Kim 2013) or limit the maximum number of parents.
Recently, RL has been used to tackle several combinatorial problems such as maximum cut and the traveling salesman problem, due to their relatively simple reward mechanisms (Khalil et al. (017)). In combination with the encoder-decoder based pointer networks (Vinyals et al. 2015). Bello et al. (2016) showed that RL can have better generalization even when the optimal solutions are used as can
labeled data in supervised way. Kool et al. (219) further used an attention based encoder-decoder model for an improved performance. These works aim to learn a policy as a combinatorial solver based on the same structure of a particular type of combinatorial problems. However, various causal discovery tasks generally have different causal relationships, data types, graph structures, etc, and moreover, are typically off-line with focus on a causal graph. As such, we use RL as a search strategy, similar to Zhu et al.] (2020); Zoph & Le 2017); nevertheless, a pretrained model or a policy can offer good starting point to speed up training, as shown in our evaluation results (Figure 3).
2

---

Under review as a conference paper at ICLR 2021
3 BACKGROUND
3.1 CAUSAL STRUCTURE LEARNING
Let 9 V. denotes a DAG, with the number of nodes, V {oi..... va) the set of nodes.
and E {(v,, "i ji.j 1, ,...., df the set of directed edges from o2 to 1z. Each node vj is associated with random variable x.. The probability model associated with 9 factorizes as p( p(X1:.... Xa) I I p(X, Pa(X))), where ((x,(P((X))) is the conditional probability distribution for X, given its parents Pa(X;) == (uK, my) E E3. We assume that the observed data xj is obtained by the Structural Equation Model (SEM) with addtive noise: Xj ;= fe, (Pa(X;)) + Ej.J = 1,..., d, where fo, parameterized by 8, is used to represent the functional relationship between X, and its parents, and Cj s denote jointly independent additive noise variables. We assume causal minimality, which is equivalent to that each fy is not constant for any X Pa(Xx) in this SEM (Peters et al. 2014).
Given sample X [ Xa]. where X Rdxd and x, is vector of observations for random variables Xy. Our goal in this paper is to find the DAG g that maximizes the BIC score (or other well studied scores), defined
10 2
Scoresic(9)
Σ
log (X) 0,)
log
where 10 is the number of free parameters in p( x Pa (Xx)) (x)))))) 0,)) For linear causal relationships,
10,1 | the number of parents, up to some constant factor.
The problem of finding a directed graph that satisfies the ayclicity constraint can be as that of finding variable ordering (Teyssier & Koller 2005 Schmidt et al. (207)). The score of an ordering is usually defined as the score of the best DAG that is consistent with the given ordering. Specifically, let II denote an ordering of the nodes in V, where the length of the ordering [TT| = [V] and II is indexed from 1. If node By E V lies in the p-th position, then 11(p) = V/, Notation II,,, denotes the set of nodes in V that precede One can establish a canonical correspondence between an ordering IT and a fully-connected DAG G'l; an example with four nodes presented in Figure 1. For a given DAG G, it can be consistent with more than one orderings and the set of these orderings is given by
1/2
14
V1
V3
node Vy in
I1 = (V2, V4, D1 V33
Figure 1: An example of the correspondence between an order- ing (down) and fully-connected DAG (top).
(II) = (II: the fully-comnected DAG gll a super- of g),
(2)
where a super-DAG of & is DAG whose edge set superset of that of G. The score of an ordering is usually defined as the score of the best DAG that consistent with the given ordering (Teyssier Koller 20055 Peters et al. 2014). We provide formal description in Proposition to show that it is possible to find the correct ordering with high probability limit. Therefore, the search for the true DAG G" can be decomposed to two phases: finding the correct ordering and performing variable selection (feature selection); the latter is to find the optimal DAG that is consistent with the ordering found in the first step.
the large sample
Proposition 1. Suppose that an identifiable SEM with true causal DAG G* on X (X,1ş= induces distribution P(X). Let glll be the fully-connected DAG that corresponds to an ordering IT. If theree is an SEM with Gn inducing the same distribution P(X), then gil must be a super-graph of g", i.e. every edge in G is covered in g
Proof. The SEM with gd may not be causally minimal but can be reduced to an SEM satisfying the causal minimality condition (Peters et al. 2014). Let ğIIl denotes the causal graph in the reduced SEM with the same distribution P(X). Since we have assumed that original SEM is identifiable. i... the distribution P(X) corresponds to unique true graph, GIIl is then identical to g". The proof is complete by noticing that glll is super-graph of ş"..
n
3

---

Under review as a conference paper at ICLR 2021
3.2 REINFORCEMENT LEARNING
Standard RL is usually formulated as an MDP over the environment state s E S and agent action E A, under an (unknown) environmental dynamics defined by transition probability T(s'|s, a). We use TH (a)s) to denote the policy, parameterized by Φ. which outputs a discrete (or continuous) distribution used to select an action from action space A based on state 8. For episodic tasks, a trajectory {st, at H where T is the finite time horizon, can be collected by executing the policy repeatedly. In many cases, an immediate reward r(s, a) can be received when agent exccutes an action. The objective of RL is to learn policy which can maximize the expected cumulative
reward along a trajectory, i.e., (Φ) Exg [Ro] with Ry Σ ( 44) and (0, 1] being a discount factor. For some scenarios, the reward is only carned at the terminal time (also called
episodic reward), and (ő) E.,, [[(]] with R(T) = (r
4
METHOD
In this section, we first introduce how to model the ordering search problem as an MDP, then we show how to use RL to find the optimal ordering. and we introduce how to process the searched ordering to obtain the final DAG, finally we provide a discussion on computational complexity.
4.1 ORDERING SEARCH AS A MARKOV DECISION PROCESS
We regard the variable ordering search problem as muti-stage decision problem with variable selected at each decision step. We sort the selected variables according to decision steps to obtain variable ordering, which is defined as the searched ordering. The decision-making process is Markovian, and its elements in the problem can be defined as follows:
State One can directly take the sample data Xj of each variable Xj as state 8j. However, preliminary experiments show that it is difficult for feed-forward neural network models to capture the underlying causal relationships directly using observed data as states, and the data pre- pre-processed by a module conventionally named encoder is helpful to finding the better ordering, see Appendix A. The encoder module embeds each X, to 8, and all the embedded states constitute the space
{81,..., sa). Adding initial state 80 to the space constitutes the complete state space S = S U so- We 8, to denote the state encountered at the 1-th decision step.
Action We select an action (variable) from the action space constituted by all the variables at each decision step, and the space size is equal to the number of variables, |A4] d. Compared to the previous RL-based method searching from the graph space with size ((2xd) (Zhu et al. 2020). the search space is smaller. Note that according to the definition of ordering, the first selected variable is the source node and the last selected node is the sink node.
State transition The specified state transition is related to the action selected at the current decision step. Specifically, if the selected variable is 2y at the f-th decision step, then the state is transferred to the state 8, E S which the j-th output from Transformer encoder, i.e., 51+1 = 8j-
Reward As we described in Section 3.1, only the variables that have been selected in previous decision steps can be the potential parents of the currently selected variable. Hence, we can design the rewards in the following cases; dense reward and episodic reward. For dense reward case, we can exploit the decomposability of the score function (BIC score) to calculate an immediate reward for the current decision step based on the potential parent variables (that have been selected), i.e., if Vy is selected at time step t, the immediate reward calculated by
log pl X( - 10,1 2
log m,
where U(Xy) denotes the potential parent variable set of X, and the set consists of the variables associated with the nodes in .L For episodic reward case, we directly calculate a score as an episodic reward for complete variable ordering regardless of whether the scoring function is decomposable or not, i.c.,
R(T) rr(6, a) Scorenic
4

---

Under review as a conference paper at ICLR 2021
where Scoregic has been defined in Equation (1), with the set Pa(Xy) of each variable X, replaced by the potential parent variable set U (X2)) here.
4.2 IMPLEMENTATION AND OPTIMIZATION WITH REINFORCEMENT LEARNING
We describe the neural network architectures implemented in our method, which consists of an encoder and a decoder as shown in Figure 2. Here we briefly describe the model architectures and leave details regarding model parameters Appendix A.
Encoder fe X + S is used to map the observed data to the embedding space S {81.....Sd Sd). . For sample efficiency, we follow Zhu jet al. (2020) to randomly draw samples from the dataset X to construct X Raxd each episode and use X instead of X. We also set the embedding 8y to be in the same dimension, i.e., 8, E IR". For conduct empirical compari- son among several possible structures such as a self- attention based encoder (Vaswani et al. 2017) and an LSTM structure, Empirically, we confirm that the self-attention based encoder in the Transformer struc- ture performs the best, which is also used in Zhu et al. (2020). Please find more details in Appendix|A..
S)
4g
Encoder
Decoder
encoder choice,
xj
80
&
ST
Figure 2: Illustration of the policy model. The encoder embeds the observed data x, into the state sj. An action can be selected by the decoder according to the given state š, at each time step I.
Decoder pdce S A maps the state space S to te the action space A. Among several decoder choices (please see also Appendix I for empirical comparison), we specifically pick an LSTM based structure that proves effective in our experiments. Although the initial state is usually generated randomly, we pick as S. = 1/2 Σi= 5; considering that the source node is fixed in correct ordering. We then restrict each node only be selected once by masking the selected nodes so to generate valid variable ordering (Vinyals et al.) 2015).
Optimization The optimization objective is to learn policy which maximizes (Φ) E, [R] where Te denotes the policy model parameterized by the paprameters {Φe, Φu) of encoder frne and decoder fdve, Based the above definition, policy gradient (Sutton & Barto 2018) is used to optimize the ordering generation model parameters. For the dense reward case, policy gradient can T be written as (Ø) Er,, Σ R. log T& (ael6e) where Ri - 7 r'ret1 denotes the return at time step t. We denote the algorithm using this reward design as CORL- For the episodic reward case, we have the following policy gradient ( = Ex, R(T) E.o lo Ts (aol&c) the algorithm using this reward design is denoted as CORL-2. Using a parametric baseline to estimate the expected score typically improves learning (Konda & Tsitsiklis 2000). Therefore, we introduce a critic network Vo_ ($:) parameterized by du., which learns the expected return given state 8, It is trained with stochastic gradient descent using Adam optimizer on mean squared eπтоr objective between its predicted value and the actual return. The details about the parameters of critic network are described in Appendix A.2
Inspired by the benefits of pretrained models in other tasks (Hinton & Salakhutdinoy 2012), we also consider to incorporate it to our method to accelerate training. Usually, one can obtain some observed data with known causal structure or correct ordering, e.g., by simulation or real data with labeled graph. Hence, we pretrain policy model with such data in supervised way.
So far we presented CORL in a general manner without specifying explicitly which distribution family is during the evaluation of rewards. In principle, any distribution family could be employed as long as its log-likelihood can be computed and no differentiability is required. However, is not always clear whether the maximization of the accumulated reward recovers the correct ordering. It will depend on both the modelling choice of reward and the underlying SEM; in fact, the causal relationships fall into the chosen model functions and right distribution family is assumed, then given infinite samples the optimal accumulated reward, corresponding to the negative log- likelihoood must be achieved by super-DAG of the underlying graph according to Proposition 1. In practice,
5

---

Under review as a conference paper at ICLR 2021
Algorithm 1 CORL.
Require: observed data X. initial parameters te. dJ and 8e, two empty buffers 'D and D scooee initial value (negative infinite) BestScore and a random ordering BestOrdeing.
1: while not terminated do
2: draw a batch of samples from X, encode them to S and calculate the initial state &o
3: for 0, 0,.,..... do
4: collect batch of data ($,, a,, Tt) with Tg: = U {(3,, ae m())
5: if (ue, II -T-,, r is not in Dsoore then
6: store (vr, II <o, re) in Dscore to avoid repeated computations
7: end if
8:
end for
9:
update l..0d. 0g according to Section 4.2
10: if -i > BestScore then
update the BestScore and BestOrdering 12: end if
11:
13:
end while
14: get the final DAG by pruning the BestOrdering
can only apply approximate model functions and also need to assume certain distribution family for caculating the reward.
Our method is summarized in Algorithm 1. In addition, record the decomposed scores for each variable U with different parental sets Ik to avoid repeated computations which are generally time-consuming (see Section 4.4). Although we cannot guarantee find the optimal ordering because policy gradient can at most guarantee local convergence (Sutton et al. (000) and also we only have access the empirical log- likelihoodd we remark that the ordering obtained from CORL still enjoy good performance in the experiments, compared with consistent methods like GES and PC.
4.3 VARIABLE SELECTION
an estimated ordering Íî consistent, then we obtain a cully-connected DAG (super-DAG) gr of the underlying DAG G. One can then pursue consistent estimation of intervention distributions based on πî without any additional need to find the true underlying DAG g* (Bühlmann et al. 2014). For other purposes, however, we need to recover the true graph from the fully fully-connected DAG. There exist several efficient methods such as sparse candidate (Teyssier & Koller (2005), significance testing of covariates (Bühlmann et al. 2014), the group Lasso (Ravikumar et al. 2009), or its improved
version with a sparsity smottheess penalty proposed in |Meier et al. (2009).
For linear data models, we apply linear regression to the obtained fully fully-comnected DAG, followed by thresholding to prune edges with small weights, similarly used by Zheng al. (018); Yu et al. (2019) : Zhu et al. (2020). For the nonlinear model, we follow the pruning process used by Bühlmann al. (211)) 4): Lachapelle
al. (2020). Specifically, for each variable Xj. one can fit generalized additive model against the current parents of X, and then apply significance testing of covariates, declaring significance the reported p-values are lower or equal to 0.001.
4.4 COMPUTATIONAL COMPLEXITY
To learn an ordering, CORL relies on the proper training of RL model. Policy gradient and stochastic gradient are adopted to train the actor and critic respectively, which are the standard choice in RL (Konda & T lis 2000). Similar to RL-BIC2 (Zhu et al., (020), CORL requires the evaluation of the rewards at each episode with O(dm² + d3) computational cost if linear functions are adopted to model the causal relations, but does not need to compute the matrix exponential term with (dd²) cost due to the use of ordering search. In addition, CORL formulates causal discovery as a multi-stage decision process and we observe that CORL performs fewer episodes than RL-BIC2 before the episode reward converges (see Appendix C). We suspect that due significant reduction in the size of action space, the model complexity of the RL policy is reduced, thus leading to higher sample efficiency. The evaluation of Transformer encoder and L.TM decoder in CORL take ((n22) and

---

Under review as a conference paper at ICLR 2021
Odd")), respectively. However, we find that computing rewards is more dominating in the total runing time (c.g., around 95% and 87% for 30- and 100-node linear data models, respectively). Speeding up the calculation of rewards would be helpful in extend our approach to larger problem, which is left as a future work.
In contrast with typical RL applications, treat RL here search strategy for causal discovery, aiming to find an ordering that achieves the best score and then applying a variable selection method to remove redundant edges. Nevertheless, for the pretraining part with the goal of a good initialization, we may want sufficient generalization ability and hence consider diverse datasets with different number of nodes, noise types, causal relationships, etc.
EXPERIMENTS
In this section, we evaluate our methods against number of methods on synthetic datasets with linear and non-linear causal relationships and also a real dataset. Specifically, these baselines include ICA-LiNGAM (Shimizu al. 2006), three ordering-based approaches LIOBS (Schmidt et al. 20077. CAM (Bühlmann et al. (2014) and A Lasso (Xiang & Kim 2013), some recent gradient- based approaches NOTEARS (Zhegg al. 2018), DAG-GNN (Yu et al. 2019) and GraN-DAG (Lachapelle et al. 2020), and the RL-baased approach RL-BIC2 (Zhu et al. 2020). For all the compared algorithms,
we use their original implementations (see Appendix B. for details) and pick the recommended
hyper-parameters unless otherwise stated.
Different types of data are generated in synthetic datasets which vary along five dimensions: level of edge sparsity, graph type, number of nodes, causal functions and sample size. Two types of graph sampling schemes: Erdös-Rényi (ER) and Scale-free (SF) are considered in our experiments. We denote d-node ER or SF graphs with average hid edges as ER/ or SFh. Two common metrics are considered: True Positive Rate (TPR) and Structural Hamming Distance (SHD). The former indicates the probability of correctly finding the positive edges among the discoveries (Jain et al., 2017). Hence, it can be used to measure the quality of an ordering, the higher the better. The latter counts the total number of missing, falsely detected or reversed edges, the smaller the better.
5.1 LINEAR DATA MODELS WITH GAUSSIAN AND NON-GAUSSIAN NOISE
We evaluate the proposed methods on Linear Gaussian (LG) with equal variance Gaussian noise and LiNGAM data models, and the true DAGs both cases are known be identifiable (Peters) & BBühlmann 2014 Shimizu et al. 2006). We set h E {2. 5 and d E (30, 50, 100} to generate observed data following the procedure done in /Zheng et al. (2018) (see Appendix B B.22for for details). For variable selection, set the thresholding as 0.3 and apply it to the estimated coefficients.
Table 1: Empirical results for ER and SF graphs of 30 nodes with LG data.
RANDOM
NOTEARS
LIOBS
Lasso
CORL-1
CORL-2
B82
1P%
(0.04)
(0.03)
0.91 (0,05)
0.4
(0.86)
0.8 (0.04)
0,99 (0.02)
0.99 (0.01)
SHD
(36.7
(9.4
(12.4)
17.8 (22.5)
(23.8)
(14.3)
(7.4)
4.4 (3.5)
ER5
(9,0)
091
(0.04)
(0.05)
(0,83)
SHD
210.2
(7.3)
68.0
45.6 (13.3)
98.6
(21.5)
37.4 (16.9)
(74.5
SF2
1PR
(0.02)
(0.02)
(0.09)
(0.02)
(0.94)
(0.02)
(0.01)
SHD)
(12.5)
6.1 (2.3)
36.8 (35.1)
3.2 (1.7)
49.7 (28.1)
27.3 (18.4)
0.0 (0.0)
0.0 (0.0)
SPS
TPR
0.44 (0.03)
0.94 (0.03)
0.89 (0.09)
0.96 (0.03)
0.79 (0.04)
0.58 (0.07)
1,00 (0.00)
1.00 (0.00)
SHD
(10.6)
23.3 (6.9)
47.8 (35.2)
(5.2)
(25.7)
40.5 (19.8)
0.0 (0.I)
0.0 (0.0)
Tables T present results only for 30- and 100-node LG data models since the conclusions do not
change with graphs of 50 nodes (see Appendix for 50-node graphs). We report the performance of the popular ICA-LiNGAM, GraN-DAG and CAM in Appendix since they are almost never on par with the best methods presented in this section. CORL-1 and CORL-2 achieve consistently good results on the LiNGAM datasets which are reported in Appendix E due to the space limit.
We now examine Tables & (the values in parentheses represent the standard deviation across datasets per task). We can see that, across all settings, CORL-1 and CORL-2 are the best performing methods, both in terms of TPR and SHD, while NOTEARS and DAG-GNN are not too far behind.
7

---

Under review as a conference paper at ICLR 2021
Table Empirical results for ER and SF graphs of 100 nodes with LG data.
RANDOM
NOTEAAR
DAG-ENN
RL-BIC2
L108S
AT Lasse
CORL-1
CORL-2
ER2
TPR
(..55
093 (0.02)
0.93 (8,03)
0.02 (0.01)
0.54 (0.02)
0.86 (0.04)
0.98 (0,02)
0.98 (0.01)
SHD
(17.4)
72.6
(19.))
270.8 (13.5
481.2 (49.9)
128.5 (38.4)
24.8 (10.1)
18.6 (5.7)
ERS
TPR
(0.04)
0.91 (0.01)
0.86 (0.16)
0.08 (0.03)
0.53 (0.02)
0.82 (0.05)
0.93 (0,02)
0.94 (0.03)
SHD
(35.7)
170.3 (34.2)
236.4 (36.5)
(46.2
547.9 (63.4)
244.0
175.3 (18.9)
164.8 (17.3)
SP2
TPR
0.48 (0.03)
0.98 (0.01)
0.89 (0.14)
0.04 (0.02)
0.57 (0.01)
0.92 (0.03)
1.00 (0.00)
1,00 (1.00)
SHD
503.4 (25.8)
(1.3)
156.8 (21.2)
(17.2
377.3 (53.4)
54,0 (22.3)
0.9 (0.0)
0.0 (0,00
SP5
1PK
0.57 (0.04)
0.9 (0,01)
0.87 (0.15)
0.05 (0.03)
0.55 (0.04)
0289 (0.03)
0.97 (0.02)
0.98 (0.01)
SHD
(19,4
90.2 (34.5)
165.2
405.2 (77.4)
503.7 (56.4)
(36.4)
(5.2)
(0.5
As we discussed previously, RL-BIC2 only achieves satisfactory results on graphs with 30 nodes. The TPR of LIOBS is lower than that of A Lasso, which indicates that LIOBS using greedy hill-climbing with tabu lists may not find a good ordering. Note that SHD of L1OBS and A Lasso reported here are the results of applying our pruning method. We observe that SHD of them is greatly improved by pruning. For example, specifically. SHD of LIOBS decreases from 171.6 (29.5), 588.0 (66.2) and 1964.5 (136.6) to 85.2 (23.8), 215.4 (26.3) and 481.2 (49.9) in LG ER2 with 30, and 100 nodes respectively, while TPR almost did not degraded. It shows the effectivvnness of pruning by model fitting with a thresholding.
We further evaluate our method 150-node LG ER2 data models. On the datasets, CORL-1 has TPR and SHD being 0.95 (0.01) and 63.7 (9.1), while CORL-2 has 0.97 (0.01) and 38.3 (14.3), respectively. CORL-2 outperforms NOTEARS with 0.94 (0.02) and 50.8 (21.8).
Although both CORL-1 and CORL-2 achieve the de- sired performance, one notice that CORL-2 achieves slightly better SHD than CORL-1. This a little different from the usual understanding that RL is usually easier to learn from dense rewards than from episodic reward case. We take the 100-node LG ER2 data models an example to show the training reward curves of the C and CORL-2 in Figure 3 One can notice that CORL-2 converges faster to a better result than CORL-1, which corresponds to the fact that CORL-2 achieves the slightly better TPR and SHD than CORL-1. We conjecture that this is because it is difficult for the critic to learn to predict the score accurately for each state; in the episodic reward case, however, it only required to learn to predict the score accurately for the initial state.
Edisode roward
CORL-1
D0RL-2
C preerrn
200
400
600
50c
1020
Episade
Figure 3: Learning curves of CORL-1, CORL-2 and CORL-2- on 100- node linear Gaussian dataset.
Pretaining We show the training reward curve of CORI-2-pretrain in Figure 3 which is CORL-22 based on pretrained model trained in the su- pervised way. The testing data task is unseen in the datasets for pretraining which contain 30-node LiNGAM ER2, 50-node LiNGAM ER2, 30-node LiNGAM SF2 and 30-node GP ERI. Obviously compared to that of CORL-2 using random initialization (CORL-2), the use of pretrained model can accelerate the model learning process. Consistent conclusion can be drawn from the experiment of CORL-1, see Appendix G In addition, we consider using the policy model learned only on 30-node LiNGAM data model as the pretrained model on 100-node LG task. We observe that the performance is similar to that of using the pretrained model obtained in the supervised way, we do not report the results repeatedly here.
5.2 NON-LINEAR MODEL WITH GAUSSIAN PROCESS
In this experiment, we consider to use Gaussian Process (GP) to model causal relationships which each causal relation f; is a function sampled from GP with radial basis function kernel of bandwidth one and normal Gaussian noises, which is known to be identifiable according to Peters et al. (2014). We set h 1 and 4 to get ER1 and ER4 graphs, respectively, and generate data by exploiting them (see Appendix B.2 details). Our method evaluated on different sample numbers, the results when m 500 are reported here (see Appendix F for additional results). Note that due to the efficiency of
8

---

Under review as a conference paper at ICLR 2021
(P10TR) GP10TR4
0.3
GPISEAT CPISEM
0.7
0.0
F 0.5
0.4
033
40
0.2
20
CORL- 2o
Cast
orak DAG ON
coff: CoRt C.
orate 080
CORL CoRtr cM
wrkk 0AG 6 ac
coft CAM
onsti 0&G
oth
(a) TPR on GP10
(b) SHD on GP10
(c) TPR on GP30
(d) SHDD on GP30
Figure 4: The empirical results GP data models with 10 and 30 nodes.
the reward calculation, we only experimented on nonlinear data models of up to 30 nodes scale. The pruning method for variable selection used here is the CAM pruning from Bühlmann et al. 2014).
The results on 10-node and 30-node datasets with ERI and ER4 graphs are shown Figure 3. Here only consider some baselines that competitive in the nonlinear data models, where CAM is very strong ordering based baseline. Although GraN-DAG achieves better results than DAG-GNN, they worse than CAM overall. We believe this is because 500 samples are so small that GraN-DAG and DAG-GNN have not learned the good model. RL-BIC2 performs well on 10-node datasets, but achieves poor results on 30-node datascts, probably due to its lack of scalability. CAM, CORL-I and CORL-2 have good results, with CORL-2 achieving the best results on 10-node graphs and slightly worse than CAM on 30-node graphs. All of these methods have better results on ERI than on ER4, especially on 30-node graphs.
5.3 REAL DATA
The Sachs dataset (Sachs et al., 2005). with 11-node and 17-edge true graph, widely used for research graphical models. The expression levels of protein and phospholipid in the dataset be used to discover the implicit protein signal network. The observational dataset has m = 853 samples and is used discover the causal structure. GP used to model the causal relationship in our method. In this experiment, CORL-1, CORL-2 and RL-BIC2 achieve the best SHD 11. CAM, GraN-DAG, and ICA-LiNGAM achieve the SHD 12, 13 and 14, respectively. Particularly, DAG-GNN and NOTEARS result in SHD 16 and 19, respectively, whereas an empty graph has an SHD 17.
6 CONCLUSION
In this work, we have proposed a RL-bbased approach for causal discovery named CORL. It searches the space of variable orderings, instead of the space of directed graphs. We formulate ordering search as an MDP and have further proposed CORL-I and CORL-2 for training the ordering generating model. For a generated ordering, it can be pruned by variable selection to the final DAG. The empirical results on the synthetic and the real datasets show that our approach is promising.
REFERENCES
Mark Bartlett and James Cussens. Integer linear programming for the bayesian network structure learning problem. Artificial Intelligence, 244:258 2211, 2017.
Irwan Bello, Hieu Pham, Quoc Le, Mohammad Norouzi, and Samy Bengio. Neural combinatorial optimization with reinforcement learning. arXiv preprint 11Xiv:1611.09940 2016.
Yoshua Bengio, Tristan Deleu, Nasim Rahaman, Rosemary Ke, Sébastien Lachapelle, Olexa Bilaniuk, Anirudh Goyal, and Christopher Pal. A meta-transfer objective for learning to disentangle causal mechanisms. In International Conference on Learning Representations (ICLR). 2020.
Daniel Bernstein, Basil Saeed, Chandler Squires, and Caroline Uhler. Ordering learning in the presence latent variables. International Conference on Artificial Intelligence and Statistics (AISTATS), pp. 4098 4108- PMLR, 2020.
Brdering-based based causal structure

---

Under review as a conference paper at ICLR 2021
Peter Bühlmann, Jonas Peters, Jan Ernest, al. CAM: Causal additive models, high-dimensional order search and penalized regression. The Annals of Statistics, 42(6))2526 2556, 2014.
David Maxwell Chickering. Learning Bayesian networks is NP-complete. Learning from Data: Artificial Intelligence and Statistics V. Springer, 1996.
David Maxwell Chickering. Optimal structure identification with greedy search. Journal of Machine Learning Research, 30755)500-555 2002.
James Cussens. Bayesian network learning with cutting planes. In Proceedings of the Twenty-Seventh Conferencee on Uncertainty Artificial Intelligence (UAI), pp. 153-160, 2011.
Cassio De Campos and Qiang Ji. Efficient structure learning of bayesian networks using constraints.
The Journal of Machine Learning Research, 12:663-689, 2011.
Nir Friedman and Daphne Koller. Being Bayesian about network structure. A Bayesian approach to structure discovery Bayesian networks. Machine learning, 50(1-2)955-125, 2003.
Dan Geiger and David Heckerman. Learning Gaussian networks. In Uncertainty Proceedings 1994, pp. 235-243. Elsevier, 1994.
Geoffrey E Hinton and Russ R Salakhutdinov. A better way to pretrain deep boltzmann machines. In Advances in Neural Information Processing Systems (NeurIPS), pp. 2447-2455, 2012.
Patrik O. Hoyer, Dominik Janzing, Joris M Mooij, Jonas Peters, and Bernhard Schölkopf. Nonlinear causal discovery with additive noise models. In Advances in Neural Information Processing Systems (NeurIPS), 2009.
Biwei Huang, Kun Zhang. Yizhu Lin, Bernhard Schölkopf, and Clark Glymour. Generalized score functions for causal discovery. In Proceedings of the 24th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining (KDD), pp. 1551-, 2018.
Tommi Jaakkola, David Sontag, Amir Globerson, and Marina Meila. Learning Bayesian network structure using IP relaxations. In Proceedings of the Thirteenth International Conference Artificial Intelligence and Statistics (AISTATS), pp. 358-365, 2010.
Shantanu Jain, Martha White, and Predrag Radivojac. Recovering true classifier performance in positive-unlabeled learning. In Proceedings of the AAAI Conference Artificial Intelligence (AAAI), 2017.
Elias Khalil, Hanjun Dai, Yuyu Zhang, Bistra Dilkina, and Le Song. Learning combinatorial optimization algorithms over graphs. In Advances Neural Information Processing Systems (NeurIPS), 2017.
Daphne Koller and Nir Friedman. Probabilistic Graphical Models: Principles and Techniques. MIT Press, Cambridge, MA, USA, 2009.
Vijay Konda and John N. Tsitsiklis. Actor-critic algorithms. Advances in Neural Information Processing Systems (NeurIPS), 2000.
Wouter Kool, Herke Van Hoof, and Max Welling. International Conference on Learning Representations (ICLR), 2019.
Attention, learn to solve routing problems! In
Sébastien Lachapelle, Philippe Brouillard, Tristan Deleu, and Simon Lacoste-Julien. Gradient-based neural DAG learning. In International Conference on Learning Representations (ICLR), 2020.
Pedro Larranaga. Cindy MH Kuijpers, Roberto H Murga, and Yosu Yurramendi. Learning bayesian network structures by searching for the best ordering with genetic algorithms. IEEE Transactions Systems, Man, and Cybernetics-Pan Part A: Systems and Humans, 26(4):487 1996.
Lukas Meier, Sara Van de Geer, Peter Bühlmann, et al. High-dimensional additive modeling. The Annals of Statistics, 37(66)3377--388 2009.
Ignavier Ng. Zhuangyan Fang. Shengyu Zhu, Zhitang Chen, and Jun Wang. Masked gradient-basee causal structure leamiing arXiv preprint arXiv: 1910.0857, 2019a.
10

---

Under review as a conference paper at ICLR 2021
Ignavier Ng, Shengyu Zhu, Zhitang Chen, and Zhuangyan Fang. A graph autoencoder approach to causal structure learning. arXiv preprint arXiv:1911.07420 2019b.
Eric Perrier, Seiya Imoto, and Satoru Miyano. Finding optimal bayesian network given a super- structure. Journal of Machine Learning Research, 9(cct: 2008.
Jonas Peters and Peter Bühlmann. Identifiabilty of gaussian structural equation models with equal error variances. Biometrika, 101(1):219-22) 228, 2014.
Jonas Peters, Joris M. Mooij, Dominik Janzing, and Bernhard Schölkopf. Causal discovery with continuous additive noise models. Journal of Machine Learning Research, 15(1):2009-2053, January 2014. ISSN 1532-4435.
Garvesh Raskutti and Caroline Uhler. Learning directed acyelic graph models based sparsest permutations. Stat, 7(1):e183, 2018.
Pradeep Ravikumar, John Lafferty, Han Liu, and Larry Wasserman. Sparse additive models. Journal of the Royal Statistical Society: Series B (Statistical Methodology), 7165):009-10000, 2009.
Karen Sachs, Omar Perez, Dana Pe' er, Douglas A Lauffenburger, and Garry P Nolan. Causal protein- signaling networks derived from multiparameter single-cell data. Science, 308(5721):523-5222
2005.
Mauro Scanagatta, Cassio de Campos, Giorgio Corani, and Marco Zaffalon. Learning bayesian networks with thousands of variables. In Advances in Neural Information Processing Systems (NeurIPS), pp. 1864 1872, 2015.
Mark Schmidt, Alexandru Niculescu-Mizi,, Kevin Murphy, et al. Learning graphical model structure using Ll-regularization paths. In Proceedings of the AAAI Conference on Artificial Intelligence (AAAI), 2007.
Gideon Schwarz al. Estimating the dimension of model. The Annals of Statistics, 6(2):461-464 1978.
Shohei Shimizu, Patrik O Hoyer, Aapo Hyvärinen, and Antti Kerminen. A linear non-Gaussian acyclic model for causal discovery. Journal of Machine Learning Research, 700c:):0033-2330 2006.
Tomi Silander and Petri Myllymäki. simple approach for finding the globally optimal bayesian network structure. In Proceedings of the Twenty-Scoon Conference on Uncertainty in Artificial Intelligence (UAI), pp. 445-452, 2006.
Liam Solus, Yuhao Wang, and Caroline Uhler. Consistency guarantees for greedy permutation-based causal inference algorithms. arXiv preprint arXiv:1702.03530, 2017.
Peter Spirtes, Clark N Glymour, Richard Scheines, and David Heckerman. Causation, Prediction, and Search. MIT press, 2000.
Richard S Sutton and Andrew G Barto. Reinforcement Learning: An Introduction. MIT press, 2018.
Richard S Sutton, David A McAllester, Satinder P Singh, and Yishay Mansour. Policy gradient meth- ods for reinforcement learning with function approximation. Advances in Neural Information Processing Systems, 2000.
Mare Teyssier and Daphne Koller. Ordering-based search: A simple and effective algorithm for learning bayesian networks. In Conference Uncertainty in Artificial Intelligence (UAI), 2005.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in Neural Information Processing Systems (NeurIPS), 2017.
Oriol Vinyals, Meire Fortunato, and Navdeep Jaitly. Pointer networks. Advances in Neural Information Processing Systems (NeurIPS), 2015.

---

Under review as a conference paper at ICLR 2021
Jing Xiang and Seyoung Kim. A lasso for learning sparse bayesian network structure for continuous variables. In Advances in Neural Information Processing Systems (NeurIPS), 2013.
Yue Yu, Jie Chen, Tian Gao, and Mo Yu. DAG-GNN: DAG structure learning with graph neural networks. In International Conference on Machine Learning (ICML), 2019.
Jiji Zhang. On the completeness of orientation rules for causal discovery in the presence of latent confounders and selection bias. Artificial Intelligence, 172(16):1873 1896, 2008.
Kun Zhang and Aapo Hyvärinen. On the identifiability of the post-nonlinear causal model. In Conference on Uncertainty Artificial Intelligence (UAI), 2009.
Xun Zheng, Bryon Aragam, Pradeep K Ravikumar, and Eric P Xing. DAGs with NO TEARS: Continuous optimization for structure learning. In Advances in Neural Information Processing Systems (NeurIPS), 2018.
Shengyu Zhu, Ignavier Ng, and Zhitang Chen. Causal discovery with reinforcement learning. In International Conference on Learning Representations (ICLR), 2020.
Barret Zoph and Quoc Le. Neural architecture search with reinforcement learning. International Conference Learning Representations (ICLR), 2017.
12

---

Under review as a conference paper at ICLR 2021
APPENDIX
A NETWORK ARCHITECTURES AND PYPER-PARAMETERS
A.1 MULTIPLE NETWORK ARCHITECTURE DESIGNS FOR ENCODER AND DECODER
There are a variety of neural network modules for encoder and decoder structures. Here we consider some representative modules: including. Multi-layer Perceptrons (MLP) module, an LSTM based recurrent neural network module, and the self- attention based encoder in the Transformer. In addition, we use the original observation data as the state directly, i.e., no encoder module used, to show the necessity of encoder, which is denoted as Null. MLP consists of 3-layer feed-forward neural networks with 256, 512 and 256 units. The architecture of LSTM and Transformer are that introduced in Appendix A.2
The empirical results of CORL-2 on 30-node LG ER2 datasets are reported in Table 3 We observe that LSTM decoder achieves better performance than that of MLP decoder. This shows that LSTM more effective than MLP in sequential decision tasks. Besides, the overall performance of neural network encoder is better than that of Null, which shows that the data pre-processed by encoder module necessary. Among all these encoders, Transformer encoder achieves the best results. Similar conclusion was drawned in Zhu al. 2222)). We hypothesize that the performance of Transformer encoder is benefit from the self-attention scheme.
Table 3: Empirical results of CORL-2 with different choices of encoder and decoder 30-node LG ER2 datasets. The smaller SHD the better, the higher TPR the better.
Encoder
Null
LSTM
MLP
Transformer
MLP Decoder
TPR
0.81 (0.07)
0.86 (0.10)
0.96 (0.02)
0.98 (0.02)
SHD
54.2 (25.1)
36.0 (26.7)
(5.3)
5.0 (3.3)
LSTM Decoder
TPR
0.94 (0.04)
0.88 (0.09)
0.97 (0.01)
0.99 (0.01)
SHD
20.6 (20.0)
29.0 (17.8)
8.6 (4.3)
4.4 (3.5)
A.2 NETWORK ARCHITECTURE AND HYPER-PARAMETI PARAMETERS OF CORL.
fc: 256
fc: 1024
S1
fc: 256
fc: 256
multi-head
256
X
$j
fc: 256
Sx
input
block@3
output
Figure 5: Illustration of the Transformer encoder. The encoder embeds the observed data x, of each variable j into the state sy. Notation block @3 denotes three blocks here.
Both CORL-1 and CORL-2 use the actor-critic algorithm. We use the Adam optimizer with learning rate 1e-4 and 1 for actor and critic respectively. The discount factor y is set 0.98. The actor
13

---

Under review as a conference paper at ICLR 2021
fc: 256
hr
S1
256
41
fc: 1
S2
LSTM (256)
fc: 256
soft soft-mal -ax
Sd
h ht-1
ad
fc:
Figure 6: Illustration of the LSTM decoder. At each time step, maps the state & to a distribution over action space A = (a1,..., ad), then action (variable) can be selected randomly according to the distribution.
consists of an encoder and decoder. We illustrate the neural network structure of the Transformer encoder used our experiments Figure consists of rerd-forward layer with 256 units and three blocks. Each block is composed of multi-head attention network with 8 heads and 2-layer feed-forward neural networks with 1024 and 256 units. Each feed-forward layer is followed by normalization layer. Given a batch of observed samples with shape b x d x where b denotes the batch size, d the node number and n the number of observed data for each variable in a batch, the final output of the encoder is batch of embedded state with shape b d x 256.
We illustrate the neural network structure the decoder in Figure 6 which is mainly a LSTM similar to the decoder proposed by Vinyals al. 2015). The LSTM takes state input and outputs embedding. The embedding is mapped to the action space A by using some feed-forward neural networks, soft-max module and the pointer mechanism (Vinyals al. 2015). The outputs of the encoder are processed as the initial hidden state hy of the decoder. The LSTM with 256 hidden units used here. All of the feed-forward neural networks used decoder have 256 units.
The critic uses 3-layer feed-forward neural networks with 512, 256 and units. It takes a state as input and outputs predicted value for the current policy given state s. For CORL-I, the critic needs to predict the score for each state. For CORL-2, the critic takes the initial state šo as input and outputs predicted value directly for complete ordering.
B BASELINES AND DATE SETS
B.1 DETAILS BASELINES
The details of all the baselines considered in our experiments are listed follows:
ICA-LiNGAM assumes linear non-Gaussian additive model for data generating procedure and applics independent component analysis to recover the weighted adjacency matrix. This method can usually achieve good performance on LiNGAM datasets. However, it does not provide guarantee for linear Gaussian datasets.
NOTEARS recovers the causal graph by estimating the weighted adjacency matrix with the least squares loss and the smooth characterization for acyclicity constraint.
DAG-GNN formulates causal discovery in the framework of variational autoencoder and optimizes weighted adjacency matrix with the evidence lower bound and a modified smooth characterization on acyclicity as loss function.
httess//qiittbbb.co////////ee
https///gittubb.com/ffiiii
14

---

Under review as a conference paper at ICLR 2021
GraN-DAG models the conditional distribution of each variable given its parents with feed- forward neural networks. It uses the smooth directed acyclic constraint from NOTEARS to find DAG that maximizes the log- likelliddooo of the observed samples. |
RL-BIC2 formulates the causal discovery as a one-step decision making process and uses score function and acyclic constraint from NOTEARS to calculate the reward for the recovered directed graph. |
CAM conducts greedy estimation procedure that starts with an empty DAG and adds at each iteration the edge (rg,, 22) between nodes v2 and v/ that corresponds to the largest gain in log-likelihood For a searched ordering, they prune it to the final DAG by applying significance testing of covariates. They perform the preliminary neighborhood selection reduce the ordering space size searched. [
LIOBS performs heuristic search (greedy hill-climbing with tabu lists) through the space of topological orderings to search a ordering with the best score, then it uses I.1 variable selection to pruning the searched ordering (fully-connected DAG) to the final DAG. M
A Lasso with a limited queue size incorporates a heuristic scheme into dynamic pro- gramming based method. It first prunes the search space by using A Lasso, then it further pruning the search space by limiting the size of the priority queue in the OPEN list of A Lasso. The queue size usually needs to be adjusted to balance the time cost and the quality of the solution. 8
B.2 DETAILS ON DATASETS GENERATION
We generate synthetic datasets which vary along five dimensions: level of edge sparsity, graph type, number of nodes, causal functions and sample size. We sampled datasets of the required number examples for each task as follows: ground truth DAG g is drawn randomly from either the Erdös-Rényi (ER) or Scale-free (SF) graph model; then, the data is generated according to a specific sampling scheme.
Specifically, for Linear Gaussian (I.G) case, we set h E (2. 5) and d E (30, 50, 100) to obtain the ER graph and SF graph (different types) with different levels of edge sparsity and different number of nodes, respectively. Then datasets of 3000 examples are generated for each task following X X + where W E dddxd denotes the weight matrix which is obtained by assigning edge weights independently from U 0.5] U (0... Note that c's are standard Gaussian noises with equal variances for each variable as to LG data model is identifiable (Peters & Bühlmann 2014).
For LiNGAM data model, the datasets are generated in the similar way with I.G but the sampling for e are different. The non-Gaussian noises are obtained by following (Shiizuu et al. (2006) which passes the noise samples from Gaussian distribution through power nonlinearity make them non-Gaussian. LiNGAM is identifiable shown in Shimizu et al. (2006).
Another data generating process is GP. We first obtain graph with different density and different number of nodes. Then the datasets with different sample sizes are generated following X, = fy (Pa(X)))) + ej with jointly independent Gaussian noises for all where the function f, is function sampled from a GP with radial basis function kernel of bandwidth one and normal Gaussian noises. This setting is known to be identifiable according to Peters et al. (2014). Note that due to the efficiency of the reward calculation, we only experimented on nonlinear data models of up to 30 nodes.
C TOTAL NUMBER OF EPISODES BEFORE CONVERGENCE
Table 4 reports the total number of episodes required for CORL-2 and RL-BIC2 to be converged, averaged over five datasets. Note that the episodic reward is evaluated once per episode. CORL
St Shnccureg.aaannng////aaaaoooossp
/
15

---

Under review as a conference paper at ICLR 2021
formulates causal discovery multi-stage decision process and we observe that CORL performs fewer episodes than RL-BIC2 before the episode reward converges. We suspect that due to a significant reduction in the size of action space, the model complexity of the RL policy is reduced thus leading to higher sample efficiency. Some runtimes about them are also provided here (CORI-2 total runtime R: 15 minutes against RI-BIC2 & 3 hours for 30-node ER2 graphs, & 4 hours against 2 14 hour for 50-node ER2 graphs, and CORL-2 2 7 hours for 100-node ER2 graphs). We set the maximal runtime up to 24 hours, but RL-BIC2 did not converge within that time on 100-node graphs, hence we did not report it here. Note that these runtime may be significantly reduced by parallelizing the evaluation of reward.
Table 4: Total number of iterations (x10%) before RL converge on LG data.
30 nodes
50 nodes
100 nodes
ER2
ER5
ER2
ER5
ER2
ER5
CORL-2
1.0 (0.3)
1.1 (0.4)
1.9 (0.3)
2.4 (0.3)
2.3 (0.5)
2.9 (0.4)
RL-BIC2
3.9 (0.5)
4.1 (0.6)
3.4 (0.4)
3.5 (0.5)
x
D ADDITIONAL RESULTS ON LINEAR GAUSSIAN DATASETS
The results for 50-node LG data models are presented in Table The conclusion is similar to the 30- and 100-node experiments. The results of ICA-LiNGAM, GraN-DAG and CAM on LG data models presented in Table 6 Their performances do not compare favorably to CORL-1 nor CORL-2 in LG datasets.
It is not surprising that ICA-LiNGAM does not perform well because the algorithm is specifically designed for non-Gaussian noise and does not provide guarantee for LG data models. We hypothesize that CAM's poor performance on LG data models is because it uses nonlinear regression instead of linear regression. As for GraN-DAG, it uses 2-layer causal relationships, which may not be able learn good linear relationship in this experiment.
neural networks to model the
Table 5: Empirical results for ER and SF graphs of 50 nodes with LG data. The higher TPR the better, the smaller SHD the better.
RANDOM
NOTEARS
DAG-GNN
RL-BIC2
LLOBS
A" Lasso
CORL-1
CORL-2
ER2
TPR
0.31 (0.03)
0.94 (0.02)
0.94 (0.04)
0.79
0,56 (0.02)
0.55 (0.00)
0.97 (0.04)
0.97 (0-82)
205.4 (78.5)
(4.5))
(27.6)
24.0
(7.9))
ER5
TPR
0.32 (0.02)
0,90 (0.01)
0.87 (0.14)
(0,0))
0.57 (0.03)
0.82 (0.00)
0.90 (0.02)
0.92 (0.02)
SHD
378.4 (24.2)
(7.5)
(109.4)
128.9 (40.4)
299.4 (53.6)
104.0 (28.3)
68.3 (10.2)
(13.1)
TPR
0.40
0.99
0.90
(0.45)
0.80 (0.00)
1.00
SHD
(14.7)
(1.6)
79.3 (93.2)
(57.4)
182.5 (33.4)
(35.2)
(0.0)
(0.0)
TPR
0.51
0.94
0.88 (0.
(0,05)
(0.03)
0.81 (0.00)
0.94
0.95
SHD
245.6 (24.3)
20.1 (14.3)
89.2 (99.2)
115:2
217.3 (36.4)
(25.3)
24.3 (11.2)
20.8 (10.1)
E RESULTS ON 30-, 50- AND 100-NODE LINGAM DATASETS
Here, we report the empirical results on 30., 50- and 100-node LiNGAM datasets in Table 7 The observed samples are generated according to the same procedure with linear Gaussian datasets (see Appendix B.2 for for details). The non-Gaussian noise is obtained by passing the noise samples from Gaussian distribution through a power nonlinearity to make them non-Gaussian. For L1OBS, we increased the authors' recommended number of evaluations 2500 to 10000. For A* Lasso, set the queue size to 10, 500 and 1000, and report the best result for all of these different parameter settings. The results of L1OBS and A Lasso reported here are that of pruning using our pruning method. For other baselines, we pick the recommended
Among all these algorithms. ICA-LiNGAM can recover the true graph most of the LiNGAM data models. This is because ICA-LiNGAM is specifically designed for non- Gan-Gussian noise. CORL-1 and CORL-2 achieve consistently good results than other baselines.
16

---

Under review as a conference paper at ICLR 2021
Table 6: Empirical results of ICA-LINGAM, GraN-DAG and CAM (against CORL-2 for reference) for ER and SF graphs with LG data. The higher TPR the better, the smaller SHD the better.
ICA-LINGAM
GraN-DAG
CAM
CORL-2
30 modes
ER2
TPR
0.75 (0.03)
0.51 (0.17)
0.47 (0.05)
0.99 (0.01)
SHD
112.3 (12.8)
(11.))
110.8 (10.3)
4.4 (3.5))
ER5
TPR
0.57 (0.03)
0.52 (0.03)
0.46
0,95 (0.03)
SHD
161.8 (9.2)
175.2 (27.4)
191.3 (32.5)
37.6 (14.5)
SF2
TPR
0.58 (0.05)
0.61 (0.04)
0.63 (0.02)
10 (0.0)
SHD
(19.8)
136.4 (21.2)
115.2 (26.7)
0.0 (0.0)
SFS
TPR
0.56 (0.04)
0.58 (0.02)
0.60 (0.03)
(0.0)
SHD
160.5
142.4 (24.3)
122.2
0.0 (0.0)
nodes
(0.03)
(0.04)
(0.06)
(0.02)
SHD
108.8 (11.3)
173.0 (22.9)
140.8 (35.4)
17.9 (10.6)
ERS
TPR
0.57 (0.01)
0.64 (0.03)
0.61 (0.02)
0.92 (0.02)
SHD
199.8 (90.7)
154.2 (36.4)
178.3 (34.8)
64.8 (13.1)
SF2
TPR
0.59 (0.04)
0.44 (0.05)
0.57 (0.02)
1.00 (0.00)
(83.2)
(34.5)
(24.4)
(.00)
SF5
TPR
0.57 (0.01)
0.49 (0.04)
0.53 (0.03)
0.95 (0.02)
SHD
216.6 (88.4)
243.9 (27.2)
235.2 (34.2)
20.8 (10.1)
100 nodes
ER2
TPR
0.73 (0.02)
0.38 (0.02)
0.43 (0.02)
0,98 (0.01)
SHD
268.4 (28.5)
191.3 (31.9)
126.4 (27.8)
18.6 (5.7)
ER5
TPR
0.57 (0.05)
0.42 (0.03)
0.47 (0.02)
0.94 (0.03)
SHD
311.1 (63.7)
208.2 (54.4)
182.3 (34.9)
164.8 (17.1)
SF2
TPR
0.69 (0.03)
0.40 (0.03)
0.44 (0.02)
1.00 (0.00)
SHD
367.6 (67.5)
239.9 (43.2)
35.2 (37.4)
0.0 (0.0)
(0.05)
(0.03)
(0.04)
(0.01)
SHD
362.3 (82.8)
219.3 (32.2)
125.2 (24.7)
10.8 (6.1)
Table Empirical results on 30-, 50- and 100-node LiNGAM ER2 datasets. The smaller SHD the better, the higher TPR the better.
sodes ER2
50 nodes ER2
100 nodes ER2
Method
TPR
SHID
TPR
SHDD
TPR
SHD
ICA-LINGAM
1.00 (0.00)
0.0 (0.0)
1.00 (0.00)
0.0 (0.0)
1.00 (0.00)
1.0 (0.9)
NOTEARS
0,94 (0.04)
17.2
0,95 (0.02)
(1.25
(0.03)
69.2 (23.2)
DAG-GNK
0.94 (0.03)
19.6
0,96 (0.01)
24.6 (2.9)
0.93 (0.03)
66.2 (19.2)
GraN-DAG
0.28 (0.09)
100.8 (14.6)
0.20 (0.01)
177.0 (25.9)
0.16 (0.04)
312.8 (25.2)
RI-BIC
0.94 (0.07)
19.8 (23.0)
0.80 (0.08)
(51.9)
0.13 (0.12)
291.3 (24.1)
CAM
0.60 (0.
310.0 (34.0)
0.33 (0.07)
(31.9)
0.53 (0.05)
247.2 (32.
LIOBS
0.72 (0.04)
(23.3)
0.47 (0.02)
0.41 (0.03)
470.5
A Lasso
0.87 (0.03)
42.3 (16.3)
0,88 (0.03)
82.6 (17.6)
0.85 (0.04)
102.5 (22.6)
CORL-1
0.99 (0.01)
3.8 (6.4)
0.96 (0.06)
24.6 (37.7)
0.98 (0.01)
20.0 (7.9)
CORL-2
0.99 (0.01)
3.9 (5.6)
0.96 (0.08)
20.2 (11.3)
0.99 (0.01)
13.8 (7.2)
F RESULTS ON 20-NODE GP DATASETS WITH DIFFERENT SAMPLE SIZES
We take the 20-node GP data models different sample numbers. We set h = 4 to get ER4 graphs and generate data by using them. We illustrate the empirical results in Table g Since previous experiments have shown that CORL-2 is slightly better than CORL-1, we only report the results of CORL-2 here. CAM as the most competitive baseline, we also report its results on these datasets. TPR reported here is calculated based on the variable ordering. We can see that, as the sample size decreases, CORL-2 ends up outperforming CAM. We believe this is because CORL-2 benefits from the exploratory ability of RL.
example to show the performance of our method on
G CORL-1 WITH A PRETRAINED MODEL
We show the training reward curves of CORL-1 1 and CORL-I-pretree which is CORL-1 based on a pretrained model in Figure 7 To obtain a pretraining model with good generalization ability, we combine various types data described in Appendix B.2 with different levels of edge sparsity, graph
17

---

Under review as a conference paper at ICLR 2021
Table 8: Empirical results on 20-node GP ERI datasets with different sample sizes. The smaller SHD the better, the higher TPR the better.
Sample size
Name
TPR
SHD
1000
CAM
0.91 (0.03)
30.0 (3.7)
CORL-2
0.87 (0.03)
36.5 (3.1)
500
CAM
0.86 (0.03)
45.0 (2.5)
CORL-2
0.85 (0.03)
46.3 (2.3)
400
CAM
0.83 (0.02)
51.0 (2.7)
CORL-2
0.84 (0.03)
50.5 (3.0)
200
CAM
0.60 (0.03)
66.3 (1.9)
CORL-2
0.75 (0.02)
63.1 (1.5)
-4.6
-4.8
Episode reward
-5.4
556
-5.8
6600
CORL-1
CORL-1-pretran
-6.2
Episode
Figure 7: Learning curves of CORL-I and CORL-I-pretrait on 100-node LG datasets.
type and number of nodes and data generating process to construct observation samples. Next, we train a policy model by supervised learning on the mixed datasets. Finally, we use the pretrained model as start point on the task that has never been seen before during pretraining.
From Figure 7 we can observe that although the pretrained model is trained on other types of datasets, it can still accelerate the training as the initial model on the LG dataset task. This shows that our policy model may learn some implicit knowledge across different tasks.
18


Reviews:
[{"review_id": "NHEEi5CnxH9", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Update: I have read the authors responses and am happy that they have addressed some of my concerns in my review, especially comparisons with previous works. However, I still have the concern that the improvement seems very minor compared to previous works. Thus, I decide to keep my score unchanged.\n\nOriginal review:\n\nIn this paper, the authors proposed a new RL based algorithm to discover causal DAG models from observational data. Unlike the previous RL based causal discovery algorithm that performs causal discovery via searching through DAG space, it propose to instead search through the smaller ordering space, thereby achieving better empirical performance.\n\nCons.\n\n1. Ordering based causal discovery is a well-established field. In addition to the empirical findings cited in this manuscript, there are also [1,2] that theoretically understand ordering based causal discovery. New algorithms have also been proposed in the two papers. The authors should cite and compare against the two research works from both theoretical and empirical perspective, which is not the case in the current manuscript.\n\n2. In the variable selection part, since the algorithms perform multiple hypothesis testing, the p-values should be adjusted (such as bonferroni correction) accordingly to adapt to the multiple testing setting.\n\n3. Proposition 1 is trivial and has been proven / taken as known knowledge in many other places, such as [1].\n\n4. The methodological improvement seems not very novel, RL based causal discovery method and ordering based method have both been intensively studied before, and combining them together seems not a very exciting improvement from my perspective.\n\n[1] https://arxiv.org/abs/1702.03530\n[2] Bernstein, Daniel, et al. \"Ordering-based causal structure learning in the presence of latent variables.\" AISTATS 2020", "extra_scores": {}, "extra_text": {}}, {"review_id": "L0z2e5IYdIx", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The paper proposes a RL formulation to learn the DAG topological order (and then the DAG structure). The reduced search space enables a better and seemly more efficiently learning than the existing RL baselines. Empirical studies has confirmed the better graph learning accuracy results than the existing RL baseline as well as some non-RL structure learning algorithms. \n\nPositive:\n1. Formulation of RL order search is interesting.  Authors discussed the advantage of order-based approach over the existing RL Algorithms, which is convincing. The motivation hence is clear.\n2. Theoretical justification with Prop 1 justified the said order-based approach. \n2. The writing is generally easy to follow. Figure 3 study is appreciated. \n\nConcerns:\n\n1. it is not accurate to state that the exact algorithms can only handle 20.  For example, B&B in (Campos & Ji, 2011) can handle more than 30 nodes, and GOBNILP can regularly handle unto hundreds of variables. \n2. Prop 1: since the statement of graphs with the same distribution P(X) is made, is faithfulness assumed? \n3.  the output dimension of encoder is d times dimension of s? What is the input dimension of decoder and s hat at each step?\n4. Some important ablation study are missing, without them it is hard to judge the exact improvement of each components. For example, what is the performance of the pretrained model when combined with RL-BIC2? Would random sampling order instead of learning it via the proposed approach gives similar performance? How about the orders from other order-based learning approaches (as indicated by authors in the paper)?\n5. pre-training model: \"consider a scenario where the policy model learned on a data model task is used as the pretrained policy model for other tasks. Such results are demonstrated in Section 5.1\" - which part in Section 5.1? how the training datasets are generated? is it generated under the same setting as testing graphs? This seems critical on the evaluation. \n6. It could be beneficial that authors show the proposed approach learns a better order than some methods discussed in the paper (Schmidt et al and others). This is important to evaluate the power of RL based approach, since the post processing step is standard. \n7. Since the method uses complex structures such as transformer, an efficiency comparison with all methods is warranted. Specially considering many methods can be tuned via hyperparameters to obtain better solutions with longer time (such as hidden dimensions of neural methods, acyclicity threshold from NOTEAR based methods).\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "JtOFbuwVEig", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper proposes an RL-based method to learn the causal ordering of variables. Specifically, it formulates the ordering search problem as a Markov decision process, and then uses different reward designs to optimize the ordering generating model. Compared to [1], which uses RL to search in the DAG space, the proposed method is more efficient. \n\nPros:\n1. The presentation is easy to follow. \n2. The experimental results seem promising. \n3. The idea of searching the causal ordering, instead of in the DAG space,  to improve the efficiency is reasonable. \n\nCons:\n1. However, my main concern is that the proposed method lacks a theoretical guarantee of the identifiability of causal ordering, which is very important in the field of causal discovery. \n2. In addition, on linear data, the authors consider large-scale graphs, e.g., the graph which contains 100 nodes. I am wondering why the graph with 100 nodes is not tested on nonlinear data. \n3. Furthermore, the authors should also report the time complexity in order to show the improvement of efficiency. \n\nIf the authors can solve the above issues, I will increase the score. \n\nAnother point is that I am very surprised at the results shown in Table 2. It shows that ICA-LINGAM achieves the true graph even when there are 100 nodes. I am curious about which implementation the authors use.\n\n[1] Shengyu Zhu, Ignavier Ng, and Zhitang Chen. Causal discovery with reinforcement learning. In International Conference on Learning Representations (ICLR), 2020.\n\n\nPost-rebuttal:\nThanks for the feedback. Some of my concerns have been addressed, but I am still worried about the consistency issue and how to handle the large-scale problem. \n\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "7IDoCO4oqL", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper describes an RL approach for DAG structure learning. In particular, the RL agent explores the space of permutations instead of space of adjacency matrices. The experimental results suggest superior performance in a few synthetic data settings. \n\nPros: \n- Using RL for causal discovery is gaining popularity recently. But to the best of my knowledge this paper is the first to search over the space of variable orderings, instead of directly searching over the space of DAGs. \n- The experimental result seems to show a significant performance improvement compared to other baselines. \n\nCons:\n- The superior performance requires some explanation. For instance, there are many design choices made for the model, but it is hard to see which part of the model has brought the performance gain. It would be nice to see some ablation study. \n- Compared to the baselines, this paper proposes a substantially more complicated model. It is surprising that it does not overfit. But it would be nice to see some justifications (other than better TPR, SHD). \n\n\nDetailed comments: \nThis is a well written paper, but I would need more information to better evaluate it. \n- The RL formulation requires more explanation. For instance, why does the state s_j really represent? Why initialize s_0 with average states? Why does the state encoder use attention networks? What is the role of state transition, if all states depend on all X? \n- More diverse set of experiments would be useful: different density, different topological properties, different noise variances, and different nonlinearities, etc.\n- Since the architecture is substantially more complex than baselines, it would be nice to include computational costs like runtime. ", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: In this paper, the authors propose an RL-based method for learning DAGs based on searching over causal orders instead of graphs. Order search for learning DAGs is a well-studied problem, and it is well-known that this can relieve some of the burden of searching through the space of DAGs. Several reviewers raised legitimate concerns regarding the experiments,  and without identifiability or theoretical results to advance the state of the art, the contribution of this work is limited.
Decision: Reject