Improve Agents without Retraining: Parallel Tree Search with Off-Policy Correction
Authors: Gal Dalal
Assaf Hallak
Steven Dalton
iuri frosio
Shie Mannor
Gal Chechik
Venue: NeurIPS 2021 Poster (2021)
Keywords: Reinforcement Learning
Tree Search
Off-policy
BFS

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. We will share the code upon publication.

Improve Agents without Retraining: Parallel Tree Search with Off-Policy Correction
Assaf Hallak
Gal Dalal *
NVIDIA Research
NVIDIA Research
ahallak@nvidia.com
italal@nvidian com
Steven Dalton NVIDIA Research
Iuri Frosio
Shie Mannor
NVIDIA Research
NVIDIA Research
sdalton@nvidia.com com
ifrosioenvidia.com
smamnor@nvidia.coo com
Gal Chechik
NVIDIA Research gchechik@nvidia... com
Abstract
Tree Search (TS) is crucial to some of the most influential successes reinforce- ment learning. Here, tackle two major challenges with TS that limit its usability: distribution shift and scalability. We first discover and analyze a counter-intuitiva 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 the exact state and reward in future steps. We show this is due 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- off-policy correction term that accounts for the mismatch between the pre-trained value and its corresponding TS policy by penalizing under-samp 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-trainne 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: GPU breadth-fire search that advances all nodes in each depth of the tree simultancously. Batch-BER BFS reduces runtime by two orders of magnitude and, beyond inference, enables also training with TS 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.
Introduction
Tree search (TS) is fundamental component of Reinforcement Learning (RL) 466 used in some
of the most successful RL systems 42 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 (NcurIPS 2021).

---

multiple iterations as the tree is built node-by by-node, but less suitable for continuous state-spaces or image-based domains like robotics and autonomous driving] [ 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, conduct the TS "on-demand" by expanding the tree up to given depth 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- pre-trained agents even without having the ability or resources to train them. For example, single AlphaGo training run is estimated to cost 35 million USD []] In other cases, even when compute budget is not 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 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.]] 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 forward model that is fed with actions to advance states and produce rewards. Here, we build on the recently published CuLE [ 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 G
Performing TS on-demand has many benefits, but it also faces limitations. We identify and analyze [wo major obstacles: distribution shift and scalability.
First, we report a counter-intuitivy 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 tree- based policy. We analyze this phenomenon by quantifying the probability of choosing sub-optimal action when the value function error is high. This occurs because for values of out-o---iittibuttio states, larger variance translates larger bias of the maximum. Our analysis leads to simple, computationally effective off-policy correction term based on the Bellman error. We refer to the resulting TS 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-traine are-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 this limitation, we propose Batch-BFS: a parallel GPU adaptation of Breadth-Firss Search (BFS), which brings the runtime down to a practical regime. We measured dideer-off-maniituee 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 tree. By combining Batch-BFS with DQN (3) 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.
steps. To overcome
Our Contributions. (1) We identify and analyze a distribution shift that impairs post-training TS. (2) We introduce correction mechanism and use it to devise BCTS: efficient algorithm that improves pre- vre-trained agents, often doubling their scores or more. (3) We create Batch-BFS, an efficient TS on GPU. (4) We use Batch-BFS train tree- rese DQN agents and obtain higher scores than DQN and Rainbow.
'In the case of continuous state spacesspces 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.
2

---

2
Preliminaries
An MDP defined as the 5-tuple AS, P,r,)), where S is state space, A is a finite action space, P(s'|s, ss, a) transition kernel, r(s, a) is reward function, Y E (0, 1) discount factor. At each step t= 0, I,..., the agent observes the last state 8t, performs action d, and receives reward rr. The next state then sampled by 8141 ~ P(-[st, 0z). For brevity, we denote A := |A].
Our framework is an Lnfinitehhorizon discounted Markov Decision Process (MDP) 39
be the state-action value of a policy т, defined in state as (s, = E" Σ -x= ('( ग (st)) so = 5, ag a) where E denotes expectation w.r.t. the distribution induced by т. Our goal is to find policy xt" yielding the optimal value Q* such that (s, a) max, r(s, a) + /
Let S + A be a stationary policy. Let Q S x A
max ("' a'). ) It is well known that
(=( = r(8, a) P(sa)()) max Q (s',a'),
T TT arg max (* (8.a).
a'
Vanilla tree search. To ease notations and make the results concise, we limit the analysis deter- ministic transitions i.e., an action sequence (an,...,,0 ad-1), starting at so leads to corresponding Nonetheless, the results can be extended to a stochastic setup by working with the marginal probability over the trajectory. Then, for policy To. let the d-step Q-function
trajectory (80:--. ...: sg).
(s,
max (nN) E.A.
Σ r'r(s,, az) + +-A (sa; oa)
80 8,0g ma
and similarly let Q7 (s, a) be the d-step Q-function estimator that uses an estimated Q-function Q instead of QTo Finally, denote by Fu the d-step greedy policy
a ( arg max Q7 (s) aEA
3 Solving the Tree Search Distribution Shift
In this section, we show how to leverage TS to improve 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 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 eтror and problem parameters, we find the exact debiasing that yields the optimal signal-tor hignal-to-noise separation.
3.1 Performance degradation with vanilla tree search
We focus applying TS at inference time, without learning. We begin with 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 n, for finite depths as well. To test this, we load Rainbow agents QT=,, pre-trained on 50M frames; they are publicly available in (24] and achieve superhuman scores in (22). We use them to test TS policies Ra with multiple depths d on several Atari benchmarks. Surprisingly, the results (red curves in Fig. 5 Sec. (.T) 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 by our analysis in Thm. 3.5
a fact later explained
We find the reason for this performance drop to the poor generalization of the value function to states outside the stationary w, 's distribution. Fig. T shows a typical, bad action selection in Atari Breakout by 1epth-1 TS. The table on the right reports the estimated Q-values of the root state (first column) and of every state at depth I (last four columns). Since the ball is dropping, 'Left' the optimal action. Indeed, this corresponds with the Q-values at the root. However, at depth 1,
"The Atari environments we experiment on here are indeed close to deterministic. Their source of randomness is the usage of random number of initial noop actions [33

---

883
a
(5o a)
((11,
(() a)
(5 a)
((1,
Left
2.6103
2.7906
2.8114
2.8114
2.8326
No-op
2.5712
2.7874
2.8082
2.8082
2.8185
Fire 2.5801
2.7837
2.8041
2.8041
2.8192
Right 2.3896
2.7748
2.7914
2.7914
2.8022
Figure 1: A failure of vanilla tree search. Left: An Atari Breakout frame. Right: Q-values of TS for the frame on the left. Rows correspond to the action taken at the considered depth, which is d for the first column and d 1 for the four others. The action at the root of the tree is color-coded: Red for 'Right,, and blue for 'Left".
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. [ 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 interpreted as an off-policy distribution shift [35]
Beyond the anecdotal example in Fig. additional evidence supports 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 To to all other actions at the tree root. When averaging over 200 episodes, we find that the error for actions chosen by To is consistently lower than for other actions: x x1.5 lower for Breakout, and 2 for Frostbite. We also measure the off-policy distribution shift between To and the TS policy (that utilizes T.) by counting disagreements 18% of the states; in Frostbite, the agreement is only 1.96%. Such a level of disagreement between 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).
actions between the policies. In Breakout, n. and л agreed only in
3.2 Analysis of the degradation
We analyze the decision process of policy Td. given a pre-trained value function estimator, in 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 given policy represented as value function @*o which feed the d-step greedy policy (2) for each action selection. Policy training a stochastic process due to random start states, exploration, replay buffer sampling, etc. The value estimator QTa can thus be regarded as random variable, with an expectation that is deterministic function QT. with a corresponding policy To. where "o" stands for 'original'.
In short, we bound the probability that sub-optimal action falsely appears more attractive than the optimal one. We thus wish to conclude with high probability whether ap arg max, Q7 (8, (4,4) is indeed optimal, i.e., 27e (s, ao) 2 Q (s,a) a) Va E A.
As we have shown in Sec. Sec.3.1 that states belonging to trajectories that follow n, 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.
Let oo. de E IR+ S.f. O < To < Ge- For action sequence (ao....,.a- ad-1) and corresponding state trajectory (So)....8a).
Assumption 1.
(N(Q ad), 0 (N(QT az), 02)) otherwise.
if (Su)
(sd, ad)

---

[ssumption presumes that different choice of the first action yields different last state sd- This commonly the case for environments with large state spaces, especially when stacking observations as done in Atari. While assuming normal distribution is simplistic, it still captures the essence of the search process. Regarding the expectation, recall that wa is originally obtained from the previous training stage via gradient descent with symmetric loss function L. Due to the symmetric loss, the estimate Q7 is unbiased, i.e., B[ QTe Regarding the variance, towards convergence, the loss is computed on replay buffer samples generated according to the stationary distribution of To. The estimate of the value function for states outside the stationary distribution is consequently characterized by a higher variance, i.c. do < Gg. We also show that this separation of variance occurs in the data, detailed in the last paragraph of Sec. 3.1
After conditioning the variance on whether To was followed, we similarly split the respective sub- trees. 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 To was selected at that state.
Assumption 2. For action sequence (0o,.... (00.... ad--1 and corresponding trajectory (50, ................................ sa),
d-1
Er(((0,,0)) ue) = ( (80) otherwise,
(Ro(so) a0 (s() (Ta ua)
(Ho(so) ag = mo(s)) (M((8)) otherwise,
t=0
with Ro: Re, Pos Me being functions of
Assumption could be replaced with a more detailed consideration of cach trajectory, but we make 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 separate between the A leaves besides the initial action.
Assuming from now on that Assumptions I of the maximal value among the leaves using Generalized Extreme Value (GEV) theory 9
[hoid, we can now explicitly express the distribution
Lemma 3.1. The estimates Q7 Ts())) and maxa/x,(s) Q70 (s, a) are GEV-distributed with parameters given in Appendix A.1
All proof's 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 no vs. all other actions.
Lemma 3.2. It holds that
E
Q (T)))
Q7 (s) Bo(o,,
EE
max
Q (s, a)
max astttaa)
+ B.
where the biases Bo, Be are given in Appendix A.2 and satisfy 0 5 (o A,d) (.(oe,A.d)
Lemma 3.2 conveys the main message of our analysis: the variance of terms being maximized translates positive shift the expectation the maximum. Hence, even if (s.(s) > (e(s) for certain s, a different action than (s) can be chosen with non-negligible probability as the bias in
(e(s) greater than the one in (s(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.
3.3 BCTS: Bellman Corrected Tree Search
Instead of selecting actions via 2 BCTS we replace Q with the corrected QCTT defined by
Q (s,a) (8, d (Balge, A, - Bolao, A,)))
if ag = (s():
BCTS (s,a)
(3)
otherwise,
and we denote the BCTS policy by
π
BCTS
arg max Q BCTS, (s,a). a).
(4)
In the following result, prove that BCTS indeed climinates undesirable bias.
5

---

Sp
frequently observed during training
si
rarely observed during training
Figure 2: BCTS Algorithm. Exploring the full tree reveals -itto--dittribution states (red). These less visited during training and tend to have highly variable scores, leading to high overestimation error. The penalty term
01
$%
11
21
2
B 86, A, d) (see (6)) cancels the excess bias. The Bellman errors de, óg are extracted from the tree at depth 1.
The max is over-estimated 3 Penalize by y B6, 8 d.d) (Eq. (Eq.(6)); (6))
Theorem 3.3. The relation T.
La
QUCTS, (s, >
ma(s))
EE max QUCTS (s, (s. a) holds if and
only if Q7: (8, (o(s)) (1) Q7 as)
The biases B, Bg include the inverse of the cumulative standard normal distribution 4-1 (Ap- pendix [A.2). We now approximate them with simple closed-form expressions that are highly accurate for d 2 2 (Appendix (..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, - Bo(o., A. 21og A o.Va (aeVd-ooVd vovd- -0)--
The bias gap in 5 depends on the ratio between de and do: this suggests that TS in different environments will be affected differently. expected, since the maximum over the sub-trees of ag + (s(s) includes noisier elements than those in the sub-tree of aq (s). Also, in (5) ae vd dominates voyd- making the bias gap grow asymptotically with Volog A. This rate reflects how the number elements being maximized over affects the bias of their maximum.
a. > oo the bias gap positive. This indeed
Next, we bound the probability choosing sub-optimal action when using BCTS. In Appendix A Thm. A.I give exact bound to that probability without assuming Ad >> 1. Here, we apply Lemma. 3.4 to give the result in terms oo: de.
the policy x BCCSS(s) (see chooses a sub-optimal action with Theorem 3.5. When A >
probability bounded by:
s(1+ 6d A (QT xo(s)) max 27 (s, nd (o3 +
Pr (s)) & arg max QJ (s,a))
The fraction in the above bound is a signal-to-no 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 d'log A because, after applying the correction from (3, a larger accurate maximum estimator. A similar amplification also
number of tree nodes amount occurs due to nodes being deeper and is captured by 7 in the denominator.
The factor nd 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 vd, 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 a, and de- As surrogate for the error, we use the Bellman error. To justify it, let us treat the values of Q at different depths as samples of Q*. Then, the following result holds for variance estimator. Note that we do not need Assumptions and 2 to prove the following result.

---

Algorithm Batch-BFS
Input: GPU environment G, value network Qo, depth d
Init tensors: state 5 [So], action A = [0, 1, 2, A - 1], reward R [0]
for dd 0 to d do
S 5 x A, R t Rx A
// Replicate state and reward tensors A times
F.S' = G(S Al)
// Feed [S, A] simulator and advance // Accumulate discounted reward
- R + 7* S 8¹
A Ax A
// Replicate action tensor A times
end for
R. R R+7" max, Qa(S.
// Accumulate discounted value of states at depth d
Return (arg max R
// Return optimal action at the root
Proposition 3.6. Let [X) be the variance estimator based samples of X. Then,
varn=2l = (Qf ((f Qu (s.aa) /2 /2= (2,
where 8(s, a) the Bellman error.
Note that during a TS, at depth we have access to 8(50, a) of all a E A without additional computation. For depths and above, the Bellman error is defined only for actions chosen by To corresponding to single trajectory down the tree. For these reasons, we base the above result on samples from depths 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 8./ /2 instead of a, and the same for de, where å, is the Bellman error corresponding to (s) at the root, and & is the average Bellman error of all other actions. Hence, the correction term
B(8,, 8, A, d) Vlog A
(8. - -8)1/BB.
visualization of the resulting BCTS algorithm is in Fig.
4 Solving Scalability via Batch-BFS
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| leaves at depth d, it has been rarely considered as viable solution. To mitigate this issue, propose Batch-BFS, an efficient, parallel, TS scheme based on BFS, built upon the ability to advance multiple environments simultaneously; see Alg. I It achieves significant runtime speed-up when forward model is implemented on GPU. Such GPU-based environments are becoming common nowadays because of their advantages (including parallelization and higher throughput) CPU counterparts. E.g., Isaac-Gym (37] provides a GPU implementation robotic manipulation tasks [46), whereas Atari-CuLE [1] is 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 36 2X] 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 all future values paves path to new types of algorithms for, e.g., risk reduction and early tree- pre-prunnng. 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 visualization of Batch-BFS in Fig.
their
4.1 Runtime experiments
To showcase the efficacy of Batch-BFS, we compare it to CPU-based BFS and to non-parallel TS, ie.., Depth-FFrst-Search (DFS). We measure the duration an average single TS operation on two environments: Atari- Muri-CuLE M|| and Deep NN (DNN) mimicking a learned forward model. The
7

---

duplication
8g
B7
5
Bg
S,
5,
B>
batch
L2
simulaiion
state duplication
a,
Qg
S
a.
0n
00
Shu
Og
St
Sh..
botch
GPU simulation
O-function estimation (DNM)
[qP9, df,, d 9
Figure 3: Visualization of Batch-BFS. The tree expansion is illustrated on the left, with the corre- sponding batch GPU operations on the right. every tree expansion, the state S, duplicated and concatenated with all possible actions. The resulting tensor is fed into the GPU forward model to generate the tensor of next states , 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.
The next-state tensor is then duplicated and
Atari-CuLE environment
Deep forward model
Time par sap [ms]]
Time per step [mb]
10a
1.'
And (Breakout)
Aud iPong)
10
A-0 (Prool))
BFSICMU), A-4 (Breakoutt
DFS (A=2)
015[C/0), A-e (Pegg)
RFRICMU), Au4 (Pioall)
10-1
SFS/
(A-10)
SFS(OPU] (A-2)
10'
12
10"
10
(iiball
10
10"
10
SFS( (A-10)
Figure 4: Average tree-search time per action selection. Left: Atari-CuLE Breakout, Pong, and VideoPinball. Right: A randomly generated neural network to mimic a learned forward-model with A E (2, 10). Note that and y axes are in log-scale.
DNN is implemented in pytorch and uses cudnn 8 It consists of three randomly-initializen hidden layers of width 100 with input size 100 for the state and 10 for the actions. The results are given in Fig. 4 We run our experiments on a 8 core Intel(R) Xeon(R) CPU ES-2698 v4 @ 2.20GHz equipped with one NVIDIA Tesla V100 16GB. Although we used a single GPU for our experiments, we expect larger parallelization (and thus computational efficiency) to be potentially achieved in the future through multi-GPU implementation. As expected, DFS scales exponentially in depth and is slower than BFS. When comparing BFS 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 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.
Experiments
In this section, we report our results two sets of experiments: the first deals solely with TS for inference, without learning, whereas the second includes the case of TS used training. Atari experiments, frame-stacking together with frame-skipping of 4 frames, conducted in [3].

---

Fresibie
Beamnaler
Spacelrraders
AsAurolds
12030
60263
16030
tarbbwr
Furrow
50283
Hantone
0000
Ru
C00
5272
Reward
40363
6000
8922
30382
0332
6030
2039
10303
2000
4035
Depth
Depth
Depth
Dosth
Depth
Figure 5: Inference only: Vanilla TS vs. BCTS. Median scores with lower (0.25) and upper (0.75) quantiles over 200 episodes, as a function of the tree depth. Surprisingly, vanilla TS often degrades the performance of the pretrained Rainbow agent. BCTS (blue) improves upon vanilla TS (red) for all depths, except for Asteroids. The improvement grows monotonically with the tree depth.
5.1 Inference with tree search
Using the pre-trained Rainbow agents in [24] (see Sec. 3}. we test vanilla TS and show (Fig. red plots) that it leads to a lower score than the baseline To. 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.c., the case of d = 0 is Rainbow itself. The results in Fig. 5 show that BCTS significantly improves the scores, monotonically baseline already for d = 1, while for d = 4, the score more than doubles. For BeamRider, BCTS with d achieves roughly x5 improvement. Notice that without the computationally efficient implementation of Batch-BFS, the results for deeper trees would not have been obtainable in practical time. We provide timing measurements per game and depth in Appendix [.]] 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. Dec.3.11m [3.T 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.
depth, in all games. It improves the Rainbow
5.2 Training with tree search
To further demonstrate the potential benefits of TS once a computationally efficient implementation (Section is available, we show how it affects training agents from scratch on CuLE-Atari environ- ments. We extend classic DQN 332 with TS using Batch-BFS for each action selection. Notice that training with TS does not suffer from the distribution shift studied in Sec. Hence, the experiments below use vanilla TS and not BCTS.
Our experiment is even more significant considering that Efroni et al. 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 toy maze. As far as we know, our work is the first to adopt this modified Bellman equation to obtain favorable results in -fate-of-the-art domains, thanks to the computationally efficient TS implementation achieved by Batch-BFS. We find this method to be beneficial several of the games we tested. In the experiments below, we treat the Bellman modification from [4] as hyper-parameter and include ablation studies of in Appendix B.3
We show the training scores in Table I and convergence plots in Appendix 9 For fair comparison of different TS depths, we stop every run after two days on the same hardware (see Appendix DD. not considering the total iteration count. To compare our results against classic DQN and Rainbow, we measure the number of iterations completed two days by DQN with TS, report the corresponding intermediate results for DQN and Rainbow reported by the original authors in 44). 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 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
In Table [ we

---

Table 1: Atari scores after two days of training. We follow the evaluation method in [2 Average of 200 testing episodes, from the agent snapshot that obtained the highest score during training.
DQN with TS, depth d d = 2
Game
DQN 32
Rainbow
20
d = 1
d = 3
d = 4
Asteroids
2, 093
2, 613
4, 794 420
17, 929
1, 664
1,594
Breakout
385
581
620
377
327
MsPacman
1, 644
2,923
3,498
4, 021
2,398
3,600
SpaceInvaders
675
1, 6,0
2,132
2, 2,550
1,
2, 162
VideoPinball 229, 129
244, 052
442,
345, 742
163, 720
641, 235
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 higher for higher depths. This essentially translates to better efficiency of compute resources. Convergence plots as function of wall-clock time are shown in Appendix C We also tested TS on two additional games not included in Table I
Boxing and Pong. Interestingly, TS with d = 4 immediately obtained the highest possible scores in both these games already the first training iteration.
6 Related work
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 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.
in RL 36
48
We leverage an efficient GPU simulator to conduct the look-ahead search. When such a simulator out of reach, it can be learned from data. There is vast literature regarding learned deep forward 23
models without optimizing a policy 26 38 27 Other works 225
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"
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 There, the focus was on the effect of planning on the learning process.
43 which utilized a learned neward-mmodd for prediction an MCTS-
12
13
14]
Finally, our distribution-shie analysis and approach draw connections to similar challenges in the off-policy 18] and offline (29] RL. literature. These address the issue of overestimation in RL due the 'max' operation in the Bellman equation in various ways. In work, 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.
Discussion
Our study of the degradation with vanilla TS implies that the learned value function is not representa- tive 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 Atari I 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.
GPU, such as
Broader Impact. The paper proposes a method to improve existing RL algorithms. As such, its main impact is to make RI. more easily and widely deployable. Since our method can be applied to policies trained with any algorithm, can be viewed generie "policy booster", and find applications with access to the environment model or its approximation.
10

---

References
[1] How much did alphago zero cost? cessed: 2021-05-20.
html Ac-
[2] Paul Serban Agachi, Mircea Vasile Cristea, Alexandra Ana Csavdari, and Botond Szilagyi. Model predictive control. De Gruyter, 2016.
[3] JM Blair, CA Edwards, and JH Johnson. Rational chebyshev approximations for the inverse of the error function. Mathematics of Computation, 30(136):827-834 1976.
[4] Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojeiech Zaremba. Openai gym. arXiv preprint urXiv:160.0015500 2016.
[5] Cameron B Browne, Edward Powley, Daniel Whitehouse, Simon M Lucas, Peter Cowling, Philipp Rohlfshagen, Stephen Tavener, Diego Perez, Spyridon Samothrakis, and Simon Colton. A survey of monte carlo tree search methods. IEEE Transactions Computational Intelligence and Al in games, 4(1):1-43, 2012.
[6] Jacob Buckman, Danijar Hafner, George Tucker, Eugene Brevdo, and Honglak Lee. Sample- efficient reinforcement learning with stochastic ensemble value expansion. arXiv preprint arXivv1807.01675. 2018.
[7] Francesco Paolo Cantelli. Sui confini della probabilita. In Arri del Congresso Internazionale dei Matematici: Bologna del 3 al 10 de settembre di 1928. pages 47-60, 1929.
[8] Sharan Chetlur, Cliff Woolley, Philippe Vandermersch, Jonathan Cohen, John Tran, Bryan Catanzaro, and Evan Shelhamer. cudnn: Efficient primitives for deep learning, 2014. cite arxiv:1410.0759.
[9] Stuart Coles, Joanna Bawa, Lesley Trenner, and Pat Dorazio. An introduction to statistical modeling of extreme values, volume 208. Springer, 2001.
[10] Rémi Coulom. Efficient selectivity and backup operators in monte-carlo tree search. In International conference computers and games, pages 72-83. Springer, 2006.
[11] Steven Dalton, Iuri Frosio, and Michael Garland. Accelerating reinforcement learning through gpu atari emulation. arXiv preprint arXiv:1907.08466; 2019, BSD 3-Clause "New" "Revised" License.
[12] Yonathan Efroni, Gal Dalal, Bruno Scherrer, and Shie Mannor. Beyond the one-step greedy approach in reinforcement learning. In International Conference on Machine Learning, pages 1387-1396. PMLR, 2018.
[13] Yonathan Efroni, Gal Dalal, Bruno Scherrer, and Shie Mannor. Multiple-step greedy policies in approximate and online reinforcement learning. In NewrIPS, 2018.
[14] Yonathan Efroni, Gal Dalal, Bruno Scherrer, and Shie Mannor. How to combine tree-search methods reinforcement learning. Proceedings of the AAAI Conference on Artificial Intelli- gence (AAAI 2019), 2019.
[15] Lasse Espeholt, Hubert Soyer, Remi Munos, Karen Simonyan, Vlad Mnih, Tom Ward, Yotam Doron, Vlad Firoiu, Tim Harley, lain Dunning, et al. Impala: Scalable distributed deep-rl with importance weighted actor-learner architectures. In International Conference on Machine Learning, pages 1407-1416. PMLR. 2018.
[16] Jesse Farebrother, Marlos C Machado, and Michael Bowling, Generalization and regularization in dgn. arXiv preprint 2XXv:1810.00123.. 2018.
[17] Vladimir Feinberg, Alvin Wan, Ion Stoica, Michael I Joseph E Gonzalez, and Sergey Levine. Model-based value estimation for efficient model-free reinforcement learning. arXiv preprint arXiv: 0010000000000,.00000000000000000 2018.

---

[18] Scott Fujimoto, Herke Hoof, and David Meger. Addressing function approximation error in actor-critic methods. In International Conference on Machine Learning, pages 1587-1596. PMLR, 2018.
[19] David Ha and Jürgen Schmidhuber. World models. arXiv preprint arXiv: 10122, 2018.
[20] Peter Hart, Nils Nilsson, and Bertram Raphael. formal basis for the heuristic determination of minimum cost paths. IEEE Transactions on Systems Science and Cybernetics, 4(2):100-107, 1968.
[21] Hado Hasselt. Double q-learning. Advances in neural information processing systems, 23:2613- 2621, 2010.
[22] Matteo Hessel, Joseph Modayil, Hado Van Hasselt, Tom Schaul, Georg Ostrovski, Will Dabney, Dan Horgan, Bilal Piot, Mohammad Azar, and David Silver. Rainbow: Combining improve- ments in deep reinforcement learning. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 32, 2018.
[23] Lukasz Kaiser, Mohammad Babacizadeh, Piotr Milos, Blazej Osinski, Roy H Campbell, Konrad Czechowski, Dumitru Erhan, Chelsea Finn, Piotr Kozakowski, Sergey Levine, et al. Model- based reinforcement learning for atari. arXiv preprint arXiv: 1933v:1903.00374, 2019.
[24] Kaixhin. Rainbow Pre-trained Agents v1.3, 2019, MIT License.
[25] Gabriel Kalweit and Joschka Boedecker. Uncertainty-driven imagination for continuous deep reinforcement learning. In Conference Robot Learning, pages 195-206. PMLR, 2017.
[26] Seung Wook Kim, Yuhao Zhou, Jonah Philion, Antonio Torralba, and Sanja Fidler. Learning to simulate dynamic environments with gamegan. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 1231-1240, 2020.
[27] Rahul G Krishnan, Uri Shalit, and David Sontag. Deep kalman filters. arXiv preprint 111.51151..05121, 2015.
[28] Steven M LaVale al. Rapidly-exploring random trees: A new tool for path planning. Technical Report. Computer Science Department, lowa State University, 1998.
(29] Jongmin Lee, Wonseok Jeon, Byung-Jun Lee, Joelle Pineau, and Kee-Eung Kim. Optidice: Offline policy optimization via stationary distribution correction estimation. arXiv preprint arXiv:2106. 10783, 2021.
[30] Jacky Liang, Viktor Makoviychuk, Ankur Handa, Nuttapong Chentanez, Miles Macklin, and Dieter Fox. Gpu-accelerated robotic simulation for distributed reinforcement learning. In Conference on Robot Learning, pages 270-282. PMLR, 2018
[31] Marlos Machado, Mare G Bellemare, Erik Talvitie, Joel Veness, Matthew Hausknecht, and Michael Bowling. Revisiting the arcade learning environment: Evaluation protocols and open problems for general agents. Journal of Artificial Intelligence Research, 61 61:5235562, 2018.
[32] Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Alex Graves, Ioannis Antonoglou, Daan Wierstra, and Martin Riedmiller. Playing atari with deep reinforcement learning. arXiv preprint 11Xiv:1312.5602, 2013.
[33] Volodymyr Mnih, Koray Kavukeuoglu, David Silver, Andrei A Rusu, Joel Veness, Mare G Bellemare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, et al. Human-level control through deep reinforcement learning. nature, 518(7540):529-532 2015.
[34] Thomas M Moerland, Anna Deichler, Simone Baldi, Joost Broekens, and Catholijn Jonker. Think too fast nor too slow: The computational trade-off between planning and reinforcement learning. arXiv preprint arXiv:2005.0744 074404 2020.
[35] Remi Munos, Tom Stepleton, Anna Harutyunyan, and Marc Bellemare. Safe and efficient off- policy reinforcement learning. In D. Lee, M. Sugiyama, U. Luxburg, I. Guyon, and R. Garnett, editors, Advances Neural Information Processing Systems, volume 29, Curran A Associates, Inc., 2016.
12

---

[36] Anusha Nagabandi, Gregory Kahn. Ronald S Fearing, and Sergey Levine. Neural network dynamics for model-based deep reinforcement learning with model-free fine-tuning. 2018 IEEE International Conference on Robotics and Automation (ICRA), pages 7559-7566. IEFE, 2018.
[37] NVIDIA. Isaac Gym preview release, 2021.
[38] Junhyuk Oh, Xiaoxiao Guo, Honglak Lee, Richard Lewis, and Satinder Singh. Action- conditional video prediction using deep networks in atari games. arXiv preprint arXiv: 2015.
[39] Martin L Puterman. Markov decision processes: discrete stochastic dynamic programming. John Wiley & Sons, 2014.
[40] John Quan and Georg Ostrovski. DQN Zoo: Reference implementations of DQN-DQNNAEEE -Qsed agents, 2020, Apache License 2.0.
[41] Sébastien Racanière, Théophane Weber, David P Reichert, Lars Buesing, Arthur Guez, Danilo Rezende, Adria Puigdomenech Badia, Oriol Vinyals, Nicolas Heess, Yujia Li, et al. Imagination- augmented agents for deep reinforcement learning. In Proceedings of the 31st International Conference on Neural Information Processing Systems, pages 5694-5705, 2017.
[42] Julian Schrittwieser, Ioannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Si- mon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, al. Mastering atari, go, chess and shogi by planning with a learned model. Nature, 588(7833):604-609) 2020.
[43] Julian Schrittwieser, Thomas Hubert, Amol Mandhane, Mohammadamin Barekatain, Ioannis Antonoglou, and David Silver. Online and offline reinforcement learning by planning with learned model. arXiv preprint arXiv: 104.06294 2021.
[44] David Silver, Aja Huang, Chris J Maddison, Arthur Guez, Laurent Sifre,, George Van Den Driess- che, Julian Schrittwieser, loannis Antonoglou, Veda Panneershelvam, Marc Lanctot, et al. Mas- tering the game of go with deep neural networks and tree search. nature, 529(7587):484 489, 2016.
[45] David Silver, Thomas Hubert, Julian Schrittwieser, Ioannis Antonoglou, Matthew Lai, Arthur Guez, Marc Lanctot, Laurent Sifre, Dharshan Kumaran, Thore Graepel, al. Mastering chess and shogi by self-play with general reinforcement learning algorithm. arXiv preprint 1721v:171..01815. 2017.
[46] Richard Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT press, 2018.
[47] Grady Williams, Nolan Wagener, Brian Goldfain, Paul Drews, James M Rehg, Byron Boots, and Evangelos A Theodorou. Information theoretic mpe for model-based reinforcement learning. In 2017 IEEE International Conference on Robotics and Automation (ICRA), pages 1714-1721. IEEE, 2017.
[48] Mario Zanon and Sébastien Gros. Safe reinforcement learning using robust mpc. IEEE Transactions Automatic Control, 2020.
13


Reviews:
[{"review_id": "vgjYfVg1Vme", "reviewer": "Reviewer_1GjW", "summary": "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:\n\n1) 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.\n\n2) A correction term that allows a new \"BCTS\" tree search algorithm to function well with a pre-trained value function.\n\n3) 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).\n\nThe resulting Batch-BFS approach outperforms DQN and Rainbow on several Atari games.", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Overall I think the paper is solid in all the core aspects (originality, quality, clarity, and significance). I do have several more detailed comments below, but do not think any of them should be a major issue to address:\n\n1) First two paragraphs of Introduction discuss how the paradigm of MCTS would be less suitable for continuous/large state spaces or image-based domains, and how the more simple \"breadth-first-style\" TS used in this paper would be more suitable without requiring any memorization, but I don't see why this is true. MCTS has no more problems with continous or large state spaces than a BFS would. With an extremely large or continuous state space, MCTS would likely struggle to grow a sufficiently large tree to fully cover that space... but that holds for any other tree search too. I wonder, with the references to image-based domains and memorization, are the authors actually thinking of problems with a large state **representation** (rather than state space), i.e. problems where a single state requires a large amount of memory to store? Again, I'm not sure that I agree that this is necessarily that much of a problem for MCTS. If it takes too much memory to stores states in nodes, it is also possible to implement MCTS by only storing actions along the edges, and use those actions to recreate the state as the algorithm traverses the tree.\n\n2) Section 4 states that Batch-BFS is proposed \"to solve this issue\" [= the issue that the number of nodes to evaluate scales exponentially with search depth d]. I think that \"solve\" may be too strong a word here. I'll happily believe that GPUs allow for significantly more parallelisation than CPUs, and can meaningfully increase the search depth that can feasibly be reached... but I don't see that as really **solving** the issue of exponential growth. There is still exponential growth, and exponential growth will at some point catch up again and even make a GPU struggle. Maybe just something like \"mitigate\" or \"address\" would be more fair. On a related note, it would be interesting to see what kinds of search depths match the numbers of leaves reported on the x-axes of Figure 3.\n\n3) In *\"Revisiting the Arcade Learning Environment: Evaluation Protocols and Open Problems for General Agents\"*, Machado et al. (2018) investigated how, due to the determinisim in ALE, many approaches turn out to be able to perform well in these environments simply by memorising strong action sequences, and propose using *sticky actions* (in training as well as evaluation) to disincentivise such solutions (as opposed to ones that actually learn to make informed decisions based on the state and generalise to new circumstances). Were sticky actions used in any of the trainings or evaluations in this paper? I think this is a particularly interesting question in the context of this paper, because the use of sticky actions during training would force a bit more exploration and force the value function to be applicable to a more diverse set of states than in a (close to) deterministic environment. To what extent would such a training practice mitigate the bias and maybe reduce or eliminate the need for a correction?\n\n---\n\nMinor comments:\n- Quotes on left-hand side of words \"Left\" and \"Right\" on bottom of page 3 should \"point\" the other way (opening instead of closing quotes)\n- Footnote 1 states the Atari environments used are \"close to deterministic\". Is it possible to elaborate on in what way(s) they are *not* deterministic?\n- Line 123: evidence support --> evidence supports\n- Line 154: Technically, I guess it does not necessarily hold exactly that the expectation of Q-hat equals Q when trained using Rainbow? There may be errors due to imperfect function approximation, probably a bias from Prioritized Experience Replay since it's customary to only partially correct for its non-uniform sampling with importance sampling, and if I recall correctly also only partial or no correcting for biases due to multi-step off-policy update steps in Rainbow. Maybe all of this is not too relevant to this paper though.\n- In Figure 2, is it correct to state that the red states were \"not visited/observed during training\"? I get that they are likely visited *less frequently*, but not entirely unvisited, right?\n- Caption Figure 4: \"a deeper tree often degrades the performance of vanilla TS\" --> is it really fair to say \"often\" here? I see this primarily being the case just in Frostbite, plus in most of the other games specifically when going from depth 0 --> 1, but not anymore when increasing depth beyond 1.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Limitations are adequately addressed in the paper, except possibly somewhat of a limitation in the experiments, see point 3) in the main review above.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "2.5"}, "extra_text": {}}, {"review_id": "my9iuUAUvyR", "reviewer": "Reviewer_drZi", "summary": "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. \n\nThe 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.\n\nThe 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. \n", "questions": "", "limitations": "", "rating": 7, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The authors attribute the performance loss of standard tree search to a distribution shift. When using the estimated value function to guide the tree search, it is evaluated in states poorly explored during training, with inaccurate value estimate and hence downgrade the performance of the whole tree search. \n\nEmpirical analysis of the degradation caused by TS:\n\nThey provide empirical proof and a further theoretical analysis. I think the empirical proof lack of details, fortunately I had observed this phenomena in my own experiments and was not to be convinced.\n-\t“We also measure the off-policy distribution shift between \\pi_0 and the TS policy by counting disagreements “ In this above sentence, do the author use pi_0 in the TS as well?\n-\tThis section on empirical proof of the discrepancy between TS and pi_0 lacks a bit of precision. The author should provide more details at least in appendix. Does the computation budget of the tree search affect these results? Did they use any exploration policy (e.g. UCB) in the tree search?\n\nRegarding the theoretical analysis: \n\nThe author assume that the value estimate follows a normal distribution with two different deviations depending on whether the action follows the original policy or not. Although simplistic this assumption seems reasonable enough to provide a first analysis of the problem and is well justified by the authors.\n\nIn assumption 2: shouldn’t the value function be Q_d^{\\pi_0} instead of Q^{\\pi_0}? The justification for assumption 2 is not as clear.\nThroughout the analysis the author separate the cases where the action follows the original policy or other branches of the tree. The dichotomy is used in all the subsequent proofs. It results in two different biases for the “pi-0 following branch” and the other branch. \n\nThe two biases are used to provide a correction term in the action selection procedure of the tree. With another approximation, the authors derive a very simple formula for approximating the correction term as a function of: the size of the action space, the depth of the tree, sigma_e, sigma_o. The two standard deviations are introduced in assumption 1. \n\nThe rest of the analysis shows that the probability of selecting a suboptimal action is bounded.\n\nThe overall analysis strongly rely on assumption 1 but seems sound and correct overall, leading to a practical correction term with two parameters to estimate. The remaining of the section provide guidance in how to estimate sigma_o and sigma_e. The author should provide an algorithm, or an intuitive explanation on how to practically estimate those two parameters, I think proposition 3.6 provides a theoretical justification but it is unclear how to practically implement it.\n\nScalability:\n\nThis part addresses a completely different problem of TS: scalability. TS is usually expensive to run as it requires many online simulations growing exponentially with the depth. The authors rely on GPUs to address the scalability issue. This contribution is not theoretical but practical.\n\nThe author provide a pseudo-code of the solution but no explanation is given. The diagram in the appendix is helpful and should in my opinion replace the provided pseudo-code.\n\nWithout surprise, the parallel implementation outperforms standard TS for a large number of leaves in terms of runtime. The batch BFS implementation consists in evaluating each level of the tree at the same time in a “batch”. It is unclear how this method can be used with TS techniques using pruning or UCB type branching. \n\nInference Experiments:\n\nThe experiments support the claim, BCTS provides consistent improvement over rainbow and standard TS. \n“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.” \nThis constant is not mentioned in the analysis and should be further discussed. If sweeping of a constant is being used, one wonders if estimating delta_e and delta_o is really useful or if we could just do parameter sweeping as well. \n\nTraining experiments:\n\nThe author don’t use the BCTS algorithm but the batch-BFS only for this set of experiments. However the author mention that they use another corrected update? Clarification needed.\nThe results in table 1 should also include the pre-trained rainbow + BCTS for a full comparison of the proposed methods.\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "This work requires models or simulators for tree search to be applied but overall it can touch a broad range of applications beyond games.", "review_text": "", "extra_scores": {"ethical_concerns": "No particular concerns.", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "1.5"}, "extra_text": {}}, {"review_id": "7Qz4WCuhI_", "reviewer": "Reviewer_z7iw", "summary": "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.", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "This work has some interesting elements, but is currently a bit unpolished and is not ready for publication. The paper would benefit from increased focus on the discussion of overestimation and its correction, the most interesting and novel components. The section on batching forward passes during search for efficiency is described as novel but is routine in the model-based RL community. The discussion of training with tree search distracts from the central topic of the paper as it does not involve the problems of overestimation and distribution shift. \n\n### Structure and text\nThe basic setting, using a learned agent at test time with a greater degree of tree search than was used at training time, is straightforward and fairly common in model-based RL; see for example the appendix of [5]. The introduction would be stronger if it did not re-introduce this idea before discussing the contribution of this work specifically.\n\nFigure 2 is nice and would be helpful if moved into the introduction to help set the scene for the reader.\n\nIn Figure 1, I believe the Q-values shown are the estimates (i.e. $\\hat Q$). It would be helpful if the \\hats were applied more consistently throughout the work.\n\nLemmas 3.1 and 3.2 are extremely unhelpful as written since they consist entirely of terms defined in the appendix. It would be much better if they were at least stated in a self-contained way, though of course the proofs can be deferred to the appendix. Even replacing them with easy-to-understand bounds would be more illuminating, especially given that we have already made some strong assumptions so the precise details are not so important.\n\n`\\widehat` would be helpful for defining the variance in Proposition 3.6.\n\nWhile the top-line numbers for BCTS are good, this work would be much more compelling with additional investigative experiments. Is overestimation on little-visited states happening in practice during tree search? How much of the time? Does that over-estimation decay as predictions are further in the future, or does performance improve with increasing $d$ due to the $\\gamma$ decay alone? (This would be surprising since $\\gamma$ is usually very near 1.) Does the overestimation scale as predicted by your theory? Overall the theory is fine, but the assumptions are such that it should only be thought of as something to build intuition; I want to know if that's _really_ what's going on with a learned agent.\n\n\n### Technical questions\nIn equation (1), I believe the max should be applied to both terms, i.e. the max over actions of (accumulated reward + bootstrap Q value)?\n\nAssumption 1 deals with state and action trajectories, and defines one variance if the first action matches the policy and a different one otherwise. What about the rest of the actions? That is, as the assumption is written now, the variance of $\\hat Q$ would be low even if every action other than the first was different than the one $\\pi_o$ would take. Similarly in Assumption 2. Is this intended, and necessary to simplify the analysis?\n\nTheorem 3.3 is maybe not quite what one would want to prove — the expected value of the estimator might be correct even while it produces incorrect actions much of the time. Indeed, by the assumptions we already have that $\\hat Q$ is unbiased, and yet optimizing against it has bad outcomes. It might be difficult to prove, but a more compelling result would be about exactly how often this algorithm takes the optimal action with respect to the true $Q_d^{\\pi_0}$.\n\nDo you use sticky actions, as described by [4]? This makes the environments substantially less deterministic.\n\nThe results shown in Table 1 on training with tree search are striking, but the comparisons are not necessarily the most meaningful. Training with tree search is a question of policy optimization against a  _known model_, whereas the full reinforcement learning setting (which the baselines are operating in) involves no prior knowledge about the environment. I'm not too bothered with this though since I would recommend removing this section entirely.\n\n\n### Novelty\nThis examination of overestimation and distribution shift is interesting, but not compellingly novel. I have seen related ideas in the literature though probably not constructed in exactly this way.\n\nThis paper touches on an interesting fundamental question that is fundamental in off-policy and offline reinforcement learning: how well does a Q function generalize, and how can we mitigate its variance for policy optimization? In the off-policy literature, there might be useful connections to draw to [2] and [3]. For offline RL, [1] might make a good starting point in the literature.\n\n### References\n1. Regularized Behavior Value Estimation, Gulcehre 2021\n2. Double Q-learning, van Hasselt 2010\n3. Addressing Function Approximation Error in Actor-Critic Methods, Fujimoto 2018\n4. Revisiting the Arcade Learning Environment: Evaluation Protocols and Open Problems for General Agents, Machado 2017\n5. Mastering Atari, Go, Chess and Shogi by Planning with a Learned Model, Schrittwieser 2019", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes", "review_text": "", "extra_scores": {"ethical_concerns": "N/A", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "5"}, "extra_text": {}}, {"review_id": "3YNWDmVxoUW", "reviewer": "Reviewer_AYhs", "summary": "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. \n\nThe second part of this paper proposes a more GPU compatible search method that seems to scale well when a GPU simulator is available. ", "questions": "", "limitations": "", "rating": 6, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Clarity:\n- the clarity of Section 3 can be greatly improeved. Can some of the lemmas be moved to the appendix? \n- also, I think more intuitive explaination of the correction term can help readers better understand the main idea. I am feeling that it is an intuitive method in the end that has been overcomplicated. \n\nOriginality:\n- the presented methods are new as I far as I know.\n\nQuality:\n- overall this submission seems technically sound.\n- what are the limitations of your methods?\n- does assumption 1 really holds in practice? \n\nSignificance:\n- this submission presents a nice solution to an interesting problem, which is the distribution shift problem when using pretrained value function in MCTS. ", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "My main suggestion is to improve the presentation of this submission, especially Section 3.\n\nSecond, I think more experimental results on other domains will also make the submission stronger. I am also interested to see search time constrained evaluation to understand to what extend the correction term can allow the performance of MCTS to be improeved with a pretrained value function. \n\nMinor:\n- is there link to be made to [1]?\n\n\n[1] Jiang, Nan, et al. \"The dependence of effective planning horizon on model accuracy.\" Proceedings of the 2015 International Conference on Autonomous Agents and Multiagent Systems. 2015.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "6"}, "extra_text": {}}]

Author rebuttal:

Meta-review: This is a solid paper and set of contributions. There are certainly areas where it could be improved; to me these are primarily around the structure of the paper and writing (not that these are poor, but there are these two essentially distinct contributions that some reviewers get tripped up on). z7iw points out that related ideas can be found in the literature and that this negatively impacts novelty, and I agree that this overlap exists but would argue that whatever minor hit novelty takes it is made up for by how this can unify understanding of where these challenges arise. There are several concrete suggestions (or areas for clarification) that the reviewers have brought up, and the authors are encouraged to include them in the final version of the paper.
Decision: Accept (Poster)