Improve Agents without Retraining: Parallel Tree Search with Off-Policy Correction
This work addresses usability issues in Tree Search for reinforcement learning practitioners, offering significant performance gains and scalability improvements, though it is incremental in building on existing TS methods.
The paper tackles two challenges in Tree Search (TS) for reinforcement learning—distribution shift and scalability—by introducing an off-policy correction term to improve pre-trained agents without retraining and a GPU-based Batch-BFS method to speed up computation. The correction more than doubles scores on Atari games for pre-trained Rainbow agents, and Batch-BFS reduces runtime by two orders of magnitude while enabling training with deeper trees.
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}.