Policy improvement by planning with Gumbel

1 INTRODUCTION . In 2018 , AlphaZero ( Silver et al. , 2018 ) demonstrated a single algorithm achieving state-of-the-art results on Go , chess , and Shogi . The community reacted quickly . Leela Chess Zero ( Linscott et al. , 2018 ) was created to reproduce AlphaZero results on chess , winning Top Chess Engine Championship in 2019 . Soon , all top-rated classical chess engines replaced traditional evaluations functions with Efficiently Updatable Neural Network ( Nasu , 2018 ) . AlphaZero was itself generalized by MuZero ( Schrittwieser et al. , 2020 ) . While AlphaZero requires a black-box model of the environment , MuZero learns an abstract model of the environment . Essentially , MuZero learns the rules of Go , chess , and Shogi from interactions with the environment . This allows MuZero to excel also at Atari and continuous control from pixels ( Hubert et al. , 2021 ) . In this work , we redesign and improve AlphaZero . In particular , we consider the mechanisms by which AlphaZero selects and uses actions , which are based upon a variety of heuristic ideas that have proven especially effective in Go , chess , and Atari ( Silver et al. , 2018 ; Schrittwieser et al. , 2020 ) . However when using a small number of simulations , some of AlphaZero ’ s mechanisms perform poorly . We use the principle of policy improvement to suggest new mechanisms with a better theoretical foundation . More specifically , we consider each mechanism in turn , alongside our proposed modifications : • Selecting actions to search at the root node . To explore different actions during training , AlphaZero selects actions by adding Dirichlet noise to its policy network , and then performs a search using the perturbed policy . However , this does not ensure a policy improvement . We instead propose to sample actions without replacement by using the Gumbel-Top-k trick ( Section 2 ) , and perform a search using the same Gumbel values to influence the selection of the best action ( Section 3.3 ) , and show that this guarantees a policy improvement when action-values are correctly evaluated . • Selecting actions at the root node . AlphaZero uses a variant of the PUCB algorithm ( Rosin , 2011 ) to select actions at the root node . This algorithm was designed to optimize cumulative regret in a bandit-with-predictor setting ( i.e . given prior recommendations from the policy network ) . However , no ancestors are dependent upon the evaluation of the root node , and the performance of the Monte-Carlo tree search therefore only depends upon the final recommended action at the root node , and not upon the intermediate actions selected during search ( Bubeck et al. , 2011 ) . Consequently , we propose to use the Sequential Halving algorithm ( Karnin et al. , 2013 ) at the root node to optimize simple regret in a stochastic bandit with a predictor ( Section 3.4 ) . • Selecting actions in the environment . Once search is complete , AlphaZero selects an action by sampling an ( annealed ) softmax distribution based upon the visit counts of root actions resulting from the search procedure . We instead propose to select the singular action resulting from the Sequential Halving search procedure . • Policy network update . AlphaZero updates its policy network towards a softmax distribu- tion based upon the visit counts of root actions . However , even if the considered actions are correctly evaluated , this does not guarantee a policy improvement , especially when using small numbers of simulations ( Grill et al. , 2020 ) . We instead propose a policy improvement based upon the root action values computed during search , and update the policy network towards that policy improvement ( Section 4 ) . • Selecting actions at non-root nodes . AlphaZero uses the PUCT algorithm to select actions at non-root nodes . We instead propose to select actions according to a policy improvement ( similar to the proposal of Grill et al . ( 2020 ) ) based upon a completion of the action values . Furthermore , rather than sampling directly from this policy improvement , we propose a deterministic action selection procedure that matches the empirical visit counts to the desired policy improvement ( Section 5 ) . The proposed modifications are applicable also to MuZero or any agent with a policy network and an expensive Q-network . The modifications are most helpful when using a small number of simulations , relative to the number of actions . When using a large number of simulations , AlphaZero works well . We tried to ensure that the new search is principled , better with a smaller number of simulations , and never worse . We succeeded on all tested domains : Go , chess , and Atari . 2 BACKGROUND . Before explaining the improved search , we will explain the Gumbel-Max trick and the Gumbel-Topk trick . The Gumbel-Max trick was popularized by Gumbel-Softmax for a gradient approximation . In this paper , we are not interested in approximate gradients . Instead , we use the Gumbel-Top-k trick to sample without replacement . Gumbel-Max trick . ( Gumbel , 1954 ; Luce , 1959 ; Maddison et al. , 2017 ; Jang et al. , 2017 ) Let π be a categorical distribution with logits ∈ Rk , such that logits ( a ) is the logit of the action a . We can obtain a sample A from the distribution π by first generating a vector of k Gumbel variables and then taking argmax : ( g ∈ Rk ) ∼ Gumbel ( 0 ) ( 1 ) A = argmax a ( g ( a ) + logits ( a ) ) . ( 2 ) Gumbel-Top-k trick . ( Yellott , 1977 ; Vieira , 2014 ; Kool et al. , 2019 ) The Gumbel-Max trick can be generalized to sampling n actions without replacement , by taking n top actions : ( g ∈ Rk ) ∼ Gumbel ( 0 ) ( 3 ) A1 = argmax a ( g ( a ) + logits ( a ) ) ( 4 ) ... An = argmax a/∈ { A1 , ... , An−1 } ( g ( a ) + logits ( a ) ) . ( 5 ) We will denote the set of n top actions by argtop ( g + logits , n ) = { A1 , A2 , . . . , An } . 3 PLANNING AT THE ROOT . We are interested in improving AlphaZero Monte-Carlo Tree Search ( MCTS ) . In this section we will focus on the action selection at the root of the search tree . 3.1 PROBLEM SETTING . Both AlphaZero and MuZero have access to a policy network . At the root of the search tree , they can explore n simulations , before selecting an action for the real environment . We will formalize the problem as a deterministic bandit with a predictor and we will later extend it to a stochastic bandit and MCTS . Bandit . A k-armed deterministic bandit is a vector of Q-values q ∈ Rk , such that q ( a ) is the Q-value of the action a . The agent interacts with the bandit in n simulations ( aka rounds ) . In each simulation t ∈ { 1 , . . . , n } , the agent selects an action At ∈ { 0 , . . . , k − 1 } and visits the action to observe the Q-value q ( At ) . The objective is to maximize the Q-value from a special last action An+1 . That means we want to maximize E [ q ( An+1 ) ] . This objective is equivalent to minimization of simple regret . The simple regret differs from the cumulative regret from all n simulations . Bubeck et al . ( 2011 ) , Hay & Russell ( 2011 ) , and Tolpin & Shimony ( 2012 ) already argued that at the root of the search tree we care about the simple regret . The problem becomes interesting when the number of possible actions is larger than the number of simulations , i.e. , when k > n. For example , 19x19 Go has 362 possible actions and we will do experiments with as few as n = 2 simulations . Fortunately , the policy network can help . Predictor . In the bandit-with-predictor setting ( Rosin , 2011 ) , the agent is equipped with a predictor : the policy network . Before any interaction with the bandit , the policy network predicts the best action by producing a probability distribution π . The agent can use the policy network predictions to make more informed decisions . Policy improvement . Naturally , we would like to have an agent that acts better than , or as well as , the policy network . We would like to obtain a policy improvement . If the agent ’ s action selection produces a policy improvement , then E [ q ( An+1 ) ] ≥ ∑ a π ( a ) q ( a ) , ( 6 ) where the probability π ( a ) is the policy network prediction for the action a.1 The policy network can then keep improving by modeling an improved policy . 3.2 MOTIVATING COUNTEREXAMPLE . We will show that the commonly used heuristics fail to produce a policy improvement . Example 1 . Acting with the best action from the top-n most probable actions fails to produce a policy improvement . Let ’ s demonstrate that . Let q = ( 0 , 0 , 1 ) be the Q-values and let π = ( 0.5 , 0.3 , 0.2 ) be the probabilities produced by the policy network . The value of the policy network is ∑ a π ( a ) q ( a ) = 0.2 . For n = 2 simulations , the set of the most probable actions is { 0 , 1 } . With that , the heuristic would select An+1 = argmaxa∈ { 0,1 } q ( a ) . The expected value of such action is E [ q ( An+1 ) ] = 0 , which is worse than the value of the policy network . You can find other counterexamples by generating random q and π vectors and testing the policy improvement ( Inequality 6 ) . The AlphaZero action selection is explained in Appendix A . 3.3 PLANNING WITH GUMBEL . We will design a policy improvement algorithm for the deterministic bandit with a predictor π . After n simulations , the algorithm should propose an action An+1 with E [ q ( An+1 ) ] ≥ ∑ a π ( a ) q ( a ) . One possibility is to sample n actions from π , and then to select from the sampled actions the action with the highest q ( a ) . Instead of sampling with replacement , we can reduce the variance by sampling without replacement . Still , the sampled actions contain a limited amount of information about π . We should exploit the knowledge of π and its logits when selecting An+1 . The main idea is to sample n actions without replacement by using the Gumbel-Top-k trick , and then to use the same Gumbel g to select the action with the highest g ( a ) + logits ( a ) + σ ( q ( a ) ) . The σ can be any monotonically increasing transformation . The pseudocode for the algorithm is in Algorithm 1 . 1Inequality 6 can be strict , if we assume that an action has a positive advantage and its π ( a ) > 0 . Algorithm 1 Policy Improvement by Planning with Gumbel Require : k : number of actions . Require : n ≤ k : number of simulations . Require : logits ∈ Rk : predictor logits from a policy network π . Sample k Gumbel variables : ( g ∈ Rk ) ∼ Gumbel ( 0 ) Find n actions with the highest g ( a ) + logits ( a ) : Atopn = argtop ( g + logits , n ) Get q ( a ) for each a ∈ Atopn by visiting the actions . From the Atopn actions , find the action with the highest g ( a ) + logits ( a ) + σ ( q ( a ) ) : An+1 = argmaxa∈Atopn ( g ( a ) + logits ( a ) + σ ( q ( a ) ) ) return An+1 The algorithm produces a policy improvement , because q ( argmax a∈Atopn ( g ( a ) + logits ( a ) + σ ( q ( a ) ) ) ) ≥ q ( argmax a∈Atopn ( g ( a ) + logits ( a ) ) ) . ( 7 ) This holds for any Gumbel g , so it holds also for expectations : E [ q ( An+1 ) ] ≥ EA∼π [ q ( A ) ] . The argmaxa∈Atopn ( g ( a ) + logits ( a ) ) is equivalent to sampling from the policy network π ( see the Gumbel-Max trick ) . By using the same Gumbel vector g in the argtop and argmax , we avoid a double-counting bias . The prior knowledge contained in the logits can help on partially observable environments , or when working with approximate or stochastic Q-values .

Summary: This paper considers MCTS with learned search guidance, as in AlphaZero, MuZero, etc. The work proposes several adjustments to the prior works, particularly regarding the way in which actions are selected at the root and non-root nodes, at training and during evaluation; and also the way in which the policy is updated after search. In a simplified bandits setting, the authors point out that the previous method of performing policy updates is not guaranteed to result in a policy improvement, and they propose using a Gumbel reparameterization trick to overcome this limitation. Experiments in Go, Chess, and Atari show that the adjustments are beneficial in the regime of low simulation count.