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

TL;DR

This paper introduces a novel off-policy correction for tree search in reinforcement learning, improving pre-trained agents without retraining and addressing scalability with a GPU-based breadth-first search method, leading to significant performance gains.

Contribution

It proposes an off-policy correction method to fix distribution shift issues in tree search and introduces Batch-BFS for scalable, efficient tree search on GPUs, enabling improved agent performance.

Findings

Significantly improved pre-trained Rainbow agents without retraining.

Enhanced scalability allowing deeper tree searches with GPU acceleration.

Improved DQN agents on Atari games using tree search techniques.

Abstract

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: \textit{distribution shift} and \textit{scalability}. 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 \url{https://github.com/NVlabs/bcts}.