Learning Scalable Causal Discovery Policies with Adversarial Reinforcement Learning

We thank Reviewer L2qU for your useful feedback. Please see our responses below.

Q1: On the readability of the paper.

A1: We apologize for the confusion caused. The concept of 'task' was introduced in section 3 under 'Causal Discovery Task & BIC Score.' Actually, Section 3 provides some background of using RL for causal discovery. We would update our paper to provide more clear and comprehensive background information. Please check the new version of the paper later.

Q2: The most important aspect of this paper is the introduction of a method using adversarial samples to guide network training. Furthermore, it designs an AD agent to generate a ground-truth graph used as a reward. However, how can we ensure that the tasks generated by this AD agent are helpful to the learning of the OL agent?

A2: Referring to Algorithm 3 (Adversarial Training Framework), the OL agent and the AD agent are trained iteratively. The AD agent's reward is correlated with the OL agent's performance on adversarial tasks it generates at the previous interation. Thus, the AD agent is incentivized to generate challenging tasks that the OL agent struggles with. This adversatial training framework helps to improve the robustness of the OL agent so that it can solve unseen and poteitially larger tasks.

Our experimental results in Table 1&2 also support this: AGCORL, with adversarial training, outperformed CORL-P, which is trained using tasks generated from a non-adversarial prior distribution.

Q3: It's well-known that the convergence of these two methods is a significant challenge. Are there any techniques that can help the convergence of the algorithm?

A3: Thanks for raising this good question. Actually, we implemented a task pool strategy to stablize training (see Figure 1). We found that training the OL agent solely on tasks generated by the AD agent in the last epoch caused instability and forgetting of previously learned DAGs. By creating a task pool that stores all tasks and from which the OL agent samples for training, we significantly stabilized the training process. Figure 4-left illustrates that the training process is successfully stablized.

Q4: The generalizability of DAG learning is a rather werid concept. Why is it that actions can be generalized from a small graph to a larger one? Does this require certain assumptions? Do the causal mechanisms need to be consistent? Can actions learned on linear data be generalized to nonlinear data?

A4: The concept of generalizability in DAG learning parallels research in combinatorial optimization where a generalizable solver could solve larger problems than that used for training. Even in the context of causal discovery, this concept is not new. For example, existing works use supervised learning to train reusable networks for solving larger tasks thank training [1,2,3]. But there is no reusable RL based methods for causal discovery. Our work fills this gap by developing a reusable and generalizable RL policy. Specifically, our work focuses on generalizing from tasks with small DAGs to those with larger DAGs, under the assumption that the local patterns of structures and the underlying SCM are shared between small the larger DAGs. This also aligns with the existing works mentioned above. We will describe the concept of generalization more clearly in the paper. Note that the generalization is not guaranteed if the data distributions are largely different (e.g., data generated under linear SCM and non-linear SCM).

Q5: Why wasn't Notears-MLP compared for nonlinear data? In my experience, this method is actually more robust.

A5: We recognize Notears-MLP[4]'s robustness in certain nonlinear contexts. Yet, we compare our AGCORL with a stronger baseline GraN-DAG [5] in the non-lienar experiments and also use Gaussian process with additive noise to generate the data. GraN-DAG has been proven much more effective than Notears-MLP. According to [5], in terms of SHD metric, GraN-DAG scores significantly better than Notears-MLP (1.7 vs. 12.2 in 10ER1 and 8.3 vs. 32.6 in 10ER4 settings). These results align with our findings, therefore we choose GraN-DAG as the more appropriate and robust baseline in our work.

Q6: Do the authors consider this time comparison to be fair? You not only need to find the order but also apply methods like CAM for further pruning. As a DAG learning method, this time also needs to be counted in. Moreover, RL-based methods are indeed quite slow. It would be helpful if the authors could provide some training logs.

A6: The run times reported in Table 1&2 are the whole time solving for the tesing tasks, including the pruning procedure. Therefore, we believe this is a fair comparison.

Unlike existing RL methods such as CORL, our method learns a reusable causal discovery policy that requires only one forward inference during testing, which significantly saves time for solving multiple testing tasks. By contrast, non-reusable methods like Notears, Grad-DAG, CAM and CORL involve time-consuming procedures like gradient descent, heuristic and reward-guided searches for each testing task.

Q7: Maximizing the BIC score to learn a causal graph is theoretically grounded and can ensure identifiability. So, how can the method learned using the AD agent ensure that the ground-truth graph is learnable?

A7: Maximizing GTR indeed ensures that the true DAG order is identified. Note that GTR directly counts the reversed edges, so if there are wrong oders in inferred DAG, it will definitely be punished in the GTR score of the inferred DAG. Therefore, GTR is an effective and reliable metric that captures the order correctness of a DAG.

Since the adversarial DAGs are generated by the AD agent during training, we know the ground truth structure of these DAGs and therefore could easily compute the GTR to measure the distance of inferred DAGs and adversarial DAGs.

[1] Li et al. (2020). Supervised Whole DAG Causal Discovery.

[2] Lorch et al. (2022). Amortized Inference for Causal Structure Learning.

[3] Ke et al. (2022). Learning to Induce Causal Structure.

[4] Zheng et al. (2019). Learning Sparse Nonparametric DAGs.

[5] Lachapelle et al. (2020). Gradient-Based Neural DAG Learning.

Thanks for your positive response! I decided to increase my score from 3 to 5.

We thank the reviewer for your constructive feedback. Please see our responses below.

Q1: Comparing with CORL.

A1: While CORL also uses RL for causal discovery tasks, it is not scalable by design. For each given testing task, CORL will run a RL training process to find the optimal causal graph. By contrast, AGCORL uses adversarial tasks generated by the AD agent to learn a policy, which could directly output solutions for a given testing task (possibly larger than the training ones). Thus, since the AGCORL policy is reusable and scalable, it significantly reduces testing time compared to CORL. We will revise the related work section to more explicitly highlight these distinctions and their impact on runtime differences.

Q2: On the transferability.

A2: We thank the reviewer for raising this valuable question. Actually, experimental results in Table 6 and Section 5.3 show that AGCORL achieves comparable transferability compared to baselines. However, note that AGCORL focuses on scalability (generalization from smaller to larger graphs) and most experiments are designed to justify this point. In addition, it is a common practice in causal discovery to pre-assume some specific SCM types and learn from synthetic data, especially when we do not have any prior knowledge on the testing tasks. However, we aggree with the reviewer that transferability is also an important question to explore. We will leave it for future work.

Q3: Why do we use variable selection on the top of the algorithm, instead of directly test how good the ordering is relative to the true ordering?

A3: Actually, the mapping from a DAG to its ordering is neither injective nor surjective, meaning multiple correct orderings can exist for a given DAG. Therefore, it is hard to design a metric to directly evaluate the inferred orderings. Therefore, followed by existing works, we first construct fully-connected DAG using the infered orderings and then use standard pruning algorithms to get the final DAG. In practice, the pruning algorithms can work efficiently with only minor mistakes, since the solution space of DAG has been substantially reduced by the infered orderings. By applying a consistent pruning algorithm across all testing methods, the comparision is more fair and reliable.

Q4: On the contribution of the adversarial component: In Fig. 4 left, the authors show the impact of the adverarail training. However, the curves in Fig 4 do not in themselves suggest that adversarial training helps. What would be useful is training AGCORL with and without the adversarial parts and comparing the metrics. What we care about is whether the metrics at test time graphs improve or not.

A4: Fig 4 illustrates how adversatial training (AD agent) helps to improve the training of the OL agent. The experiments on AGCORL with and without the adversarial parts, as the reviewer mentioned, is shown in Table 1 and Table 2, where the CORL-P is actually AGCORL without adversarial part. To train CORL-P, we sample training tasks from a prior distribution without adversary. The empirical results show that the CORL-P performs worse than AGCORL, which demonstrates the contribution of the adversarial component. We will rename CORL-P as AGCORL-without-AD to avoid the confusion caused. Thanks for pointing out this.

Q5: Why do we use reinforcement learning?

A5: Compared with supervised learning based approaches which directly output a DAG solution for a given causal discovery task, RL based methods decompose the solution construction as a sequential decision making process. The decomposition of the task lowers the difficulty of the problem and potentially leads to higher quality of solutions. In addition, deterministic transition and episodic reward are very common in RL environments such as robotics. So this does not weaken the motivation of using RL.

Q6: On the runtime: In Table 2, why does AGCORL take 11m to run for a 30-node DAG? Whereas in Table 1, it takes 19.2s for a 100-node graph. Why is it slower for a 30-node DAG?

A6: The runtime difference between a 30-node DAG and a 100-node graph in AGCORL is primarily due to the nature of the data being processed. For nonlinear data, as in the 30-node case, the pruning process is more time-intensive compared to linear data, which is the case for the 100-node graph. Since AGCORL, along with CAM and CORL, first determines the node order and then applies a variable selection algorithm for pruning, the complexity of the data type significantly impacts the overall runtime. Hence, the 30-node nonlinear task requires more time than the 100-node linear task.

We thank Reviewer Wgj3 for your constructive feedback.

Q1: Using multiple OL agents compromise the fast running time, which is a main advantage of the proposed work.

A1: It is common in literature to assume some specific SCMs and generalize to real-world data, since the variable relationships are usually unknown. Actually, using multiple OL agents to infer candidate orderings is a practical approach to adapt to different real-world data types. While this may increase the computational load, the time complexity for inference remains efficient at $O(ld^2)$, where $l$ is the number of OL agents and $d$ is the task size. This is notably more efficient than other methods with complexities of $\Omega(d^3)$, especially for large $d$. Thus, our approach maintains a balance between adaptability to real-world data and the advantage of fast runtime, which is a key feature of our method.

Q2: Whether GTR can always be a surrogate for BIC.

A2: The key difference between GTR and BIC is that the computation of GTR requires ground-truth DAG, which is generated by the AD agent. Of course, in small scale tasks we can also apply AGCORL to achieve competitive performance (see Table 1&2). However, the necessity of appying AGCORL in small-scale tasks is not that significant, because AGCORL is designed for improving scalability.

Q3: Theorem 1 in Appendix A only shows that a random adversarial agent is the worst for training a generalizable order learning agent. It does not provide any indication of how good the authors' proposed AD agent is. For example, a theoretical analysis revealing how far the proposed AD agent's $\zeta$ is from 0.5, with some probability, would be helpful.

A3: Since the AD agent's $\zeta$ is closely related to experimental settings such as the number of nodes and SCMs, we can hardly estimate $\zeta$ accurately. However, Theorem 1 reveals that improving $\zeta$ by the AD agent could significantly reduce the number of required training tasks, hence successfully justifies the necessity of using the AD agent.

Q4: Is there any measure other than the Structural Hamming Distance (SHD) that can distinguish between missing edges and reversed edges? The authors argue that the OL agent is generalizable to large tasks based on the observation that the True Positive Rate (TPR) only decreases slightly (Table 4). However, in the same table, the SHD increases. If the SHD increases because of a large number of reversed edges, it is difficult to conclude that the OL agent generalizes to large tasks.

A4: The increase in SHD with larger graph sizes and more edges is a typical occurrence in causal discovery. To provide a more nuanced perspective, considering normalized SHD (SHD divided by the number of edges) offers a clearer view. This normalized metric exhibits only a slight increase on larger tasks, which is consistent with the changes observed in the True Positive Rate (TPR). This alignment suggests that the increase in SHD does not predominantly stem from a large number of reversed edges, supporting the generalizability of the OL agent to larger tasks.

Q5: Does AGCORL outperform previous BIC score-based algorithms on small-scale tasks too? If so, can we safely say that AGCORL replaces previous BIC score-based algorithms in all cases?

A5: In our experiments, AGCORL achieves competitive performance to existing BIC score-based algorithms on small-scale tasks (see Table 1&2). However, since the AGCORL is designed to improve scalability, we cannot guarantee that AGCORL outperforms baselines in all small-scale tasks. Therefore, it is more safe to try both algorithms in small-scale tasks and use AGCORL in large-scale tasks where BIC-based algorithms are not applicable.

Q6: For the same $d$ and $\theta$, how does the time complexity of computing the BIC score differ from that of Algorithm 1?

A6: In the linear case, computing the BIC score for one node with $n$ parents is $O(n^3)$, making the overall complexity for an order $O(d^4)$. By contrast, Algorithm 1 uses a Transformer model, which has a time complexity of $O(t^2)$ at step $t$, resulting in a total complexity of $O(d^3)$.

We thank Reviewer yLLZ for your helpful feedback. Please see our responses below.

Q1: Code is not shared so the results are irreproducible.

A1: We promise to share our code upon acceptance of this paper.

Q2: So many nodes but so little edges. E.g. Only 2 and 5 Table 1. and 1 to 4 with Table 2.

A2: The term 'edge' in our experiment refers to the average degree of a node in a DAG, not the total edge count. So a 50-node-2-edge graph will have 50*2/2=50 edges in average. We will clarify this by replacing 'edge' with 'degree' in the paper.

Q3: The construction of adversarial graph is unclear.

A3: In our approach, the construction of adversarial graphs is based on the causal principle that the data of a child node is generated from its parent nodes according to Structural Causal Models (SCMs). The AD agent creates these DAGs where $\mathbf{a}^i_{adv}$ specifies the parent nodes of the $i+1$-th node. For instance, in Fig 3, $\mathbf{a}^3_{adv} = {0,1,0}$ indicates that $X_4$ has $X_2$ as its sole parent node. Consequently, the data for $X_4$ is generated using the SCM: $X_4 = f(X_2) + \epsilon_4$. We apologize for any confusion caused by a labeling error in Fig 3, where $X_4$ was incorrectly marked as $X_2$. We will correct this in the figure caption to enhance clarity. This general procedure of adversarial graph construction underpins our methodology and aligns with established causal principles.

Q4: It would be good give high level intuition on how the employed pruning algorithm works.

A4: Our pruning algorithm, in essence, constructs a fully-connected DAG and then selectively removes edges to form the final DAG. This process involves evaluating the edge significance (e.g., using regression techniques), with a threshold to determine which edges to be pruned. For linear cases, we simple apply linear regression and a selected threshold (0.3 in our experiments). For nonlinear cases, we use a more complex pruning method outlined in CAM [1]. Since the pruning algorithms are not the focus of this work, we did not elaborate them in our paper. But we will add high-level descriptions of them in the deployment part. Thanks for your advice!

[1] Peter Bühlmann, Jonas Peters, and Jan Ernest. CAM: Causal additive models, high-dimensional order search and penalized regression. The Annals of Statistics, 42(6):2526–2556, 2014.