Monte Carlo Tree Search With Iteratively Refining State Abstractions

Decision-time planning is the process of constructing a transient , local policy with the intent of using it to make the immediate decision . Monte Carlo tree search ( MCTS ) , which has been leveraged to great success in Go , chess , shogi , Hex , Atari , and other settings , is perhaps the most celebrated decision-time planning algorithm . Unfortunately , in its original form , MCTS can degenerate to one-step search in domains with stochasticity . Progressive widening is one way to ameliorate this issue , but we argue that it possesses undesirable properties for some settings . In this work , we present a method , called abstraction refining , for extending MCTS to stochastic environments which , unlike progressive widening , leverages the geometry of the state space . We argue that leveraging the geometry of the space can offer advantages . To support this claim , we present a series of experimental examples in which abstraction refining outperforms progressive widening , given equal simulation budgets . 1 Introduction . In control problems , an agent makes sequential decisions toward the end of accumulating a large amount of reward [ 21 ] . Many reinforcement learning algorithms approach this task by computing a policy , which , given a state , dictates the agent ’ s behavior ; others operate under the paradigm of decision-time planning , in which , after the agent has arrived at a state , it spends additional computation constructing or revising its policy for the immediate decision ( s ) . The decision-time planning paradigm can be advantageous because it allows for an additional step of policy improvement that is not available to agents acting according to a predetermined policy . Of the numerous ways to perform decision-time planning , Monte Carlo tree search ( MCTS ) is among the most influential [ 6 , 10 , 15 ] . Algorithms powered by MCTS , like AlphaZero and MuZero , yield strong performance in games with complex value function landscapes , like Go , chess , Hex , and shogi [ 2 , 20 , 19 ] . MuZero even shows strong performance on Atari games , where model-free algorithms had previously been the most performant . Unfortunately , in its original form , MCTS can degenerate to a one-step search when the transition function has an infinite support . Even if the support is finite but non-trivial , MCTS may be relegated to building shallow search trees . However , despite the importance of stochasticity in real-world problems , relatively little attention has been paid to this issue , perhaps as a result of the fact that many popular benchmarks are deterministic or close-to-deterministic . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . One way to address this issue is by using progressive widening [ 9 ] . Progressive widening works by alternating between adding new children and selecting among existing children such that , asympototically , the search tree becomes both infinitely deep and infinitely wide . Progressive widening is also advantageous because it includes hyperparameters that control its propensity to add new children ( as opposed to selecting among existing children ) . These hyperparameters can interpolate between width-wise expansion only ( vanilla MCTS ) and depth-wise expansion only ( transition determinization ) . As a result , with proper tuning , progressive widening is capable of performing well in both settings with important stochasticity and in settings with unimportant ( or non-existent ) stochasticity . However , progressive widening ’ s decision rule for expansion is the same for every transition . In settings with non-uniform stochasticity , we argue that this property is undesirable , as it requires progressive widening to compromise between focusing on unimportant stochasticity and ignoring important stochasticity . In this work , we propose a new method , called abstraction refining , for extending MCTS to stochastic settings . Abstraction refining uses random and iteratively refining state abstractions defined by the geometry of the state space . If a newly sampled state is similar to a state that is already in the tree , it is discarded in favor of the pre-established state ; otherwise , the newly sampled state is added to the tree . Because the criteria for similarity becomes more strict the more often an abstraction is used , abstraction refining guarantees that , in the limit , the search tree will grow both infinitely wide and infinitely deep . In addition to proposing the abstraction refining algorithm , our contributions are twofold . First , we present both a proof that , when used for policy evaluation , abstraction refining converges almost surely in finite MDPs . Second , we make an empirical case that , provided a good notion of state similarity , abstraction refining can outperform progressive widening in settings with stochasticity of varying importance , given an equal simulation budget . To make this case , we present a series of domains in which i ) stochasticity is sometimes , but not always , relevant to performing well on the task and ii ) the naive notion of distance between states is reflective of behavioral similarity ( though we also include some settings in which we use a learned notion of distance ) . In these domains , we find that abstraction refining can outperform progressive widening . 2 Background . We consider a Markov decision process ( MDP ) setting 〈S , A , T , R , γ〉 , with state space S , action spaceA , transition function T , reward functionR and discount factor γ . The objective is to determine a policy π that achieves a large expected cumulative reward Eπ [ ∑∞ t=0 γ tR ( st , at ) ] . Rather than explicitly computing a function π , decision-time planning algorithms ’ policies are defined implicitly by their planning procedures . These planning procedures are given generative access to ( i.e. , the ability to sample from ) the transition function T . We also consider a Markov reward process ( MRP ) setting 〈S , T , R , γ〉 , with state space S , transition function T , reward functionR , and discount factor γ . An MRP is like an MDP , except that the agent is a passive observer , rather than a participator . Equivalently , one could also think of an MRP as an MDP in which the policy has been fixed and baked into the transition and reward functions . For the purposes of this work , the goal of an agent in an MRP is to evaluate the expected return from state—i.e. , to perform policy evaluation . 2.1 Monte Carlo Tree Search . MCTS works by iteratively building a look-ahead tree of states . For each action a at a state s , the algorithm keeps track of the number of times the action has been selected at that state N ( s , a ) and the average of the value assessments of that action Q ( s , a ) . At each iteration , the agent ( i ) selects a leaf node to expand ; ( ii ) expands the tree by adding a child node of the leaf node ; ( iii ) evaluates the newly added state , and ( iv ) backpropagates the evaluation from the new child to the root . Selection In the selection phase , starting from the root node , which corresponds to the state s0 currently occupied by the agent , MCTS traverses the tree until a leaf node is reached . The tree traversal is dictated by a selection policy . The most commonly used selection policies are based on the UCT algorithm [ 15 ] , or its variants , like PUCT [ 18 ] . Although it is less common , MCTS can also be used as a policy evaluation algorithm . In this case , the selection policy is set to the target policy . Expansion Once a leaf node is reached , the tree is expanded by adding a child of the leaf node to the tree . If the leaf node is a terminal state , no children are added . Evaluation Whenever MCTS adds a new state to the tree , it evaluates the state . Traditionally , MCTS practitioners performed evaluation by generating a rollout to the end of the game using a fast-toevaluate simulation policy . Modern implementations often use a parameterized value function in lieu of simulation rollouts . Backpropagation After MCTS evaluates a state , it updates the nodes along the path from the root to the evaluated state . Denote the path from the root to the leaf using s0 , a0 , r0 , . . . , sT . Then for each pair ( sk , ak ) , MCTS updates Q ( sk , ak ) using a weighted average between the existing value Q ( sk , ak ) and the leaf evaluation summed with the accumulated reward rk+1+ · · ·+rT . Additionally each visit counter N ( sk , ak ) gets incremented by one . Selecting an Action To Execute After the search is over , MCTS selects an action for the agent to take . Common ways of doing this include selecting the action at the root node with the highest average Q-value , the action with highest lower confidence value or the action with the highest visit count . When used in an evaluation context , MCTS returns the average value of the root node . What about Stochasticity ? The description of MCTS above assumes that the environment is determinstic . In games with stochasticity , the selection process has an additional complication . Specifically , after an action a is selected at a state s , there is no uniquely determinable next state s′ , but rather , a distribution over next states T ( s , a ) . The most obvious way to extend MCTS to this setting is to sample s′ ∼ T ( s , a ) . If there is already a child node for s′ , then that node is selected and the selection process continues . Otherwise , if there is no child node for s′ , one is added to the tree and evaluated . The consequence of this modification is that , for a fixed number of iterations , the search tree can not be as deep as it would be in a deterministic setting . In the most extreme case , where there are an infinite number of possible next states , this modification can degenerate MCTS into one-step search ( because it may be the case that no state is sampled twice ) . This means that the performance of MCTS in games in which deep search trees are crucial for good performance could be ruined by arbitrarily small ( but infinite support ) perturbations to state representations . 2.2 Progressive Widening . Progressive widening [ 9 ] is one way to augment MCTS in such a way that it can build deep search trees in arbitrarily stochastic domains.1 The basic idea behind progressive widening is to alternate between sampling new next states and sampling among the next states that are already in the tree . If the state-action pairN ( s , a ) has been tried a large number of times relative to the number of successor states in the tree , progressive widening adds a new state . On the other hand , if there are a large number of successor states relative to the number of times N ( s , a ) has been tried , progressive widening samples among the successor states already in the tree . Whether or not to sample is determined by the boolean num_children [ s , a ] < k ·N ( s , a ) α where α ∈ [ 0 , 1 ] and k ∈ R+ are hyperparameters . Simplified pseudocode for the sampling step of progressive widening is shown in Algorithm 1 . Informally , the hyperparameter α can be thought of as a propensity to select among existing children , rather than adding new children . When k = 1 , α = 0 , progressive widening reduces to transition determinization—the first time that a state-action pair is visited , a successor state is sampled and added to the tree ; thereafter , the same successor state is selected every time the state-action pair is visited . When k = 1 , α = 1 , progressive widening reduces to vanilla MCTS—every time a state action pair is encountered , a new successor state is sampled and added to the tree . If k and α can be properly tuned , progressive widening offers flexibility . In domains in which stochasticity is important , α can be set to one or close to one . In domains in which stochasticity can be ignored , α can be set to zero or close to zero . Otherwise , an intermediate value often works well . However , there is a downside to progressive widening . In particular , it can not discriminate between stochasticity that matters and stochasticity can safely be ignored . This means that , if some stochasticity 1Note that we call progressive widening here is generally referred to as double progressive widening . Double progressive widening uses progressive widening over both the action space and the state space . For simplicity , this work restricts its attention to small action spaces , and therefore uses the term progressive widening to refer to the state space behavior . is important but some is unimportant , the best thing that progressive widening can do requires either placing too much emphasis on unimportant stochasticity , or too little on important stochasticity . Algorithm 1 Progressive Widening procedure SAMPLE ( s , a ) N ← num_visits [ s , a ] if num_children [ s , a ] < kNα then s′ ∼ T ( s , a ) if s′ 6∈ children [ s , a ] then children [ s , a ] .add ( s′ ) return s′ else return sample ( children [ s , a ] ) Algorithm 2 Abstraction Refining procedure SAMPLE ( s , a ) s′ ∼ T ( s , a ) s′′ ← nearest_neighbor ( s′ , children [ s , a ] ) if d ( s′ , s′′ ) < num_visits [ s , a , s′′ ] then return s′′ else children [ s , a ] .add ( s′ ) return s′

Summary: The paper proposes a novel extension of MCTS to tackle stochastic state transitions with infinite support. It does so by leveraging the geometry of the state space to "abstract" "similar" states together during the tree building phase. Abstraction refining is introduced to assure that in the limit of infinite samples, the search-tree converges to the full tree, by introducing a schedule on the maximum abstraction distance, based on the number of samples observed. The authors provide a proof of asymptotic convergence of abstraction refining in the policy evaluation setting. A brief discussion is done on the choice of distance function to employ during the state abstraction. Finally, an experimental campaign is conducted in several evaluation and control tasks, that show that abstraction refining consistently outperforms it's main competitor (progressive widening) in handling stochastic transitions with infinite or very large support.