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

1 Introduction . 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 can not 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 . 2 Preliminaries . 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 ) 3 Solving the Tree Search Distribution Shift . 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 . 3.1 Performance degradation with vanilla tree search . 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 ) .

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.