# Tree Search in DAG Space with Model-based Reinforcement Learning for Causal Discovery

arXiv: https://arxiv.org/abs/2310.13576 · HTML: https://mlsanity.com/paper/2310.13576/

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 21 of its 64 references that are in the map.
- First posted on arXiv 20 Oct 2023.

## Builds on (most cited first, 21 of 21)

- [DAGs with NO TEARS: Continuous Optimization for Structure Learning](https://mlsanity.com/paper/1803.01422/) (2018)
- [Neural Combinatorial Optimization with Reinforcement Learning](https://mlsanity.com/paper/1611.09940/) (2016)
- [Learning Combinatorial Optimization Algorithms over Graphs](https://mlsanity.com/paper/1704.01665/) (2017)
- [Graph Convolutional Policy Network for Goal-Directed Molecular Graph Generation](https://mlsanity.com/paper/1806.02473/) (2018)
- [DAG-GNN: DAG Structure Learning with Graph Neural Networks](https://mlsanity.com/paper/1904.10098/) (2019)
- [Adversarial Attack on Graph Structured Data](https://mlsanity.com/paper/1806.02371/) (2018)
- [Thinking Fast and Slow with Deep Learning and Tree Search](https://mlsanity.com/paper/1705.08439/) (2017)
- [Learning Deep Generative Models of Graphs](https://mlsanity.com/paper/1803.03324/) (2018)
- [Gradient-Based Neural DAG Learning](https://mlsanity.com/paper/1906.02226/) (2019)
- [CAM: Causal additive models, high-dimensional order search and penalized regression](https://mlsanity.com/paper/1310.1533/) (2013)
- [Benchmarking Model-Based Reinforcement Learning](https://mlsanity.com/paper/1907.02057/) (2019)
- [Beware of the Simulated DAG! Causal Discovery Benchmarks May Be Easy To Game](https://mlsanity.com/paper/2102.13647/) (2021)
- [Causal Discovery with Reinforcement Learning](https://mlsanity.com/paper/1906.04477/) (2019)
- [Learning Scheduling Algorithms for Data Processing Clusters](https://mlsanity.com/paper/1810.01963/) (2018)
- [Ordering-Based Causal Discovery with Reinforcement Learning](https://mlsanity.com/paper/2105.06631/) (2021)
- [Differentiable DAG Sampling](https://mlsanity.com/paper/2203.08509/) (2022)
- [On the Convergence of Continuous Constrained Optimization for Structure Learning](https://mlsanity.com/paper/2011.11150/) (2020)
- [Goal-directed graph construction using reinforcement learning](https://mlsanity.com/paper/2001.11279/) (2020)
- [Computably Continuous Reinforcement-Learning Objectives are PAC-learnable](https://mlsanity.com/paper/2303.05518/) (2023)
- [DAG Learning on the Permutahedron](https://mlsanity.com/paper/2301.11898/) (2023)
- [Planning Spatial Networks with Monte Carlo Tree Search](https://mlsanity.com/paper/2106.06768/) (2021)

## Built on by (most cited first, 3 of 3)

- [Large Language Models are Effective Priors for Causal Graph Discovery](https://mlsanity.com/paper/2405.13551/) (2024)
- [Graph Reinforcement Learning for Combinatorial Optimization: A Survey and Unifying Perspective](https://mlsanity.com/paper/2404.06492/) (2024)
- [Accelerating Atomic Fine Structure Determination with Graph Reinforcement Learning](https://mlsanity.com/paper/2509.16184/) (2025)

## Abstract (arXiv)

Identifying causal structure is central to many fields ranging from strategic decision-making to biology and economics. In this work, we propose CD-UCT, a model-based reinforcement learning method for causal discovery based on tree search that builds directed acyclic graphs incrementally. We also formalize and prove the correctness of an efficient algorithm for excluding edges that would introduce cycles, which enables deeper discrete search and sampling in DAG space. The proposed method can be applied broadly to causal Bayesian networks with both discrete and continuous random variables. We conduct a comprehensive evaluation on synthetic and real-world datasets, showing that CD-UCT substantially outperforms the state-of-the-art model-free reinforcement learning technique and greedy search, constituting a promising advancement for combinatorial methods.