mlsanityData to 7 Oct 2026

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

arXiv 2107.01715 ↗First posted 4 Jul 2021cs.AIcited by 10

Not on the map. Why?

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.

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

Builds on

25 of 44 references.

References · 25

  1. 2013Playing Atari with Deep Reinforcement Learningcited by 3,211
  2. 2018Addressing Function Approximation Error in Actor-Critic Methodscited by 2,302
  3. 2019Mastering Atari, Go, Chess and Shogi by Planning with a Learned Modelcited by 1,327
  4. 2018IMPALA: Scalable Distributed Deep-RL with Importance Weighted Actor-Learner Architecturescited by 866
  5. 2017Rainbow: Combining Improvements in Deep Reinforcement Learningcited by 774
  6. 2017Mastering Chess and Shogi by Self-Play with a General Reinforcement Learning Algorithmcited by 723
  7. 2019Model-Based Reinforcement Learning for Ataricited by 495
  8. 2017Neural Network Dynamics for Model-Based Deep Reinforcement Learning with Model-Free Fine-Tuningcited by 288
  9. 2017Revisiting the Arcade Learning Environment: Evaluation Protocols and Open Problems for General Agentscited by 267
  10. 2016Safe and Efficient Off-Policy Reinforcement Learningcited by 260
  11. 2014cuDNN: Efficient Primitives for Deep Learningcited by 240
  12. 2015Action-Conditional Video Prediction using Deep Networks in Atari Gamescited by 171
13 more
  1. 2018Sample-Efficient Reinforcement Learning with Stochastic Ensemble Value Expansioncited by 167
  2. 2015Deep Kalman Filterscited by 163
  3. 2018Model-Based Value Estimation for Efficient Model-Free Reinforcement Learningcited by 152
  4. 2017Imagination-Augmented Agents for Deep Reinforcement Learningcited by 137
  5. 2021OptiDICE: Offline Policy Optimization via Stationary Distribution Correction Estimationcited by 121
  6. 2018Generalization and Regularization in DQNcited by 117
  7. 2021Online and Offline Reinforcement Learning by Planning with a Learned Modelcited by 90
  8. 2020Learning to Simulate Dynamic Environments with GameGANcited by 86
  9. 2018GPU-Accelerated Robotic Simulation for Distributed Reinforcement Learningcited by 77
  10. 2019Safe Reinforcement Learning Using Robust MPCcited by 72
  11. 2018Beyond the One Step Greedy Approach in Reinforcement Learningcited by 20
  12. 2018How to Combine Tree-Search Methods in Reinforcement Learningcited by 16
  13. 2020Think Too Fast Nor Too Slow: The Computational Trade-off Between Planning And Reinforcement Learningcited by 5

Built on by

Cited by year

Citing papers per year: 2022: 3; 2023: 3; 2024: 3; 2025: 1.

10 papers.

  1. 2024Inference-Aware Fine-Tuning for Best-of-N Sampling in Large Language Modelscited by 86
  2. 2022Generalised Policy Improvement with Geometric Policy Compositioncited by 11
  3. 2022Planning and Learning with Adaptive Lookaheadcited by 8
  4. 2024Policy Mirror Descent with Lookaheadcited by 6
  5. 2024SPO: Sequential Monte Carlo Policy Optimisationcited by 4
  6. 2023Policy Gradient with Tree Expansioncited by 3
  7. 2023The Effective Horizon Explains Deep RL Performance in Stochastic Environmentscited by 3
  8. 2025Counting Hours, Counting Losses: The Toll of Unpredictable Work Schedules on Financial Securitycited by 2
  9. 2022SoftTreeMax: Policy Gradient with Tree Searchcited by 1
  10. 2023Acceleration in Policy Optimizationcited by 0