Improve Agents without Retraining: Parallel Tree Search with Off-Policy
  Correction
  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.
