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: The paper proposes a number of principled algorithmic modifications to state-of-the-art planning algorithms (AlphaZero, MuZero) for improving performance in settings with many actions and a relatively small computation and / or sample budget. The main contributions are algorithmic and empirical. The key ideas include the use of the Gumbel-max and top-k tricks along with the use of sequential halving to improve online planning. The paper also proposes a planning-learning loop wherein a policy using the estimated (completed) Q values is learned as well as a different selection policy at non-root nodes. The experiments show that the Gumbel variants of AlphaZero and MuZero perform well in low search budget settings in the domains of Go, Chess and Atari.