Paper: Learning Scalable Causal Discovery Policies with Adversarial Reinforcement Learning (ICLR-2024)

Reviewer comment (Writing):
I still do not understand what 'tasks' refers to in this paper.

Author response:
**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.

Reviewer comment (Evaluation):
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?

Author response:
**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.

Reviewer comment (Theory):
In the paper, a combination of adversarial learning and reinforcement learning is used for training. 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?

Author response:
**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.

Reviewer comment (Theory):
The generalizability of DAG learning is a rather weird 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?

Author response:
**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).

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

Author response:
**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.

Reviewer comment (Evaluation):
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.

Author response:
**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.

Reviewer comment (Theory):
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?

Author response:
**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.

Reviewer comment (Novelty):
The related work section does not explain the prior work CORL in enough detail, and the text does not make sufficient effort to explicitly distinguish the proposed work from CORL. It is not currently clear what the precise innovation is relative to CORL. This also makes it harder to interpret why AGCORL and CORL have such different runtimes.

Author response:
**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.

Reviewer comment (Experiments):
In the experiments, the authors only attempt to test generalization to larger graph instances. A limitation of AGCORL is that the training data itself is synthetic (limited to linear Gaussian or GP SCMs), and the authors only test on the same SCM distributions. It is unclear what would happen if a different distribution of SCMs were used at test time (e.g., different GP kernel from training), or if the graph distribution at test time is not ER or scale-free. More experiments understanding generalization would be useful.

Author response:
**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.  

Re response to transferability: 

Thank you for acknowledging the value of our work in demonstrating generalizability to larger graph structures. We recognize that this represents only one dimension of generalization. Indeed, training exclusively on synthetic data, as in our method, may inherently limit generalization to a specific set of Structural Causal Models (SCMs). This raises valid concerns, such as how our method would perform when encountering linear SCMs with parameter variations not seen during training.

While it's true that encountering significantly different SCM parameters could impact performance, we believe this challenge can be mitigated by allocating more computational resources to train the agent across a broader spectrum of SCM distributions. Such an approach could enhance the method's adaptability to diverse scenarios. Thus, while our current focus is on graph structure generalizability, future iterations of our method, with adequate computational support, hold the potential to address a wider range of generalization challenges, including those highlighted in your query.

Reviewer comment (Evaluation):
Since AGCORL only learns the topological ordering, the authors do not use metrics that directly test how good the ordering is relative to the true ordering. Using variable selection on top of the learned ordering makes the results harder to interpret, as it is unclear whether errors are due to the wrong ordering or wrong variable selection.

Author response:
**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.  

Re metrics and pruning: 

Thank you for your suggestions regarding the evaluation metrics for DAG ordering. However, our primary objective is the accurate identification of DAGs rather than their specific orderings. This focus aligns with the majority of prior work in the field, which typically emphasizes DAG accuracy over ordering precision. Besides, there exist non-ordering-based methods in baselines.  Consequently, we have chosen to prioritize metrics that assess the correctness of the DAGs themselves, in line with these established methodologies.

Regarding nonlinear data, we acknowledge the increased difficulty in order inference. Despite these challenges, our method has demonstrated superior performance in comparison to all baselines in handling such data. This suggests that, while the complexity of the task increases with nonlinear data, our approach remains effective in accurately identifying the underlying DAG structures.

Reviewer comment (Experiments):
In Figure 4 (left), the authors show the impact of adversarial training, but the curves do not suggest that adversarial training helps. A more useful approach would be to train AGCORL with and without the adversarial parts and compare the metrics to see whether performance on test-time graphs improves.

Author response:
**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.

Reviewer comment (Theory):
Regarding the reinforcement learning formulation, it is unclear why the authors chose to formulate this as an RL problem instead of a supervised learning task. The state is a tuple of $< s^{+}, s^{-} >$, meaning that for a given action $a$, the next state can be deterministically computed, which differs from a standard RL setup where states are stochastically generated using $p(s_{t+1}|s_t, a_t)$. Additionally, since a reward is only received after the full ordering is produced, it seems this could be treated as a supervised problem where an ordering is directly output, with many existing losses available for this task. A supervised formulation might be more efficient given that the RL approach is essentially zeroth-order optimization.

Author response:
**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.

Reviewer comment (Evaluation):
In Table 2, AGCORL takes 11 minutes to run for a 30-node DAG, whereas in Table 1 it takes 19.2 seconds for a 100-node graph. Given that at inference time the trained policy agent should only require d-steps to compute the ordering, it is unclear why it is slower for a 30-node DAG.

Author response:
**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.

Reviewer comment (Experiments):
The OL agent is trained on synthetic data generated by an explicit Structural Causal Model (SCM). Therefore, the OL agent may not be suitable for real-world data tasks where we do not know the mathematical relationship between the variables. In such cases, the authors suggest using multiple OL agents trained with different function types to infer candidate orderings, and then calculate their BIC score to select the best one (Appendix E). However, this approach could compromise the fast running time, which is a main advantage of the proposed work.

Author response:
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.

Reviewer comment (Theory):
The paper lacks theoretical and/or qualitative explanations that compare the BIC score and the newly proposed GTR. It seems important to clarify whether GTR can always be a surrogate for BIC. If so, AGCORL could replace frameworks that calculate the BIC score, even for small-scale tasks. If not, in certain circumstances, it will still be necessary to use BIC score based algorithms, even for large-scale tasks.

Author response:
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.

Reviewer comment (Evaluation):
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.

Author response:
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.

Reviewer comment (Experiments):
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?

Author response:
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.

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

Author response:
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)$.

Reviewer comment (Reproducibility):
Code is not shared so the results are irreproducible.

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

Reviewer comment (Experiments):
So many nodes but so little edges. E.g. Only 2 and 5 Table 1. and 1 to 4 with Table 2. This makes me skeptical about the relevance of the results. (See questions)

Author response:
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.

Reviewer comment (Presentation):
The construction of adversarial graph is unclear to me. Caption of Figure 3 attempts to explain how does it work, but it is not clear to me still. It should be improved. (See the question) In general would be great to write down the general procedure. (if it's already there, apologies).

Author response:
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.

Reviewer comment (Writing):
It would be good give high level intuition on how the employed pruning algorithm works (especially at the end of the deployment subsection). Currently, just referring to the appendix without sharing the idea obscures it.

Author response:
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.

Reviewer comment (Experiments):
So many nodes but so little edges. E.g. Only 2 and 5 Table 1. and 1 to 4 with Table 2. I wonder how does your results are when it comes to different sparsity, node vs edge is more balanced, or dense?

Author response:
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.