Spending Thinking Time Wisely: Accelerating MCTS with Virtual Expansions

1 INTRODUCTION . When artificial intelligence was first studied in the 1950s , researchers seek to answer the question of what is the solution to the question if the agent were “ perfect rational ” . The term “ perfect rational ” here refers to the decision made with infinite amounts of computations . However , without taking into consideration the practical computation time , one can only solve small-scale problems , since classical search algorithms usually exhibit exponential running time . Recent AI researches no longer seek to achieve “ perfect rational ” , but instead carefully trade-off computation versus the level of rationality . People have developed computational models like “ bounded optimality ” to model these settings ( Russell & Subramanian , 1994 ) . The increasing level of rationality under the same computational budget has given us a lot of AI successes nowadays . Notable algorithms include the Monte-Carlo sampling algorithms , the variational inference algorithms , and using neural networks as universal function approximators ( Coulom , 2006 ; Chaslot et al. , 2008 ; Gelly & Silver , 2011 ; Silver et al. , 2016 ; Hoffman et al. , 2013 ) . More recently , MCTS-based RL algorithms have achieved a lot of success , mainly in board games . The most notable achievement is that AlphaGO beats Hui Fan in 2015 ( Silver et al. , 2016 ) . This is the first time that a computer program beats a human professional player . After that , AlphaGo beats two top-ranking human players , Lee Sedol in 2016 and Jie Ke in 2017 , the latter of which rank first worldwide at the time . Later , the MCTS-based RL algorithms are further extended to other board games , as well as the Atari video games ( Schrittwieser et al. , 2020 ) . EfficientZero ( Ye et al. , 2021 ) greatly improves the sample efficiency of MCTS-based RL algorithms , shedding light on its future applications in real-world applications like robotics and self-driving . Despite the impressive performance of MCTS-based RL algorithms , they require massive computations to train and evaluate . For example , Schrittwieser et al . ( 2020 ) used 1000 TPUs trained for 12 hours to learn the game of GO , and for a single Atari game , it needs 40 TPUs to train 12 hours . Compared to previous algorithms on the Atari games benchmark , it needs around two orders of magnitude more compute . This prohibitively large computational requirement has slowed down both the further development of MCTS-based RL algorithms , as well as its practical use . Under the hood , MCTS-based RL algorithms are model-based methods , that imagine what the futures look like when doing different future action sequences . However , this imaging process for the current method is not computationally efficient . For example , AlphaGo needs to look ahead 1600 game states to place a single stone . On the contrary , top human professional players can only think through around 100-200 game states per minute ( Silver et al. , 2016 ) . Besides being computationally inefficient , the current MCTS algorithm deals with easy cases and hard ones with the same computational budget . On the other hand , human knows to use their time when it is most needed . In this paper , we aim to design new algorithms that save the computational time of the MCTS-based RL methods . More specifically , we are interested in pushing the Pareto front of the rationality level - computation curve . Empirical results show that our method can achieve comparable performance while requiring less than 50 % simulations to search on average . 2 RELATED WORK . 2.1 MULTI-ARMED BANDIT PROBLEM . RL algorithms are always brought into the exploration and exploitation dilemma . Multi-armed bandit ( MAB ) problem ( Berry & Fristedt , 1985 ; Auer et al. , 2002 ; Lattimore & Szepesvári , 2020 ) is one of the most extensively studied but fundamental instances . TheK-armed MAB problem is a sequential game with a collection of K unidentified but independent reward distributions , each associated with the corresponding arms . For each round , the learner pulls an arm and receives a reward sampled from the corresponding distributions . The optimal policy of the learner for the MAB problem is to maximize the cumulative rewards obtained from the sequential decisions . In the case where the cost of pulling arms is little , the learner is allowed to trial and error for enough times until convergence . A series of upper confidence bound ( UCB ) algorithms ( Auer et al. , 2002 ; Bubeck & Cesa-Bianchi , 2012 ) are proposed to solve the stochastic MAB problem and they have theoretical bounds . When there exist costs for each trial , pure exploration attempts to make the best use of the finite trials ( Bubeck et al. , 2011 ; Lattimore & Szepesvári , 2020 ) . Kocsis & Szepesvári ( 2006 ) proposed UCT to adapt UCB algorithms to the tree structures , which is the basis of MCTS . 2.2 REINFORCEMENT LEARNING WITH MCTS . For a long time , Computer Go is regarded as a very challenging game ( Bouzy & Cazenave , 2001 ; Cai & Wunsch , 2007 ) . Researchers attempt to use Monte-Carlo techniques that evaluate the value of the node state through random playouts ( Bouzy & Helmstetter , 2004 ; Gelly & Silver , 2007 ; 2008 ; Silver et al. , 2016 ) . Afterwards , UCT has generally replaced those earlier heuristic methods for Monte-Carlo tree search ( MCTS ) . UCT algorithms ( Kocsis & Szepesvári , 2006 ) apply UCB1 to select action at each node of the tree . Recently , MCTS-based methods ( Silver et al. , 2016 ; 2017 ; 2018 ; Schrittwieser et al. , 2020 ) have become increasingly popular and achieved super-human performances on board games because of the strong ability to search . Modern MCTS-based RL algorithms include four stages in each search iteration , namely simulation : selection , expansion , evaluation , and backpropagation . The selection stage targets selecting a new leaf node with UCT . The expansion stage expands the selected node and updates the search tree . The evaluation stage evaluates the value of the new node . The backpropagation stage propagates the newly computed value to the nodes along the search path to obtain more accurate Q-values with Bellman backup . However , search is quite time-consuming , which prevents MCTS-based methods to be used in wider scenarios . 2.3 ACCELERATION OF MCTS . There are two kinds of bottlenecks in MCTS in the aspect of speed : the evaluation/selection stage of each iteration and the outer search loop . In the previous research , people attempted to evaluate the node value by random playouts to the end of the game , which makes the evaluation stage quite expensive . In addition , compared to other model-free RL methods like PPO ( Schulman et al. , 2017 ) and SAC ( Haarnoja et al. , 2018 ) , MCTS-based algorithms have much larger computations due to the search loop . Therefore , a lot of works are devoted to accelerating MCTS . Some heuristic pruning methods ( Gelly et al. , 2006 ; Wang & Gelly , 2007 ; Sephton et al. , 2014 ; Baier & Winands , 2014 ; 2018 ) are developed to make the selection or evaluation more effectively . Lorentz ( 2015 ) proposed early playout termination of MCTS ( MCTS-EPT ) to stop the random playouts early and use an evaluation function to assess win or loss , which is an improvement in the evaluation stage of normal MCTS . And Hsueh et al . ( 2016 ) applied MCTS-EPT to the Chinese dark chess . Afterwards , MCTSEPT similar ideas have been applied in the evaluation stage of AlphaGoZero ( Silver et al. , 2017 ) and later MCTS-based methods ( Silver et al. , 2018 ; Schrittwieser et al. , 2020 ; Ye et al. , 2021 ) including our baseline models . They evaluate the Q-values through evaluation networks instead of running playouts to the end . However , these methods focus on the specific stage of the search iteration to accelerate the MCTS . We propose Virtual MCTS , which aims to terminate the outer search iteration adaptively in MCTS under distinct circumstances without sacrificing the final policy quality . 3 BACKGROUND . The AlphaGo series of work ( Silver et al. , 2016 ; 2017 ; 2018 ; Schrittwieser et al. , 2020 ) are all MCTS-based reinforcement learning algorithms . Those algorithms assume the environment transition dynamics are known or learn the environment dynamics . Based on the dynamics , they use the Monte-Carlo tree search ( MCTS ) as the policy improvement operator . I.e . taking in the current policy , MCTS returns a better policy with the search algorithm . The systematic search allows the MCTS-based RL algorithm to quickly improve the policy , and perform much better in the setting where a lot of reasoning is required . MCTS is the core component in the algorithms like AlphaGo . 3.1 MCTS . In this part , we give a brief introduction to the MCTS method implemented in reinforcement learning applications . MCTS takes in the current MDP state and runs a search algorithm guided by the current policy function . It outputs an improved policy of the current state . The improved policy is later to select an action in the environment . In the selection stage , an action will be selected by maximizing over UCB . Specifically , AlphaZero ( Silver et al. , 2018 ) and MuZero ( Schrittwieser et al. , 2020 ) are developed based on a variant of UCB , P-UCT ( Rosin , 2011 ) and have achieved great success on board games and Atari games . The formula of P-UCT in the two methods is the Eq ( 1 ) : ak = argmax a∈A Q ( s , a ) + P ( s , a ) √∑ bN ( s , b ) 1 +N ( s , a ) ( c1 + log ( ∑ bN ( s , b ) + c2 + 1 c2 ) ) , ( 1 ) where k is the index of the iterative step , A is the action set , Q ( s , a ) is the estimated Q-value , P ( s , a ) is the policy prior obtained from neural networks and N ( s , a ) is the visit counts to select the action a from the state s. The output of MCTS is the visit count of each action of the root node . After N search iterations , the final policy π ( s ) is defined as the normalized root visit count distribution πN ( s ) , where πk ( s , a ) = ( N ( s , a ) ) / ∑ b∈AN ( s , b ) = N ( s , a ) /k , a ∈ A . For simplification , we use πk in place of πk ( s ) sometimes . In our method , we propose to approximate the final policy πN ( s ) with π̂k ( s ) , which we name as a policy candidate , through a new expansion method and a termination rule . In this way , the number of iterations in MCTS can be reduced from N to k . 3.2 COMPUTATION REQUIREMENT . Most of the computations in MCTS-based RL are in the MCTS procedure . For each action taken by MCTS , it needs N times neural network evaluations , where N is the number of search iterations in MCTS . Traditional RL algorithms , such as PPO ( Schulman et al. , 2017 ) or DQN ( Mnih et al. , 2015 ) , only need a single neural network evaluation per action . Thus , MCTS-based RL is roughly N times computationally more expansive than traditional RL algorithms . In practice , training a single Atari game needs 12 hours of computation time on 40 TPUs ( Schrittwieser et al. , 2020 ) . The computation need is roughly two orders of magnitude more than traditional RL algorithms ( Schulman et al. , 2017 ) , although the final performance of MuZero is much better .

Summary: This paper proposes an approach for a significant speedup of Monte-Carlo Tree Search (MCTS) at a relatively small cost in playing strength. The basic idea is that, when the change in the distribution of visit counts between two different time points $\frac{k}{2}$ and $k$ is less than some constant $\epsilon$, it can also be shown that the remaining change in distributions between $k$ and $N$ (where $N$ is the maximum visit count allowed by some budget) will be bounded below some value, and if we consider such a maximum possible error to be sufficiently small we can just terminate the search early at time $k$. The paper also proposes a very simple but important idea called Virtual Expansions, which basically consists of running $N - k$ additional iterations of the bandit algorithm used by MCTS solely in the root node, providing the current estimated average values as rewards for every pull (i.e., leaving average reward estimates unchanged), in order to transform the distribution of visits at iteration $k$ into a prediction of the distribution we would end up with after $N$ iterations.  Several empirical evaluations in Go ($9\times9$ board) and a few Atari games, including ablation studies, demonstrate that the approach can substantially reduce the search times and self-play training times, while only decreasing playing strength by a small amount.