# Differentiable Tree Search Network

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

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Placed in Latent dynamics models and model-based RL, one of 11 directions in World models.
- Builds on 31 of its 38 references that are in the map.
- First posted on arXiv 22 Jan 2024.

## Placed in

- World models / [Latent dynamics models and model-based RL](https://mlsanity.com/wm/latent-dynamics-models-and-model-based-rl/)

## Builds on (most cited first, 30 of 31)

- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [Playing Atari with Deep Reinforcement Learning](https://mlsanity.com/paper/1312.5602/) (2013)
- [Trust Region Policy Optimization](https://mlsanity.com/paper/1502.05477/) (2015)
- [Conservative Q-Learning for Offline Reinforcement Learning](https://mlsanity.com/paper/2006.04779/) (2020)
- [Deep Reinforcement Learning with Double Q-learning](https://mlsanity.com/paper/1509.06461/) (2015)
- [Dream to Control: Learning Behaviors by Latent Imagination](https://mlsanity.com/paper/1912.01603/) (2019)
- [Mastering Atari, Go, Chess and Shogi by Planning with a Learned Model](https://mlsanity.com/paper/1911.08265/) (2019)
- [Learning Latent Dynamics for Planning from Pixels](https://mlsanity.com/paper/1811.04551/) (2018)
- [Dota 2 with Large Scale Deep Reinforcement Learning](https://mlsanity.com/paper/1912.06680/) (2019)
- [IMPALA: Scalable Distributed Deep-RL with Importance Weighted Actor-Learner Architectures](https://mlsanity.com/paper/1802.01561/) (2018)
- [Rainbow: Combining Improvements in Deep Reinforcement Learning](https://mlsanity.com/paper/1710.02298/) (2017)
- [Dueling Network Architectures for Deep Reinforcement Learning](https://mlsanity.com/paper/1511.06581/) (2015)
- [Recurrent World Models Facilitate Policy Evolution](https://mlsanity.com/paper/1809.01999/) (2018)
- [Deep Reinforcement Learning in a Handful of Trials using Probabilistic Dynamics Models](https://mlsanity.com/paper/1805.12114/) (2018)
- [Model-Based Reinforcement Learning for Atari](https://mlsanity.com/paper/1903.00374/) (2019)
- [Data-Efficient Reinforcement Learning with Self-Predictive Representations](https://mlsanity.com/paper/2007.05929/) (2020)
- [Agent57: Outperforming the Atari Human Benchmark](https://mlsanity.com/paper/2003.13350/) (2020)
- [Neural Network Dynamics for Model-Based Deep Reinforcement Learning with Model-Free Fine-Tuning](https://mlsanity.com/paper/1708.02596/) (2017)
- [Mastering Atari Games with Limited Data](https://mlsanity.com/paper/2111.00210/) (2021)
- [Value Iteration Networks](https://mlsanity.com/paper/1602.02867/) (2016)
- [Gradient Estimation Using Stochastic Computation Graphs](https://mlsanity.com/paper/1506.05254/) (2015)
- [Value Prediction Network](https://mlsanity.com/paper/1707.03497/) (2017)
- [The Predictron: End-To-End Learning and Planning](https://mlsanity.com/paper/1612.08810/) (2016)
- [An investigation of model-free planning](https://mlsanity.com/paper/1901.03559/) (2019)
- [Path Planning using Neural A* Search](https://mlsanity.com/paper/2009.07476/) (2020)
- [Learning to Search with MCTSnets](https://mlsanity.com/paper/1802.04697/) (2018)
- [Learning Differentiable Programs with Admissible Neural Heuristics](https://mlsanity.com/paper/2007.12101/) (2020)
- [Gated Path Planning Networks](https://mlsanity.com/paper/1806.06408/) (2018)
- [On Inductive Biases in Deep Reinforcement Learning](https://mlsanity.com/paper/1907.02908/) (2019)
- [Learning model-based planning from scratch](https://mlsanity.com/paper/1707.06170/) (2017)

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

- [Why LLMs Aren't Scientists Yet: Lessons from Four Autonomous Research Attempts](https://mlsanity.com/paper/2601.03315/) (2026)

## Abstract (arXiv)

In decision-making problems with limited training data, policy functions approximated using deep neural networks often exhibit suboptimal performance. An alternative approach involves learning a world model from the limited data and determining actions through online search. However, the performance is adversely affected by compounding errors arising from inaccuracies in the learned world model. While methods like TreeQN have attempted to address these inaccuracies by incorporating algorithmic inductive biases into the neural network architectures, the biases they introduce are often weak and insufficient for complex decision-making tasks. In this work, we introduce Differentiable Tree Search Network (D-TSN), a novel neural network architecture that significantly strengthens the inductive bias by embedding the algorithmic structure of a best-first online search algorithm. D-TSN employs a learned world model to conduct a fully differentiable online search. The world model is jointly optimized with the search algorithm, enabling the learning of a robust world model and mitigating the effect of prediction inaccuracies. Further, we note that a naive incorporation of best-first search could lead to a discontinuous loss function in the parameter space. We address this issue by adopting a stochastic tree expansion policy, formulating search tree expansion as another decision-making task, and introducing an effective variance reduction technique for the gradient computation. We evaluate D-TSN in an offline-RL setting with a limited training data scenario on Procgen games and grid navigation task, and demonstrate that D-TSN outperforms popular model-free and model-based baselines.