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 the Virtual MCTS (V- MCTS), a variant of MCTS that mimics the human behavior, and is 50% more sample efficient, by performing a type of forward pruning. MCTS is characterized as a model-based reinforcement learning algorithm, that imagines what the future would look like, using terminology borrowed from Sutton in his description of Dyna, a model-based algorithm.  I agree with the planning part, but am less certain if I agree with the learning characterization of MCTS. Often, MCTS is used inside a model-based approach as the planning component. MCTS is usually not regarded as a full model-based approach.