# Improve Agents without Retraining: Parallel Tree Search with Off-Policy Correction

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

## Key facts

- Cited by 10 papers in the map (data through 14 Oct 2026).
- Builds on 25 of its 44 references that are in the map.
- First posted on arXiv 4 Jul 2021.

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

- [Playing Atari with Deep Reinforcement Learning](https://mlsanity.com/paper/1312.5602/) (2013)
- [Addressing Function Approximation Error in Actor-Critic Methods](https://mlsanity.com/paper/1802.09477/) (2018)
- [Mastering Atari, Go, Chess and Shogi by Planning with a Learned Model](https://mlsanity.com/paper/1911.08265/) (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)
- [Mastering Chess and Shogi by Self-Play with a General Reinforcement Learning Algorithm](https://mlsanity.com/paper/1712.01815/) (2017)
- [Model-Based Reinforcement Learning for Atari](https://mlsanity.com/paper/1903.00374/) (2019)
- [Neural Network Dynamics for Model-Based Deep Reinforcement Learning with Model-Free Fine-Tuning](https://mlsanity.com/paper/1708.02596/) (2017)
- [Revisiting the Arcade Learning Environment: Evaluation Protocols and Open Problems for General Agents](https://mlsanity.com/paper/1709.06009/) (2017)
- [Safe and Efficient Off-Policy Reinforcement Learning](https://mlsanity.com/paper/1606.02647/) (2016)
- [cuDNN: Efficient Primitives for Deep Learning](https://mlsanity.com/paper/1410.0759/) (2014)
- [Action-Conditional Video Prediction using Deep Networks in Atari Games](https://mlsanity.com/paper/1507.08750/) (2015)
- [Sample-Efficient Reinforcement Learning with Stochastic Ensemble Value Expansion](https://mlsanity.com/paper/1807.01675/) (2018)
- [Deep Kalman Filters](https://mlsanity.com/paper/1511.05121/) (2015)
- [Model-Based Value Estimation for Efficient Model-Free Reinforcement Learning](https://mlsanity.com/paper/1803.00101/) (2018)
- [Imagination-Augmented Agents for Deep Reinforcement Learning](https://mlsanity.com/paper/1707.06203/) (2017)
- [OptiDICE: Offline Policy Optimization via Stationary Distribution Correction Estimation](https://mlsanity.com/paper/2106.10783/) (2021)
- [Generalization and Regularization in DQN](https://mlsanity.com/paper/1810.00123/) (2018)
- [GPU-Accelerated Robotic Simulation for Distributed Reinforcement Learning](https://mlsanity.com/paper/1810.05762/) (2018)
- [Online and Offline Reinforcement Learning by Planning with a Learned Model](https://mlsanity.com/paper/2104.06294/) (2021)
- [Learning to Simulate Dynamic Environments with GameGAN](https://mlsanity.com/paper/2005.12126/) (2020)
- [Safe Reinforcement Learning Using Robust MPC](https://mlsanity.com/paper/1906.04005/) (2019)
- [Beyond the One Step Greedy Approach in Reinforcement Learning](https://mlsanity.com/paper/1802.03654/) (2018)
- [How to Combine Tree-Search Methods in Reinforcement Learning](https://mlsanity.com/paper/1809.01843/) (2018)
- [Think Too Fast Nor Too Slow: The Computational Trade-off Between Planning And Reinforcement Learning](https://mlsanity.com/paper/2005.07404/) (2020)

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

- [Inference-Aware Fine-Tuning for Best-of-N Sampling in Large Language Models](https://mlsanity.com/paper/2412.15287/) (2024)
- [Generalised Policy Improvement with Geometric Policy Composition](https://mlsanity.com/paper/2206.08736/) (2022)
- [Planning and Learning with Adaptive Lookahead](https://mlsanity.com/paper/2201.12403/) (2022)
- [Policy Mirror Descent with Lookahead](https://mlsanity.com/paper/2403.14156/) (2024)
- [SPO: Sequential Monte Carlo Policy Optimisation](https://mlsanity.com/paper/2402.07963/) (2024)
- [Policy Gradient with Tree Expansion](https://mlsanity.com/paper/2301.13236/) (2023)
- [The Effective Horizon Explains Deep RL Performance in Stochastic Environments](https://mlsanity.com/paper/2312.08369/) (2023)
- [Counting Hours, Counting Losses: The Toll of Unpredictable Work Schedules on Financial Security](https://mlsanity.com/paper/2504.07719/) (2025)
- [SoftTreeMax: Policy Gradient with Tree Search](https://mlsanity.com/paper/2209.13966/) (2022)
- [Acceleration in Policy Optimization](https://mlsanity.com/paper/2306.10587/) (2023)

## Abstract (arXiv)

Tree Search (TS) is crucial to some of the most influential successes in reinforcement learning. Here, we tackle two major challenges with TS that limit its usability: textitdistribution shift and textitscalability. We first discover and analyze a counter-intuitive phenomenon: action selection through TS and a pre-trained value function often leads to lower performance compared to the original pre-trained agent, even when having access to the exact state and reward in future steps. We show this is due to a distribution shift to areas where value estimates are highly inaccurate and analyze this effect using Extreme Value theory. To overcome this problem, we introduce a novel off-policy correction term that accounts for the mismatch between the pre-trained value and its corresponding TS policy by penalizing under-sampled trajectories. We prove that our correction eliminates the above mismatch and bound the probability of sub-optimal action selection. Our correction significantly improves pre-trained Rainbow agents without any further training, often more than doubling their scores on Atari games. Next, we address the scalability issue given by the computational complexity of exhaustive TS that scales exponentially with the tree depth. We introduce Batch-BFS: a GPU breadth-first search that advances all nodes in each depth of the tree simultaneously. Batch-BFS reduces runtime by two orders of magnitude and, beyond inference, enables also training with TS of depths that were not feasible before. We train DQN agents from scratch using TS and show improvement in several Atari games compared to both the original DQN and the more advanced Rainbow. The code for BCTS can be found in urlhttps://github.com/NVlabs/bcts.