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: 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.