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

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: distribution shift and 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.

heading: 1 Introduction
text: Tree search (TS) is a fundamental component of Reinforcement Learning (RL) [46] used in some of the most successful RL systems [42, 44]. For instance, Monte-Carlo TS (MCTS) [10] achieved superhuman performance in board games like Go [44], Chess [45], and Bridge [5]. MCTS gradually unfolds the tree by adding nodes and visitation counts online and storing them in memory for future traversals. This paradigm is suitable for discrete state-spaces where counts are aggregated across
∗Equal contribution (random order)
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
multiple iterations as the tree is built node-by-node, but less suitable for continuous state-spaces or image-based domains like robotics and autonomous driving 1. For the same reason, MCTS cannot be applied to improve pre-trained agents without collecting their visitation statistics in training iterations.
Instead, in this work, we conduct the TS “on-demand” by expanding the tree up to a given depth at each state. Our approach handles continuous and large state-spaces like images without requiring any memorization. This on-demand TS can be performed both at training or inference time. Here, we focus our attention on the second case, which leads to score improvement without any re-training. This allows one to better utilize existing pre-trained agents even without having the ability or resources to train them. For example, a single AlphaGo training run is estimated to cost 35 million USD [1]. In other cases, even when compute budget is not a limitation, the setup itself makes training inaccessible. For example, when models are distributed to end-clients with too few computational resources to train agents in their local custom environments.
We run TS for inference as follows. For action selection, we feed the states at the leaves of the spanned tree to the pre-trained value function. We then choose the action at the root according to the branch with the highest discounted sum of rewards and value at the leaves. Our approach instantly improves the scores of agents that were already trained for long periods (see Sec. 5.1). Often, such improvement is possible because the value function is not fully realizable with a function approximator, e.g., a deep neural network; TS can then overcome the limitation of the model. In practice, TS requires access to a forward model that is fed with actions to advance states and produce rewards. Here, we build on the recently published CuLE [11] – an Atari emulator that runs on GPU. This allows us to isolate the fundamental properties of TS without the added noise of learned models such as those described in Sec.6.
Performing TS on-demand has many benefits, but it also faces limitations. We identify and analyze two major obstacles: distribution shift and scalability.
First, we report a counter-intuitive phenomenon when applying TS to pre-trained agents. As TS looks into the future, thus utilizing more information from the environment, one might expect that searching deeper should yield better scores. Surprisingly, we find that in many cases, the opposite happens: action selection based on vanilla TS can drastically impair performance. We show that performance deteriorates due to a distribution shift from the original pre-trained policy to its corresponding treebased policy. We analyze this phenomenon by quantifying the probability of choosing a sub-optimal action when the value function error is high. This occurs because for values of out-of-distribution states, larger variance translates to a larger bias of the maximum. Our analysis leads to a simple, computationally effective off-policy correction term based on the Bellman error. We refer to the resulting TS as the Bellman-Corrected Tree-Search (BCTS) algorithm. BCTS yields monotonically improving scores as the tree depth increases. In several Atari games, BCTS even more than doubles the scores of pre-trained Rainbow agents [22].
The second limitation is scalability: the tree grows exponentially with its depth, making the search process computationally intensive and limiting the horizon of forward-simulation steps. To overcome this limitation, we propose Batch-BFS: a parallel GPU adaptation of Breadth-First Search (BFS), which brings the runtime down to a practical regime. We measured orders-of-magnitude speed-up compared to alternative approaches. Thus, in addition to improving inference, it also enables training tree-based agents in the same order of training time without a tree. By combining Batch-BFS with DQN [32] and training it with multiple depths, we achieve performance comparable or superior to Rainbow – one of the highest-scoring variants of the DQN-based algorithm family.
Our Contributions. (1) We identify and analyze a distribution-shift that impairs post-training TS. (2) We introduce a correction mechanism and use it to devise BCTS: an efficient algorithm that improves pre-trained agents, often doubling their scores or more. (3) We create Batch-BFS, an efficient TS on GPU. (4) We use Batch-BFS to train tree-based DQN agents and obtain higher scores than DQN and Rainbow.
1In the case of continuous state-spaces or image-based domains, MCTS can be used by reconstructing trajectories from action sequences only in deterministic environments. Also, MCTS requires a fixed initial state because the root state has to either exist and be found in the MCTS tree or alternatively a new tree has to be built and reiterated for that initial state.

heading: 2 Preliminaries
text: Our framework is an infinite-horizon discounted Markov Decision Process (MDP) [39]. An MDP is defined as the 5-tuple (S,A, P, r, γ), where S is a state space, A is a finite action space, P (s′|s, a) is a transition kernel, r(s, a) is a reward function, γ ∈ (0, 1) is a discount factor. At each step t = 0, 1, . . . , the agent observes the last state st, performs an action at and receives a reward rt. The next state is then sampled by st+1 ∼ P (·|st, at). For brevity, we denote A := |A|. Let π : S → A be a stationary policy. Let Qπ : S × A → R be the state-action value of a policy π, defined in state s as Qπ(s, a) ≡ Eπ [∑∞ t=0 γ tr(st, π(st)) ∣∣s0 = s, a0 = a], where Eπ denotes expectation w.r.t. the distribution induced by π. Our goal is to find a policy π∗ yielding the optimal value Q∗ such that Q∗(s, a) = maxπ r(s, a) + γEs′∼P (·|s,a) maxa′ Qπ(s′, a′). It is well known that
Q∗(s, a) = r(s, a) + γEs′∼P (·|s,a) max a′ Q∗(s′, a′), π∗(s) = arg max a Q∗(s, a).
Vanilla tree search. To ease notations and make the results concise, we limit the analysis to deterministic transitions2, i.e., an action sequence (a0, . . . , ad−1), starting at s0 leads to a corresponding trajectory (s0, . . . , sd). Nonetheless, the results can be extended to a stochastic setup by working with the marginal probability over the trajectory. Then, for a policy πo, let the d-step Q-function
Qπod (s, a) = max (ak)dk=1∈A [ d−1∑ t=0 γtr(st, at) + γ dQπo(sd, ad) ] s0=s,a0=a , (1)
and similarly let Q̂πod (s, a) be the d-step Q-function estimator that uses an estimated Q-function Q̂ πo instead of Qπo . Finally, denote by πd the d-step greedy policy
πd(s) := arg max a∈A
Q̂πod (s, a). (2)

heading: 3 Solving the Tree Search Distribution Shift
text: In this section, we show how to leverage TS to improve a pre-trained policy. We start by demonstrating an intriguing phenomenon: The quality of agents degrades when using vanilla TS. We then analyze the problem and devise a solution. The core idea of our approach is to distinguish between actions that are truly good, and those that are within the range of noise. By quantifying the noise using the Bellman error and problem parameters, we find the exact debiasing that yields the optimal signal-to-noise separation.

heading: 3.1 Performance degradation with vanilla tree search
text: We focus on applying TS at inference time, without learning. We begin with a simple experiment that quantifies the benefit of using TS given a pre-trained policy.
A TS policy has access to future states and rewards and, by definition, is optimal when the tree depth goes to infinity. Hence, intuitively, one may expect a TS policy to improve upon πo for finite depths as well. To test this, we load Rainbow agents Q̂πo , pre-trained on 50M frames; they are publicly available in [24] and achieve superhuman scores as in [22]. We use them to test TS policies πd with multiple depths d on several Atari benchmarks. Surprisingly, the results (red curves in Fig. 5, Sec. 5.1) show that TS reduces the total reward, sometimes to scores of random policies, in various games and depths. The drop is particularly severe for TS policies with d = 1 — a fact later explained by our analysis in Thm. 3.5.
We find the reason for this performance drop to be the poor generalization of the value function to states outside the stationary πo’s distribution. Fig. 1 shows a typical, bad action selection in Atari Breakout by a depth-1 TS. The table on the right reports the estimated Q-values of the root state (first column) and of every state at depth 1 (last four columns). Since the ball is dropping, ‘Left’ is the optimal action. Indeed, this corresponds with the Q-values at the root. However, at depth 1,
2The Atari environments we experiment on here are indeed close to deterministic. Their source of randomness is the usage of a random number of initial noop actions [33].
the Q-values of the future state that corresponds to choosing ‘Left’ (second column) are the lowest among all depth-1 future states. Subsequently, the depth-1 TS policy selects ‘Right’.
During training, towards convergence, the trained policy mostly selects ‘Left’ while other actions are rarely sampled. Therefore, expanding the tree at inference time generates states that have been hardly observed during training and are consequently characterized by inaccurate Q-value estimates. In the case of Fig. 1, the Q-value for ‘Left’ should indeed be low because the agent is about to lose the game. As for the other, poorly sampled states, regression towards a higher mean leads to over-estimated Q-values. Similar poor generalization has been observed in previous studies [16] and is interpreted as an off-policy distribution shift [35].
Beyond the anecdotal example in Fig. 1, additional evidence supports our interpretation regarding the distribution shift. We first consider the error in the value estimate captured by the Bellman error minimized in training. We compare the average Bellman error of the action chosen by πo to all other actions at the tree root. When averaging over 200 episodes, we find that the error for actions chosen by πo is consistently lower than for other actions: ×1.5 lower for Breakout, and ×2 for Frostbite. We also measure the off-policy distribution shift between πo and the TS policy (that utilizes πo) by counting disagreements on actions between the two policies. In Breakout, πo and π1 agreed only in 18% of the states; in Frostbite, the agreement is only 1.96%. Such a level of disagreement between a well-trained agent and its one-step look-ahead extension is surprising. On the other hand, it accounts for the drastic drop in performance when applying TS, especially in Frostbite (Fig. 5).

heading: 3.2 Analysis of the degradation
text: We analyze the decision process of a policy πd, given a pre-trained value function estimator, in a probabilistic setting. Our analysis leads to a simple correction term given at the end of the section. We also show how to compute this correction from the TS. Formally, we are given a policy represented as a value function Q̂πo , which we feed to the d-step greedy policy (2) for each action selection. Policy training is a stochastic process due to random start states, exploration, replay buffer sampling, etc. The value estimator Q̂πo can thus be regarded as a random variable, with an expectation that is a deterministic function Qπo with a corresponding policy πo, where ‘o’ stands for ‘original’.
In short, we bound the probability that a sub-optimal action falsely appears more attractive than the optimal one. We thus wish to conclude with high probability whether a0 = arg maxa Q̂ πo d (s, a) is indeed optimal, i.e., Qπod (s, a0) ≥ Qπod (s, a) ∀a ∈ A. As we have shown in Sec. 3.1 that states belonging to trajectories that follow πo have lower value estimation noise, we model this effect via the following two assumptions. Here, we denote by t = 0 the time when an agent acts and not as the time step of the episode.
Assumption 1. Let σo, σe ∈ R+ s.t. 0 < σo < σe. For action sequence (a0, . . . , ad−1) and corresponding state trajectory (s0, . . . , sd),
Q̂πo(sd, ad) ∼ {N (Qπo(sd, ad), σ2o) if a0 = πo(s0) N (Qπo(sd, ad), σ2e) otherwise.
Assumption 1 presumes that a different choice of the first action a0 yields a different last state sd. This is commonly the case for environments with large state spaces, especially when stacking observations as done in Atari. While assuming a normal distribution is simplistic, it still captures the essence of the search process. Regarding the expectation, recall that πo is originally obtained from the previous training stage via gradient descent with a symmetric loss function L. Due to the symmetric loss, the estimate Q̂πoθ is unbiased, i.e., E[Q̂ πo θ ] = Q
πo . Regarding the variance, towards convergence, the loss is computed on replay buffer samples generated according to the stationary distribution of πo. The estimate of the value function for states outside the stationary distribution is consequently characterized by a higher variance, i.e. σo < σe. We also show that this separation of variance occurs in the data, as detailed in the last paragraph of Sec. 3.1.
After conditioning the variance on whether πo was followed, we similarly split the respective subtrees. To split, we assume the cumulative reward along the tree and values at the leaves depend only on (i) the root state and (ii) whether πo was selected at that state. Assumption 2. For action sequence (a0, . . . ad−1) and corresponding trajectory (s0, . . . , sd),
d−1∑ t=0 γtr(st, at) = { Ro(s0) if a0 = πo(s0) Re(s0) otherwise, Qπo(sd, ad) = { µo(s0) if a0 = πo(s0) µe(s0) otherwise,
with Ro, Re, µo, µe being functions of s.
Assumption 2 could be replaced with a more detailed consideration of each trajectory, but we make it for simplicity. This assumption considers the worst-case scenario: the rewards are unhelpful in determining the optimal policy, and all leaves are equally likely to mislead the d-step greedy policy. That is, there is no additional signal to separate between the Ad leaves besides the initial action.
Assuming from now on that Assumptions 1 and 2 hold, we can now explicitly express the distribution of the maximal value among the leaves using Generalized Extreme Value (GEV) theory [9].
Lemma 3.1. The estimates Q̂πod (s, πo(s)) and maxa6=πo(s) Q̂ πo d (s, a) are GEV-distributed with parameters given in Appendix A.1.
All proofs are deferred to Appendix A. Using the GEV distribution, we can now quantify the bias stemming from the maximization in each of two sub-trees corresponding πo vs. all other actions. Lemma 3.2. It holds that
E [ Q̂πod (s, πo(s)) ] = Qπod (s, πo(s)) + γ dBo(σo, A, d)
E [
max a 6=πo(s)
Q̂πod (s, a) ] = max a6=πo(s) Qπod (s, a) + γ dBe(σe, A, d),
where the biases Bo, Be are given in Appendix A.2, and satisfy 0 ≤ Bo(σo, A, d) < Be(σe, A, d).
Lemma 3.2 conveys the main message of our analysis: the variance of terms being maximized translates to a positive shift in the expectation of the maximum. Hence, even if µo(s) > µe(s) for a certain s, a different action than πo(s) can be chosen with non-negligible probability as the bias in µe(s) is greater than the one in µo(s). To compensate for this bias that gives an unfair advantage to the noisier nodes of the tree, we introduce a penalty term that precisely cancels it.

heading: 3.3 BCTS: Bellman Corrected Tree Search
text: Instead of selecting actions via (2), in BCTS we replace Q̂πod with the corrected Q BCTS d defined by
Q̂BCTS,πod (s, a) := { Q̂πod (s, a) if a0 = πo(s0), Q̂πod (s, a)− γd (Be(σe, A, d)−Bo(σo, A, d)) otherwise,
(3)
and we denote the BCTS policy by
πBCTSd := arg max a Q̂BCTS,πod (s, a). (4)
In the following result, we prove that BCTS indeed eliminates undesirable bias.
Theorem 3.3. The relation E [ Q̂BCTS,πod (s, πo(s)) ] > E [ maxa6=πo(s) Q̂ BCTS,πo d (s, a) ] holds if and
only if Qπod (s, πo(s)) > maxa6=πo(s)Q πo d (s, a).
The biases Bo, Be include the inverse of the cumulative standard normal distribution Φ−1 (Appendix A.2). We now approximate them with simple closed-form expressions that are highly accurate for d ≥ 2 (Appendix A.7). These approximations help revealing how the problem parameters dictate prominent quantities such as the correction term in (3) and the probability in Thm. 3.5 below.
Lemma 3.4. When Ad−1 1, the correction term in (3) can be approximated with
Be(σe, A, d)−Bo(σo, A, d) ≈ √ 2 logA ( σe √ d− σo √ d− 1 ) − (σe − σo)/2. (5)
The bias gap in (5) depends on the ratio between σe and σo; this suggests that TS in different environments will be affected differently. As σe > σo, the bias gap is positive. This is indeed expected, since the maximum over the sub-trees of a0 6= πo(s) includes noisier elements than those in the sub-tree of a0 = πo(s). Also, in (5), σe √ d dominates σo √ d− 1, making the bias gap grow
asymptotically with √ d logA. This rate reflects how the number of elements being maximized over affects the bias of their maximum.
Next, we bound the probability of choosing a sub-optimal action when using BCTS. In Appendix A, Thm. A.1, we give an exact bound to that probability without assuming Ad−1 1. Here, we apply Lemma. 3.4 to give the result in terms of σo, σe. Theorem 3.5. When Ad−1 1, the policy πBCTSd (s) (see (4)) chooses a sub-optimal action with probability bounded by: Pr ( πBCTSd (s) /∈ arg max
a Qπod (s, a)
) ≤ ( 1 + 6d logA ( Qπod (s, πo(s))−maxa6=πo(s)Qπod (s, a) )2 γ2dπ2 (σ2o + σ 2 e) )−1 .
The fraction in the above bound is a signal-to-noise ratio: The expected value difference in the numerator represents the signal, while the variances in the denominator represent the noise. In addition, the signal is “amplified” by d logA because, after applying the correction from (3), a larger number of tree nodes amount to a more accurate maximum estimator. A similar amplification also occurs due to nodes being deeper and is captured by γd in the denominator.
The factor γd also appears in the correction term from (3). This correction is the bias gap from (5), scaled by γd. Hence, asymptotically, while the bias gap grows with √ d, the exponential term is more dominant; so, overall, this product decays with d. This decay reduces the difference between the vanilla and BCTS policies for deeper trees by lowering the portion of the estimated value compared to the exact reward. As we show later, this phenomenon is consistent with our experimental observations.
Computing the correction term requires estimates of σo and σe. As a surrogate for the error, we use the Bellman error. To justify it, let us treat the values of Q̂πod at different depths as samples of Q̂
πo . Then, the following result holds for its variance estimator. Note that we do not need Assumptions 1 and 2 to prove the following result.
Algorithm 1 Batch-BFS Input: GPU environment G, value network Qθ, depth d Init tensors: state S̄ = [s0], action Ā = [0, 1, 2, .., A− 1], reward R̄ = [0] for id = 0 to d− 1 do S̄ ← S̄ ×A, R̄← R̄×A // Replicate state and reward tensors A times r̄, S̄′ = G([S̄, Ā]) // Feed [S̄, Ā] to simulator and advance R̄← R̄+ γid r̄, S̄ ← S̄′ // Accumulate discounted reward Ā← Ā×A // Replicate action tensor A times
end for R̄← R̄+ γd maxaQθ(S̄, a) // Accumulate discounted value of states at depth d Return b(arg max R̄)/Ad−1c // Return optimal action at the root
Proposition 3.6. Let v̂arn[X] be the variance estimator based on n samples of X. Then, v̂arn=2[Q̂πo(s, a)] = ( Q̂πo1 (s, a)− Q̂πo0 (s, a) )2 /2 = δ2(s, a)/2,
where δ(s, a) is the Bellman error.
Note that during a TS, at depth 1 we have access to δ(s0, a) of all a ∈ A without additional computation. For depths 2 and above, the Bellman error is defined only for actions chosen by πo, corresponding to a single trajectory down the tree. For these reasons, we base the above result on samples from depths 0 and 1.
Thanks to Prop. 3.6, we can estimate the bias correction term in Lemma 3.4 directly from the TS operation. Specifically, we substitute δ̂o/ √ 2 instead of σo and the same for σe, where δ̂o is the
Bellman error corresponding to a = πo(s) at the root, and δ̂e is the average Bellman error of all other actions. Hence, the correction term is
B(δ̂e, δ̂o, A, d) = √ logA ( δ̂e √ d− δ̂o √ d− 1 ) − (δ̂e − δ̂o)/ √ 8. (6)
A visualization of the resulting BCTS algorithm is in Fig. 2.

heading: 4 Solving Scalability via Batch-BFS
text: The second drawback of TS is scalability: exhaustive TS is impractical for non-trivial tree depths because of the exponential growth of the tree dimension. As TS requires generating |A|d leaves at depth d, it has been rarely considered as a viable solution. To mitigate this issue, we propose Batch-BFS, an efficient, parallel, TS scheme based on BFS, built upon the ability to advance multiple environments simultaneously; see Alg. 1. It achieves a significant runtime speed-up when a forward model is implemented on a GPU. Such GPU-based environments are becoming common nowadays because of their advantages (including parallelization and higher throughput) over their CPU counterparts. E.g., Isaac-Gym [37] provides a GPU implementation robotic manipulation tasks [46], whereas Atari-CuLE [11] is a CUDA-based version of the AtariGym benchmarks [4]. Batch-BFS is not limited to exact simulators and can be applied to learned deep forward models, like those in [36, 23]. Since Batch-BFS simultaneously advances the entire tree, it enables exhaustive tree expansion to previously infeasible depths. It also allows access to the cumulative reward and estimated value over all the tree nodes. Such access to all future values paves a path to new types of algorithms for, e.g., risk reduction and early tree-pruning. Efficient pruning can be done by maintaining an index array of unpruned states which are updated with each pruning step. These indices are then used for tracing the optimal action at the root. We leave such directions to future work. In addition to Alg. 1, we also provide a visualization of Batch-BFS in Fig. 3.

heading: 4.1 Runtime experiments
text: To showcase the efficacy of Batch-BFS, we compare it to a CPU-based BFS and to non-parallel TS, i.e., Depth-First-Search (DFS). We measure the duration of an average single TS operation on two environments: Atari-CuLE [11] and a Deep NN (DNN) mimicking a learned forward model. The
t+1 ). The next-state tensor is then duplicated and
concatenated again with all possible actions, fed into the forward model, etc. This procedure is performed until the final depth is reached, in which case the Q-function is applied per state.
DNN is implemented in pytorch and uses cudnn [8]. It consists of three randomly-initialized hidden layers of width 100 with input size 100 for the state and 2 or 10 for the actions. The results are given in Fig. 4. We run our experiments on a 8 core Intel(R) Xeon(R) CPU E5-2698 v4 @ 2.20GHz equipped with one NVIDIA Tesla V100 16GB. Although we used a single GPU for our experiments, we expect a larger parallelization (and thus computational efficiency) to be potentially achieved in the future through a multi-GPU implementation. As expected, DFS scales exponentially in depth and is slower than BFS. When comparing BFS on CPU vs. GPU, we see that CPU is more efficient in low depths. This is indeed expected, as performance of GPUs without parallelization is inferior to that of CPUs. This issue is often addressed by distributing the simulators across a massive number of CPU cores [15]. We leverage this phenomenon in Batch-BFS by finding the optimal “cross-over depth” per game and swap the compute device in the middle of the search from CPU to GPU.

heading: 5 Experiments
text: In this section, we report our results on two sets of experiments: the first deals solely with TS for inference, without learning, whereas the second includes the case of TS used in training. In all Atari experiments, we use frame-stacking together with frame-skipping of 4 frames, as conducted in [31].

heading: 5.1 Inference with tree search
text: Using the pre-trained Rainbow agents in [24] (see Sec. 3), we test vanilla TS and show (Fig. 5, red plots) that it leads to a lower score than the baseline πo. The largest drop is for d = 1 as supported by our analysis. The game that suffered the most is Frostbite. This can be explained from it having the largest number of actions (A = 18), which increases the bias in (5). As for BCTS, we found that multiplying its correction term (see (6)) by a constant that we sweep over can improve performance; we applied this method for the experiments here. Recall that BCTS is a TS applied on the pre-trained Rainbow baseline; i.e., the case of d = 0 is Rainbow itself. The results in Fig. 5 show that BCTS significantly improves the scores, monotonically in depth, in all games. It improves the Rainbow baseline already for d = 1, while for d = 4, the score more than doubles. For BeamRider, BCTS with d = 4 achieves roughly ×5 improvement. Notice that without the computationally efficient implementation of Batch-BFS, the results for deeper trees would not have been obtainable in a practical time. We provide timing measurements per game and depth in Appendix B.1. Finally, notice that the advantage provided by BCTS is game-specific. Different games benefit from it by a different amount. In one of the games tested, Asteroids, vanilla TS was as useful as BCTS. Our findings reported in the last paragraph of Sec. 3.1 hint why certain games benefit from BCTS more than others. Nonetheless, a more thorough study on how the dynamics and state distribution in different games affect TS constitutes an interesting topic for future research.

heading: 5.2 Training with tree search
text: To further demonstrate the potential benefits of TS once a computationally efficient implementation (Section 4) is available, we show how it affects training agents from scratch on CuLE-Atari environments. We extend classic DQN [32] with TS using Batch-BFS for each action selection. Notice that training with TS does not suffer from the distribution shift studied in Sec. 3. Hence, the experiments below use vanilla TS and not BCTS.
Our experiment is even more significant considering that Efroni et al. [14] recently proved that the Bellman equation should be modified so that contraction is guaranteed for tree-based policies only when the value at the leaves is backed. However, this theory was not supported by empirical evidence beyond a toy maze. As far as we know, our work is the first to adopt this modified Bellman equation to obtain favorable results in state-of-the-art domains, thanks to the computationally efficient TS implementation achieved by Batch-BFS. We find this method to be beneficial in several of the games we tested. In the experiments below, we treat the Bellman modification from [14] as a hyper-parameter and include ablation studies of it in Appendix B.3.
We show the training scores in Table 1 and convergence plots in Appendix 9. For a fair comparison of different TS depths, we stop every run after two days on the same hardware (see Appendix D), not considering the total iteration count. To compare our results against classic DQN and Rainbow, we measure the number of iterations completed in two days by DQN with TS, d = 0. In Table 1 we report the corresponding intermediate results for DQN and Rainbow reported by the original authors in [40]. In most games, it amounts to roughly 30 million iterations. Note that DQN and Rainbow do not utilize the model of the environment while DQN with TS does; the former are brought as a reference for assessing the potential of using a TS with identical computation budget. As already shown in the case of inference with no training, the achieved score increases monotonically with
the tree depth. In four of the five games, DQN with TS even surpasses the more advanced Rainbow. Since all results were obtained for identical runtime, improvement per unit of time is higher for higher depths. This essentially translates to better efficiency of compute resources. Convergence plots as a function of wall-clock time are shown in Appendix C. We also tested TS on two additional games not included in Table 1: Boxing and Pong. Interestingly, TS with d = 4 immediately obtained the highest possible scores in both these games already in the first training iteration.

heading: 6 Related work
text: The idea of searching forward in time has been employed extensively in control theory via methods such as A* [20], RRT [28], and MPC [2]. The latter is quite popular in the context of efficient planning in RL [36, 47, 48]. MPC-based controllers rely on recourse planning by solving an optimization program given the continuous structure of the dynamics. In our setup, the controls are discrete and the forward model is a black-box that cannot be directly used in an optimization scheme.
We leverage an efficient GPU simulator to conduct the look-ahead search. When such a simulator is out of reach, it can be learned from data. There is vast literature regarding learned deep forward models without optimizing a policy [26, 38, 27]. Other works [25, 23, 19, 41] used the learned model for planning but not with a TS. Few considered roll-outs of learned dynamics [6, 17] but only for evaluation purposes. Additional relevant works are MuZero [42] and “MuZero Unplugged” [43] which utilized a learned forward-model for prediction in an MCTS-based policy. In [34], the trade-off between learning and planning using a TS was empirically tested. Look-ahead policies for RL were also studied theoretically; bounds on the suboptimality of the learned policy were given in [12, 13, 14]. There, the focus was on the effect of planning on the learning process.
Finally, our distribution-shift analysis and approach draw connections to similar challenges in the off-policy [21, 18] and offline [29] RL literature. These address the issue of overestimation in RL due to the ‘max’ operation in the Bellman equation in various ways. In our work, we show that with TS the overestimation problem is exacerbated, and can lead to counter-intuitive performance degradation. Our analysis is different from previous papers, but intuitively the core solution is similar: “be careful” of states and actions not seen during training, and penalize them accordingly.

heading: 7 Discussion
text: Our study of the degradation with vanilla TS implies that the learned value function is not representative of the actual states that can be visited. This conclusion can be helpful in debugging RL systems. It can also be used to improve robustness, tune the approximation architecture, or guide exploration.
Our solution to the above performance degradation is an off-policy correction that penalizes high-error trajectories. It can be further improved with different notions of uncertainty for the value function, e.g., Bayesian or bootstrapped models. Also, while some simulators are available on GPU, such as Atari [11] and robotic manipulation [30], in other cases, a learned model can be used. In these cases, one could include the simulator quality in the off-policy penalty term. Finally, a limitation of our work is that we focus on problems with a discrete action space. Handling problems with continuous action tasks is a challenging direction for future work.
Broader Impact. The paper proposes a method to improve existing RL algorithms. As such, its main impact is to make RL more easily and widely deployable. Since our method can be applied to policies trained with any algorithm, it can be viewed as a generic “policy booster”, and find applications with access to the environment model or its approximation.

Review summaries:
This work investigates the problem of improving pretrained agents on discrete RL problems by using tree search. Using a ground-truth environment model, it finds that naively implemented tree search sometimes results in worse performance than no tree search at all. The authors provide analysis of this phenomenon in a simplified setting which shows that performing search on noisy Q estimates, especially with greater noise for states outside the span of the policy, results in taking incorrect actions. They introduce a correction term for tree search and show that it provides superior performance. The paper also includes a discussion of using batching to make tree search faster on accelerators, including using it to train agents.

In the setting where a pretrained value function is used in combination with MCTS, the authors show that because of the distribution shift problem, using the pretrained value function as the evaluation function in MCTS results in a performance worse than directly using the value function. Then, a correction term that in the end is based on the bellman error is proposed to resolve this issue. Experiments show that this correction term is effective.   The second part of this paper proposes a more GPU compatible search method that seems to scale well when a GPU simulator is available. 

This paper is about combinations of Tree Search (TS) and trained value functions, where the value functions may already have been trained without involving any TS during training, in RL problems such as Atari games. Claimed contributions:  1) Identification and analysis of a distribution-shift phenomenon, where, if a value function was trained standalone (without TS), adding TS on top of it at evaluation time can actually reduce an agent's performance (relative to, say, a plain agent that is greedy w.r.t. the trained value function). The basic intuition for why this happens is as follows: during training, states that are more likely to be visited under the greedy policy (w.r.t. value function) show up more frequently in collected experience, and hence also get a more accurately-trained value function. A tree search (especially an exhaustive one) will run the value function significantly more often to also evaluate states that the greedy policy would rarely/never reach, and in these states the value function will have larger errors (which may be overestimating or underestimating errors). Because there are usually more paths through the tree that the greedy policy would not visit than there are paths that the greedy policy would visit, the number of branches with large errors are likely to outnumber the branches with smaller errors. This outnumbering leads to a maximisation bias, where, when there are multiple such branches with large errors, likely at least one of them will be overestimating, causing the tree search to prioritise that branch over one with a more accurately-trained value function.  2) A correction term that allows a new "BCTS" tree search algorithm to function well with a pre-trained value function.  3) An approach named Batch-BFS to efficiently, on a GPU, run a breadth-first-style tree search with trained value function (assumes, I think, that the forward model of the environment, like an Atari game, itself can also be implemented to run on the GPU).  The resulting Batch-BFS approach outperforms DQN and Rainbow on several Atari games.

The authors introduce a new tree search technique that can be used to improve the score of trained RL agents online. They propose an analysis of the counter-intuitive phenomenon that applying standard MCTS with a pre-trained value function is less performant than using directly the value function in spite having access to more information.   The first contribution is an analysis of TS and the introduction of an off-policy correction term based on the analysis. The second contribution is a GPU adaptation of TS making it fast at inference time but also possible to train tree based agents in a scalable way.  The paper is well written and addresses an interesting problem and the proposed method could be used in many RL algorithms and applications. The GPU implementation seems somewhat disconnected to the rest of the paper.