Planning in Stochastic Environments with a Learned Model
Authors: Ioannis Antonoglou
Julian Schrittwieser
Sherjil Ozair
Thomas K Hubert
David Silver
Venue: ICLR 2022 Spotlight (2022)
Keywords: model-based reinforcement learning
deep reinforcement learning
tree based search
MCTS

Abstract: Model-based reinforcement learning has proven highly successful. However, learning a model in isolation from its use during planning is problematic in complex environments. To date, the most effective techniques have instead combined value-equivalent model learning with powerful tree-search methods. This approach is exemplified by MuZero, which has achieved state-of-the-art performance in a wide range of domains, from board games to visually rich environments, with discrete and continuous action spaces, in online and offline settings. However, previous instantiations of this approach were limited to the use of deterministic models. This limits their performance in environments that are inherently stochastic, partially observed, or so large and complex that they appear stochastic to a finite agent. In this paper we extend this approach to learn and plan with stochastic models. Specifically, we introduce a new algorithm, Stochastic MuZero, that learns a stochastic model incorporating afterstates, and uses this model to perform a stochastic tree search. Stochastic MuZero matched or exceeded the state of the art in a set of canonical single and multi-agent environments, including 2048 and backgammon, while maintaining the same performance as standard MuZero in the game of Go.

Published as a conference paper at ICLR 2022
PLANNING IN STOCHASTIC ENVIRONMENTS WITH A LEARNED MODEL
Ioannis Antonoglou"
Julian Schrittwieser'
Sherjil Ozair'
Thomas Hubert'
David S2lver
'DeepMind, London, UK University College London
ABSTRACT
Model-based reinforcement learning has proven highly successful. However, learn- ing model in isolation from its during planning is problematic in complex environments. To date, the most effective techniques have instead combined value- equivalent model learning with powerful tree-search methods. This approach is exemplified by MuZero, which has achieved state-of-the-ar performance a wide range of domains, from board games to visually rich environments, with discrete and continuous action spaces, in online and offline settings. However, previous instantiations of this approach were limited the use of deterministic models. This limits their performance in environments that are inherently stochastic, partially observed, or so large and complex that they appear stochastic this paper we extend this approach to learn and plan with stochastic models. Specif- ically, we introduce a new algorithm, Stochastic MuZero, that learns stochastic model incorporating afterstates, and uses this model to perform stochastic tree search. Stochastic MuZero matched or exceeded the state the art set canonical single and multi-agent environments, including 2048 and backgammon, while maintaining the superhuman performance standard MuZero in the game Go.
finite agent. In
INTRODUCTION
Constructing plans and exccuting them is an important feature of human and animal behaviour. In the field of artificial intelligence there has been great amount of research into adding planning capabilities to intelligent agents.
planning algorithms have shown lot of success in a
wide variety of environments such card games (Moravčík et al.) (017).. board games (Campbell et al., 2002; Silver et al. (016) and more recently video games (Schrittwieser et al.) 2020) and continuous control tasks (Hubert et al., (021). Most tree search methods assume that the agent has access to a perfect simulator of the environment, whereas real-world environments are typically unknown.
Model-based reinforcement learning algorithms combine a model-learning component, which es- timates the dynamics of the environment, with a planning component, using the learned model as simulator. However, learning model in isolation from its use during planning has proven to be problematic in complex environments (van Hasselt al. 2019). Instead, value-equivalent model- learning methods (Silver et al. 2017 Farahmand et al. 2017: Oh et al. 2017, Grimm et al. 2020) identify a model that reconstructs only those quantities required for planning. The most successful method, MuZero (Schrittwieser et al. 2020) learns model that reconstructs reward, value and policy, and uses this model to perform a powerful Monte Carlo tree search. MuZero achieved superhuman results in Go, chess, shogi and Atari without any prior knowledge of the rules, and has also achieved -fate-of-the-ar performance in large and continuous action spaces (Hubert et al.) 2021) and offline reinforcement learning (Schrittwieser al. 202I).
However, value equivalent methods such as MuZero have in practice been limited to a deterministic class of models, which severely limits their applicability. Many environments are inherently stochastic

---

Published as a conference paper at ICLR 2022
and may be poorly approximated by deterministic model. Partially observed environments may also be perceived by the agent as stochastic, whenever aliased states cannot be disambiguated, Similarly, large and complex environments may appear stochastic to small agent with finite capacity.
In this paper we introduce the first empirically effective approach for handling stochasticity in value equivalent model-leaming and planning. The model is factored to first transition deterninistically from state to an afterstate, and then to branch stochastically from the afterstate to the next state. This factored model is trained end-to-end so as maintain value equivalence for both state value function and action value function respectively, and a stochastic planning method is applied to the model.
We apply these ideas MuZero, using discrete generative network represent the model, and modifying the Monte Carlo tree search to effectively use the factored model. We apply our method, Stochastic MuZero, to several environments in which handling stochasticity important. First, we consider the popular stochastic puzzle game 2048, in which the prior state of the art exploits a perfect simulator and significant handcrafted domain knowledge. achieved better results without any domain knowledge. Second, we consider the classic stochastic two-player game of backgammon, in which near-optimal play has been achieved using a perfect simulator. Stochastic MuZero matches this performance without any prior knowledge of the game rules. Finally, we evaluated our method in the deterministic board game of Go. There our method matched the performance of MuZero, demonstrating that Stochastic Mulero extends MuZero without sacrificing performance.
our experiments, Stochastic MuZero
2 RELATED WORK
Observation models (Oh et al. 2015 Chiappa et al. 2017 Lukasz Kaiser et al. 2020 explicitly
learn the dynamics of an environment by fitting model of observations and rewards observed transitions. Subsequently, these models can be combined with model-free learning rule in Dyna fashion (Sutton 199I). However, modeling high dimensional image observations can be computationally prohibitive, prone to high error accumulation as the model is unrolled for multiple steps, and limiting since the capacity of the model could be spent on background features which are not helpful for the problem at hand. These issues make such models unconducive for planning. Finally, van Hasselt et al. (2019) argues that Dyna- Dyna-based methods are unlikely to outperform model- approaches that replay buffer.
Latent models (Schritwwisser et al., 2020; Oh et al. 2017 Hafner et al. 2021 Henaff et al. 2017)
attempt to overcome the limitations of observation models by learning recurrent networks that operate on latent states. In this framework, the model is conditioned on the current observation and future trained to make predictions about rewards,
actions and unrolled for is steps. Subsequently, values, policies or observations at each timestep based on the current latent state. The model can then be combined with tree- bree- planning algorithm or used to generate synthetic trajectories. Recently, MuZero has shown that it is possible to use this approach to achieve Mate-off-hhe-rt performance in many challenging domains (Hubert et al. (02T) while using less data (Schrittwieser al.) (02T). However, most approaches, including MuZero, use a deterministic function to model the environment dynamics, which limits their applicability to deterministic or weakly stochastic environments.
Stochastic latent models are stochastic models of the environment dynamics that operate on latent states. In (Hafner et al., 202T) the authors propose a recurrent state-space model which consists of three main modules, recurrent module which generates the deterministic recurrent state ht, a representation model which combines ht with the current observation to generate a distribution over stochastic states 8 and plays the role of the posterior, and transition predictor which depends only on h and acts as the prior of the model. By combining the deterministic and stochastic states he and &t the model is trained to predict the current observation or, the transition reward r, and the discount dg. The next deterministic recurrent state is generated using ht, 8, and action de. The stochastic states 8 are modeled as multidimensional multinomial variables. The learned model is then used to generate synthetic data which are used train an actor-critic model-free agent. The authors show that their approach outperforms pure model-free methods but it fails to achieve the performance of MuZero which combines its learned model with planning.
2

---

Published as a conference paper at ICLR 2022
In (Ozair al. 2021) the authors learn stochastic transition model using VQ-VAE generative network (van den Oord et al. 2017) and subsequently combine it with MCTS. They show that their method match the performance of MuZero in chess, while viewing the problem as single- player task and implicitly learning to model the behaviour of the opponent. Despite its promise their approach was only applied in supervised setting using expert data, and did not address the challenges of learning a stochastic model in the reinforcement learning setting. Moreover, the learned model was trained explicitly predict the observation at every step, which can be limiting factor in terms of computation and model efficiency when dealing with high dimensional observations. Finally, the authors used two stage training process: first, a model learns latent representations of the observations, then these representations are used to learn a transition model. This makes it hard to apply this approach in the reinforcement learning setting.
BACKGROUND
3.1 MuZero
MuZero is model-based general reinforcement learning agent which combines a learned model of the environment dynamics with a Monte Carlo free search planning algorithm. The model is conditioned on the history of observations 0g: at timestep t and a sequence of future actions dra+KK and it is trained to predict the search policies Tr.+KK values Vi and intermediate rewards Tua+KK at each future timestep. MuZero uses deterministic functions for its model, and thus it implicitly assumes that the underlying environment dynamics are also deterministic. MuZero uses its dynamics model to plan ahead at each time step and the outcome of its MCTS search to select an action and as targets for its policy improvement operator.
Model MuZero's 's learned model consists of functions: a representation function h, a dynamics function g and prediction function f. The representation function maps the current history of observations ogt into a latent state sf. The dynamics function receives the previous latent state sf and combines it with an action az+* to produce the next latent state s and the reward rf. Finally, the prediction function f receives each latent state sf as input and computes the policy of and value ox. Given sequence of policy Te.r, value 2 and reward u4r targets, the model is trained minimize the loss shown [
K Σ ((TT++)PP) + Σ (' of) + Σ k=0
K
L Muleero
((c+kki7))
(I)
k=
The policy targets Te+h correspond to the MCTS policy that generated when searching from observation 05t+k. The value targets 2++& are computed using n-step returns (Sutton & Barto 2018). Finally, the reward targets u++ correspond to the real instantaneous rewards observed when this sequence was generated.
Search MuZero uses variant of the MCTS tree based algorithm first proposed in (Silver et al. 2018). The tree is constructed recursively through a number of simulations. Each simulation consists of phases: selection, expansion and backpropagation. During the selection phase the tree is traversed starting from the root node until a leaf edge is reached. At each internal node s the algorithm selects the action which maximizes the upper confidence bound proposed in (Silver et al.) 2016) and shown in equation 2
a = arg max Q(s.a)+ s). VI+ ((s 1+ ((,
N,, a2 +1 (
a1 + log
(2)
Here, Q(s, a) is the value estimate for action a, (s, a) the visit count, P(a s) the prior probability of selecting action and a;, 02 are constants which control the relative importance of the Q(s.-) estimates and prior probabilities P(- s). In the next phase expansion, the leaf edge is expanded by querying the MuZero model and new node is added to the tree. Finally, during the backpropagation phase the value estimate of the newly added edge is backpropagated up the tree using the n-step return estimate.
3

---

Published as a conference paper at ICLR 2022
3.2 VECTOR QUANTISED VARIATIONAL AUTOENCODER
Vector Quantised Variational AutoEncoder (VQ-VAE, van den Oord et al. (2017)) is a generative modeling technique which uses four key components: an encoder neural network a decoder neural network d. vector quantisation layer vg. and an autoregressive model m. Given an input the encoder produces an embedding ef e(zz). The quantisation layer comprises of {ci) called the codebook, and quantises the encoder's output embedding of by returning the nearest code Q = ct. along with its index k = arg min; ||c; cf I|- Additionally, in the backwards pass, this quantisation is treated as an identity function, referred to straight-through gradient estimation (Bengio et al. 2013). The decoder produces a reconstruction of the input o; = d(c.). The autoregressive model predicts distribution p(ky ((k/)c<e) mm(c() over the code index at time t using the quantised embeddings Cct of the previous timesteps. The VQ-VAE equations are shown in Equations 3
of M codes
Encoder
ci = e(xx))
Quantisation
Ca k2 vg(c)
(3)
Decoder
â, d(ca)
Model
p(ke (ca) m(cer)
Typically, the encoder, decoder, and codebook trained first and then frozen train the autoregres- sive model in an additional second stage. The total loss for the VQ-VAE is
N
L Equae
+.8 |c0
Gll -
log (hellcee))
(4)
reconstruction
comminment
stage
Stochastic MuZero
this section we present our novel algorithm Stochastic MuZero. Our approach combines learned stochastic transition model of the environment dynamics with variant of Monte Carlo tree search (MCTS). First, we describe the new model and subsequently how it is combined with MCTS for planning.
4.1 STOCHASTIC MODEL
Afterstates We consider the problem of modeling the dynamics of a stochastic environment. Simi- larly to MuZero, the model receives an initial observation ogr at time step I and sequence actions Oua+K. and needs to make predictions about the future values, policies and rewards. In contrast to MuZero which only considers latent states which correspond to real states of the environment, Stochastic MuZero makes use of the notion of afterstates (Sutton & Barto 2018) to capture the stochastic dynamics. An afterstate as, is the hypothetical state of the environment after an action is applied but before the environment has transitioned to true state:
(18;
8t++1
By using afterstates we can separate the effect of applying an action to the environment and of the chance transition given an action. For example in backgammon, the afterstate corresponds to the board state after one player has played its action but before the other player had the chance to roll the dice. It also possible to define the value of an afterstate as V(ast) = Q(+t, az) and the transition probabilities of the environment dynamics Pr(s asz) Pr( 8t, at). An afterstate can lead to multiple states based on a chance event. In our work we assume that there is a finite number of possible states M that the environment can associate each transition with chance outcome ci. An example of a chance outcome could be the result of the dice in game of backgammon. By defining afterstates as, and chance outcomes
transition to, given afterstate, and this way
4

---

Published as a conference paper at ICLR 2022
Ct. model a chance transition using deterministic model 8141, Pt+1 _M(ast, ci) and distribution Pr($t+1 as;) = Pr(ce ast). The task of learning a stochastic model is then reduced to the problem of learning afterstates and chance outcomes
Model The stochastic model of Stochastic MuZero consists of functions: representation function which maps the current observation Oge to a latent state si an afterstate dynamics function ø which given a state s and an action Oc+*k produces the next latent afterstate as a dynamics function which given an afterstate ast and a chance outcome Ct4-k-1 produces the next latent state it k+1 and a reward prediction , prediction function f which given state sf generates the value of and policy P predictions, and a afterstate prediction function & which given an afterstate as generates value prediction Qf and a distribution of = Pr( asf) over possible future chance outcomes C++k+1. The model equations are shown in
Representation
5f h(oge)
Prediction
pe. 1t f(x))
Afterstate Dynamics
as ol a(tk)
(5)
Afterstate Prediction
of O + (as)))
Dynamics
A+I = g( ast co+R+1)
During inference, given an initial observation ose and sequence of actions at:+KK we can generate trajectories from the above model by recurrently unrolling it and by sampling chance outcomes from the distributions C++k+1 off
Chance outcomes Stochastic MuZero models the chance outcomes by using a novel variant of the VQ-VAE method. Specifically, we consider VQ-VAE with constant codebook of size M. Each entry the codebook a fixed one-hot vector of size M. By using a fixed codebook one hot vectors,
simplify the equations of the VQ-VAE In this case, model the encoder embedding of as a categorical variable, and selecting the closest code o is equivalent to computing The resulting encoder can also be viewed as a stochastic function of the observation which makes use of the Gumbel softmax reparameterization trick (Jang et al. 2016) with zero temperature during the forward pass and a straight through estimator during the backward. There no explicit decoder in our model, and contrary to previous work (Ozair et al.) 2021) we do not make use of reconstruction loss. Instead the network is trained end-to-end in fashion similar MuZero. In the following section we explain the training procedure more detail.
the expression one hot (arg (ax (ct
Model training The stochastic model is unrolled and trained an end-to-end fashion similar to MuZero. Specifically, given trajectory of length K with observations Ogta+. actions 0t.++K value targets St.tKK policy targets k and rewards Ut+ 1:t+K the model unrolled for K steps as shown in figure and is trained to optimize the sum of two losses shown in equation 6 MuZero loss and a chance loss for learning the stochastic dynamics of the model.
I, total
M.Mueero
T, chonce
The MuZero loss is the same as the one described in MuZero (see equation 3.1 The chance loss is applied to the predictions O and of which correspond to the latent afterstates as The Q value trained to match the value target z++R and the x² trained towards the one hot chance code C++++1 one hot (arg (e(o +++++++1) )) produced by the encoder. Finally, following the standard VQ-VAE practice, we use a VQ- RQEV commitment cost to ensure that the output of the encoder c e(og++441) is close to the code c++ Equation 7 shows the chance loss used to train
the model.
ボ Σ k=0
K-1
K-1
L chance y
19 (244k: Q.)+ +
Σ
1a (co+841, +
Σ k=0
||-0+1 - E
VQ-VAE commitment cost
5

---

Published as a conference paper at ICLR 2022
A
B
Kitko Zeo1
Ri42, Zeo2
pP, v
of, Qa
pl. Va
p.. V
Dou2
09, Qo - 4p
as
pl
d/, Q' +-
as'
Figure 1: Stochastic MuZero. (A) Monte Carlo Tree Search used in Stochastic MuZero, where diamond nodes represent chance nodes and cireular nodes represent decision nodes. During the selection phase edges are selected by applying the pUCT formula in the case of decision nodes, and by sampling the prior a in the case of chance nodes. (B) Training of stochastic model in Stochastic MuZero. Here for given trajectory of length 2 with observations Ogt:r++2, actions 0e:t422 value targets 2t:t+-2, policy targets 4 and rewards W K. the model is unrolled for steps. During the unroll, the encoder receives the observation ogi+k as input and generates a chance code c++ deterministically. The policy, value and reward outputs of the model are trained towards the targets Te-k, 2++h and Ut+ respectively. The distributions of over future codes trained to predict the code produced by the encoder.
4.2 STOCHASTIC SEARCH
Stochastic MuZero extends the MCTS algorithm used in MuZero by introducing chance nodes and chance values to the search. In the stochastic instantiation of MCTS, there are two types of nodes: decision and chance (Couetoux 2013). The chance and decision nodes are interleaved along the depth of the tree, so that the parent of each decision node is chance node. The root node of the tree always a decision node.
In our approach, each chance node corresponds to latent afterstate querying the stochastic model, where the parent state and an action are provided as an input and the model returns value for the node and prior distribution over future codes Pr(c ass. After chance node is expanded its value is backpropagated up the tree. Finally, when the node is traversed during the selection phase, a code is selected by sampling the prior distribution | In Stochastic MuZero each internal decision node is again expanded by querying the learned model, where the state of the chance parent node and sampled code returns a reward, a value and a policy. Similarly to Mulero the value of the newly added node is backpropagated up the tree, and the pUCT (2) formula is used to select edge. The stochastic search used by Stochastic MuZero is shown schematically in figure]
and it is expanded by
provided
input, and the model
5 EXPERIMENTS
We applied our algorithm to variety of challenging stochastic and deterministic environments. First, we evaluated our approach in the classic game 2048, a stochastic single player game. Subsequently, we considered a two player zero-sum stochastic game, Backgammon, which belongs to the same
In practice we follow the same quasi-random sampling approach Ozair et al. 2021 (A.3), where the Pr(clas)
code is selected using the formula arg max, N[e)+1

---

Published as a conference paper at ICLR 2022
class of board games such Go, chess Shogi where MuZero excels, but with stochasticity induced by the use a die. Finally, we evaluated our method in the deterministic game of Go, to measure any performance loss caused by the use of a stochastic model and search in deterministic environments in comparison to Muleroo
In each environment assess algorithm's ability to learn transition model and effectively it during search. To this end, we compare Stochastic MuZero (using a stochastic learned model) to MuZero (using a deterministic learned model), AlphaZero (using a perfect simulator), and strong baseline method (also using perfect simulator). In the following sections we present our results for each environment separately.
5.1
2048
2048 Training
Scaling in 2048
600k
600k
ApphaZero
MiZero
500k
Stachedtic Muzer
500k
Episode Reward
Episode Score
400k
400k
300k
300k
Jackowski
Jedkowski 2016 2
200k
200k
Jedenwnti 2016
Jedemwnti 2016 5-pry
AlphaZero 91-2000
100k
100k
AlphaZera 924000
Stochesticc NuZero
0KO
0
10
10a
10a
10
Environment frames in billions
Simulations per Move
Figure Planning in 2048. a) Stochastic MuZero, trained using 100 simulations of planning with learned stochastic model, matched the performance of AlphaZero, using 100 simulations of a perfect stochastic simulator, while a deterministic learned model (MuZeero) performed poorly. b) Evaluation of final agent using different levels of search. Stochastic MuZero scales well during evaluation to intermediate levels of search (roughly comparable to 3-ply lookahead), exceeding the playing strength of the state-of- oatte-of-thh-ar baseline (Jaskowski 2016). However, the number of simulations increases we observe diminishing returns due to imperfections of the learned model.
The game of 2048 (inspired by the game of Threes!) is a stochastic, single player, perfect information puzzle game played on 4x4 board. The objective of the game is slide numbered tiles on grid to combine them to create tile with the number 2048; one can continue to play the game after reaching the goal, creating tiles with larger numbers. The episode reward is the sum of all created tile numbers.
There is a plethora of previous work (Szubert & Jaskowski] 2014 Yeh et al. 2017 Oka & Matsuzaki 2016 Rodgers Levine 2014 Neller 2015) on combining reinforcement learning and tree search methods for tackling 2048. Despite its simplicity, model-free approaches have traditionally struggled to achieve high performance, while planning- of the simulator. To date, the best performing agent used the planning-based approach proposed in (Jaskowski] 2016). This method used an expectimax tree search over a perfect simulator, combined with Coomain-speciffi knowledge and number of novel algorithmic ideas that exploited the structure of this specific problem.
approaches have exploited perfect knowledge
In contrast our method uses learned model and no prior knowledge about the environment. Figure 2 compares the performance of Stochastic MuZero in 2048 to AlphaZerro, MuZero and the state-of- the-art Jaskowski 2016 agent. Our method outperformed Jaskowski 2016, while using only a quarter of the training data. Stochastic MuZero also achieved the same performance as AlphaZerro (using perfect simulator), despite learning the model, and performed far better than MuZero (using deterministic model).
5.2 BACKGAMMON
Backgammon is a classic two player, zero-sum, stochastic board game; it was popularized as a standard testbed for reinforcement learning and artificial intelligence by TD-gammon (Tesauro 1995).
7

---

Published as a conference paper at ICLR 2022
Backgammon Training
Backgammon Scaling
0.1
0.0
0.0
Score vs GNUb9
-0.5
Score vs GNUbg
-1.0
-0.1
-155
-02
-2.0
G4Lbg Altaieero
-0.
MiFwro
-2.5
Stechweric NuShem
-0.4
2000
4000
6000
8000
10"
10
102
10a
10a
Train steps in thousands
Simulations per Move
Figure 3: Stochastic MuZero in Backgammon, a) Stochustic MuZero, trained using 1600 simula- tions of planning with a learned stochastic model, matched the performance of AlphaZero, trained using 1600 simulations of perfect stochastic simulator, well matching the superhuman-level program GNUbg Grandmaster. A deterministic learned model (MuZero) performed poorly. b) Stochastic MuZero's 's model scaled well to large searches, and exceeded the playing strength of GNUbg Grandmaster when using more than 10 simulations.
Here we focus on the single game setting, where the final score takes the values + for a simple win or loss, +2 for gammon and +3 for a backgammon.
In all experiments we compared to GNUbg Grandmaster (Free Software Foundation, 2004).
superhuman-level open-source backgammon player. GNUbg combines a learned value function based on handcrafted features with a specialized min-max tree search using a perfect stochastic simulator. GNUbg Grandmaster uses 3-ply look-ahead search over branching factor of 20 legal moves average and 21 chance transitions.
Stochastic
MuZero, using a learned stochastic model the environment and only 1600 simulations per move, achieved the same playing strength as GNUbg, as shown in Figure 5p. The model learned by Stochastic MuZero of high quality: it reached the same playing strength AlphaZero (using a perfect stochastic simulator), and much higher strength than MuZero (using deterministic learned model).
The model also robustly scaled to larger planning budgets (Figure 5p) MuZero improved with increasing number of simulations per move, and ultimately exceeded the playing strength of GNUbg Grandmaster.
the performance of Stochastic
Given the high dimensionality the action space in Backgammon (see appendix for details), our Backgammon experiments used the sample-based search introduced by Hubert et al. (2021).
5.3 Go
Go is classic, two player, perfect information, zero-sum board game, that has been studied heavily in the field of artificial intelligence. AlphaZero and subsequently. MuZero have been the only algorithms which have managed achieve super-human performance, purely through selfplay, this challenging domain. Since the goal of Stochastic MuZero is to extend the applicability of MuZero to stochastic environments while maintaining the latter's 's performance in deterministic environments, we compared the performance of the two algorithms in the game of Go. Figure shows the Elo (Coulom) 2008) achieved by Stochastic MuZero and MuZero during training. Although, Stochastic MuZero requires twice the number of network expansions in comparison to MuZero to achieve the same performance, due to the use stochastic MCTS instead of deterministic one, we ensure that the methods are computationally equivalent by halving the network depth for the chance and dynamic parts of the Stochastic MuZero's is network.
5.4 REPRODUCIBILITY
In order to evaluate the robustness of our method all different environments, we replicated our experiments using nine different initial random seeds (see figure 5.4). We observe that our method is
8

---

Published as a conference paper at ICLR 2022
9x9 Go
19x19
2000
2000
1750
1500
1500
1000
1250
1000
500
750
Muzem
No2em
Muzlers
500
Muzee
100
200
300
400
100
200
300
400
500
600
700
Train steps thousands
Train steps thousands
Figure 4: Stochastic MuZero in Go. Comparison of Stockastic MuZero and MuZero in the game of Go. a) Stochastic MuZero and MuZero when compared 9x9 Go. MuZero has a search budget of 200 simulations during training of 800 during evaluation, while Stochastic MuZero uses 400 simulations during training and 1600 during evaluation. The Elo scale was anchored so that the performance of the final MuZero baseline corresponded an Elo of 2000. b) Stochastic MuZero and MuZero when compared in 19x19 Go. MuZero has a search budget of 400 simulations during training of 800 during evaluation, while Stochastic MuZero uses 800 simulations during training and 1600 during evaluation. The Elo scale was anchored so that the performance of the final MuZero baseline corresponded to an Elo of 2000.
2048
Backgammon
9x9 Go
400000
2000
1750
-0.5
Episode Reward
300000
GNUb9
1500
1250
-1.0
200000
1000
-1.5
750
100000
500
-2.0
250
0
1000
2000 3000
100
300
Environment frames in billions
Train steps in thousands
Train steps in thousands
Figure 5: Stochastic MuZero reproducibility across all domains. We ran our method Stochastic MuZero in all environments using 9 different seeds to measure its robustness to random initialization. We observed that there is minimal variation the performance Stochastic MuZero for all different seeds. Due to the computational cost of each experiment we used smaller number of training steps for each experiment.
robust to the random initialization and there minimal variation in performance between multiple runs. Due to the computational cost of each experiment we used a smaller number of training steps for each experiment.
CONCLUSIONS
In this work, we proposed new method for learning stochastic model of the environment, in fully online reinforcement learning setting, and showed that the learned model can be effectively combined with planning. Our approach builds top of MuZero, model-based reinforcement learning agent that has been widely successful in a range of environments and settings, but its applicabiliity is limited deterministic weakly stochastic environments. We have shown that our algorithm, Stochastic MuZero, can overcome the limitations MuZero, significantly outperforming it in stochastic environments, and it can achieve the same better performance than AlphaiZero

---

Published as a conference paper at ICLR 2022
which makes use of perfect simulator for the environment. Finally, we have demonstrated that Stochastic MuZero matches or exceeds the performance of previous methods that use perfect stochastic simulator, in pure reinforcement learning setting without using any prior knowledge about the environment.
7 REPRODUCIBILITY STATEMENT
In order to ensure the reproducability of our results by the research community, we have included detailed pseudocode, references all environments and datasets used as well detailed description of the hyperparameters used (see Appendix). We did not release the full code as it relies on a lot of proprietary internal infrastructure, limiting usefulness. We also provide study of the robustness of our method under different random initialization conditions (see 5.4).
REFERENCES
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E. Hinton. Layer normalization, 2016.
Yoshua Bengio, Nicholas Léonard, and Aaron Courville, Estimating or propagating gradients through stochastic neurons for conditional computation. arXiv preprint 130.. 2013.
James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milne and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018. URL
http: /t///ggto github........ coom/gogll jax
M. Campbell, J. Hoane, and Hsu. Deep Blue. Artificial Intelligence, 134:57-83, 2002.
Silvia Chiappa, Sébastien Racanière, Daan Wierstra, and Shakir Mohamed. Recurrent environment simulators. 04 2017.
Adrien Couetoux. Monte Carlo Tree Search for Continuous and Stochastic Sequential Decision Making Problems. Theses, Université Paris Sud Paris XI, September 2013. URL https
R. Coulom. Whole-history rating: A Bayesian rating system for players of time-varying strength. In International Conference on Computers and Games, pp. 113-124, 2008
Amir-Massoud Farahmand, Andre Barreto, and Daniel Nikovski. Value-Aware Loss Function for Model-based Reinforcement Learning. In Aarti Singh and Jerry Zhu (eds.), Proceedings of the 20th International Conference on Artificial Intelligence and Statistics, volume 54 Proceedings of
Machine Learning Research, pp. 1486 1494, Fort Lauderdale, FL, USA, 20 20-22 Apr 2017. PMLR.
URL URL/http://proceeiiiii wwwp://sss....r....seeeems.e.eee /aape///ss//aaaaaaa......
mlr
Inc. Free Software Foundation. Gnu backgammon, 2004. URL https: https:///www.ooo org
software/gnuber
Google, 2018. Cloud TPU.
Accessed: 2019.
Christopher Grimm, André Barreto, Satinder Singh, and David Silver. The value equivalence principle for model-based reinforcement learning. CoRR, abs/2011.03506 2020. hhttps
//axxiv.oro/abs/222 03506
Danijar Hafner, Timothy Lillicrap. Mohammad Norouzi, and Jimmy Ba. Mastering atari with discrete world models. In International Conference on Learning Representations, 2021. URL
httts:://ppeneew// wwt.:netttttte..ettt o
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Identity mappings in deep residual networks. CoRR, abs/1603.05007, 2016. htthttp:
arxiv. http:///axxxv.....e.. absg://b//li abwpaabs...a./////a/ 05027
Mikael Henaff, William Whitney, and Yann LeCun. Model-based planning with discrete and continuous actions. arXiv preprint 17X1v: 2017.
10

---

Published as a conference paper at ICLR 2022
Tom Hennigan, Trevor Cai, Tamara Norman, and Igor Babuschkin. Haiku: Sonnet for JAX, 2020.
URL htLhhtt::/ /ttt::///tthb com/ deepminein sereemiid///ee wtmddmmaaaa haiku
Thomas Hubert, Julian Schrittwieser, Ioannis Antonoglou, Mohammadamin Barekatain, Simon Schmitt, and David Silver. Learning and Planning in Complex Action Spaces. arXiv e-prints, April 2021.
Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparameterization with gumbel-softmax, 11 2016.
Wojeiech Jaškowski. Mastering 2048 with delayed temporal coherence learning, multi-state weight promotion, redundant encoding and carousel shaping. IEEE Transactions on Computational Intelligence and Al in Games, PP, 04 2016. doi: 10.1 10.1109/CCIIIG2201..660
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. CoRR, htss/442..980, 2015.
Ilya Loshchilov and Frank Hutter. Fixing weight decay regularization in adam. CoRR, 2017.
Matej Moravčík, Martin Schmid, Neil Burch, Viliam Lisý, Dustin Morrill, Nolan Bard, Trevor Davis, Kevin Waugh, Michael Johanson, and Michael Bowling. Deepstack: Expert-level artificial intelligence in heads-up no-limit poker. Science, 356(6337):508-51) 2017.
Todd W. Neller. Pedagogical possibilities for the 2048 puzzle game. J. Comput. Sci. Coll., 30(3):
38-46, January 2015. ISSN 1937-4771.
Junhyuk Oh, Xiaoxiao Guo, Honglak Lee, Richard L Lewis, and Satinder Singh. Action-conditions video prediction using deep networks in atari games. In C. Cortes, N.D. Lawrence, D.D. Lee, M. Sugiyama, R. Garnett, and R. Gar- nett (eds.), Advances in Neural Information Processing Systems 28, pp. 2845- 2853. Curran Associates, Inc., 2015. URL http::/ www.rrssiissccppppeeee cc/paper/
5859 all--acttoon--ee condicionatttt...cm conditononailionaall -iidvovieo cedee--preiction using deep networks in attar-aaiooo games.
pdf
Junhyuk Oh, Satinder Singh, and Honglak Lee. Value prediction network. In I. Guyon, U. V. Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 30. Curran Asso- ciates, Inc., 2017. URL hattpa://proceecinges pdf
Kazuto Oka and Kiminori Matsuzaki. Systematic selection of n-tuple networks for 2048. In Aske Plaat, Walter Kosters, and Jaap den Herik (eds.), Computers and Games, pp. 81-92, Cham, 2016. Springer International Publishing. ISBN 978-3-319-50935-8.
Sherjil Ozair, Yazhe Li, Ali Razavi, Ioannis Antonoglou, Aaron Van Den Oord, and Oriol Vinyals. Vector quantized models for planning. In Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pp. 8302-8313. PMLR, 18-24 Jul 2021. URL
mlr
Philip Rodgers and John Levine. An investigation into 2048 ai strategies. In 2014 IEEE Conference on Computational Intelligence and Games, pp. 1-2, 2014. doi: /01./////........999
Tom Schaul, John Quan, Ioannis Antonoglou, and David Silver. Prioritized experience replay. In International Conference Learning Representations, Puerto Rico, 2016.
Julian Schrittwieser, Ioannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Simon Schmitt. Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, Timothy P. Lillicrap, and David Silver. Mastering Atari, Go, Chess and Shogi by Planning with a Learned Model. Nature, 588(7839)6604 609, 2020.
11

---

Published as a conference paper at ICLR 2022
Julian Schrittwieser. Thomas Hubert, Amol Mandhane, Mohammadamin Barekatain, loannis Antonoglou, and David Silver. Online and Offline Reinforcement Learning by planning with a learned model. arXiv e-prints, April 2021.
David Silver, Aja Huang, Chris J. Maddison, Arthur Guez, Laurent Sifre, George van den Driessche, Julian Schrittwisser, Ioannis Antonogiou, Veda Panneershelvam, Marc Lanctot, Sander Dieleman, Dominik Grewe, John Nham, Nal Kalchbrenner, Ilya Sutskever, Timothy Lillicrap, Madeleine Leach, Koray Kavukcuoglu, Thore Graepel, and Demis Hassabis. Mastering the game of Go with deep neural networks and tree search. Nature, 529(7587):484 489, January 2016.
David Silver, Hado van Hasselt, Matteo Hessel, Tom Schaul, Arthur Guez, Tim Harley, Gabriel Dulac- Amold, David Reichert, Neil Rabinowitz, Andre Barreto, and Thomas Degris. The predictron: End-to-end learning and planning. In Doina Precup and Yee Whye Tch (eds.). Proceedings of the 34th International Conference Machine Learning, volume 70 of Proceedings of Machine Learning Research, pp. 3191-3199. PMLR, 06-11 Aug 2017. URL mlr. wwtpss/vvossssseeeeeeeeeeep
David Silver, Thomas Hubert, Julian Schrittwieser, Ioannis Antonoglou, Matthew Lai, Arthur Guez, Marc Lanctot, Laurent Sifre, Dharshan Kumaran, Thore Graepel, et al. A general reinforcement learning algorithm that masters chess, shogi, and Go through self-play. Science, 362(6419): 1140-1144, 2018.
Richard S. Sutton. Dyna, an integrated architecture for learning, planning, and reacting, SIGART Bull., 2(4): July 1991. ISSN 0163-5719. doi: 10.1 (0.1145/122344.12231 URL https!
/25//22244422 122377
Richard Sutton and Andrew G. Barto. Reinforccement Learnic An Introduction. The MIT Press. second edition, 2018. URL httt:://cccommeteedees...co net.bokkeeeetoooo.oooooooooo.oo wecthee--oo-eee - html
Marcin Szubert and Wojeiech Jaskowski. Temporal difference learning of n-tuple networks for game 2048. In 2014 IEEE Conference Computational Intelligence and Games, pp. 1-8, 2014. doi: StT.:/II...21..6922220
Gerald Tesauro. Temporal difference learning and td- g-gammonn. Commun ACM, ACM, 38()):58 68, March 1995. ISSN 0001-0782.
doi: 10. 2..145/203330.2033333 wtt.hhtsss/oi..oggg////a
203330. 203343
Aaron van den Oord, Oriol Vinyals, and Koray Kavukcuoglu. Neural discrete representation learn- ing. In I. Guyon, U. V. Luxburg. S. Bengio, H. Wallach, R. Fergus, Vishwanathan, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 30. Curran As- sociates, Inc., 2017. URL httttts::////ceeeettees 7a98af1 7e (e9Sal7eee53e0eeee22662333333333
.eerris httpssssppppppe/./////e//.// pdf
Hado van Hasselt, Matteo Hessel, and John Aslanides. When to parametric models in rein- forcement learning? CoRR, abs/1906.05243. 2019. URL http::tp:///xi..rrrttr/t
1906
05243
Kun-Hao Yeh, I-Chen Wu, Chu-Hsuan Hsueh, Chia-Chuan Chang, Chao-Chin Liang, and Han Chiang. Multistage temporal difference learning for 2048-like like games. JEEE Transactions on Computational Intelligence and Al in Games, 94):369-380, 2017. doi: 10. (0.1100/CCIAGG22216.299777 10.
Lukasz Kaiser, Mohammad Babaeizadeh, Piotr Milos, Blažej Osiński, Roy H Campbell, Konrad Czechowski, Dumitru Erhan, Chelsea Finn, Piotr Kozakowski, Sergey Levine, Afroz Mohiuddin, Ryan Sepassi, George Tucker, and Henryk Michalewski. Model based reinforcement learning for atari. In International Conference on Learning Representations, 2020. URL https:// open ew nwt www.:e/e/ooofiomoomm fto.:/ouummmiddioooo X
12

---

Published as a conference paper at ICLR 2022
A Stochastic MuZero LOSSES
For our Go and Backgammon experiments we used cross entropy loss for the policy and a MSE loss for the value and reward predictions shown below:
(T. -" log
(2, = (z v)
(8)
(u, (u r)²
For the extra chance losses used by Stochastic MuZero, we used a categorical loss to train the probability distribution over the chance codes and MSE loss to train the Q afterstate value. The losses are shown bellow:
1" (c, = lo) ()) Q2
(9)
In our 2048 experiments, similarly to the MulZeo Schrittwieser et al. (2020) approach for general MDPs, used the invertible transform h(x) sign(z)([
cz) with 0.001 for
scaling the value and reward targets and subsequently we obtained categorical representation for each. Subsequently, we trained the value and reward losses using categorical loss as shown below:
[(22, v) = Cat(h(=)) log
("(,, r) = Cal(h(u)) log log
(10)
(= = Cal(h(z))) log Q
B AlphaZero IMPLEMENTATION
In all our experiments compared the performance of Stochastic MuZero against custom Alp- haZero agent. We adapted our AlphaZero implementation to achieve good performance in stochastic and single player environments. Specifically:
We assumed that AlphaZero had access to the true distribution of chance outcomes at each transition.
We used the same conventions of training the policy and value functions of AlphaZero in MuZero and Stochastic MuZero (i.c. n-step bootstrapping, distributional value representation, value scaling etc.).
The network architecture for the representation function used by AlphaZero was the same as in MuZero and Stochastic MuZero.
AlphaZero had access to learned Q value each chance node, similarly to Stochastic MuZero. We found that this significantly improved the performance of the AlphaZero agent.
C 2048 EXPERIMENTS
The game of 2048 is a single player, perfect information, stochastic puzzle game. The board is represented by a 4x4 grid with numbered tiles, and at each step the player has four possible actions which correspond to the four arrow keys (up, down, right, left). When the player selects an action, all the tiles in the board slide in the corresponding direction until they reach the end of the board or another tile of different value. Tiles of the same value are merged together to form a new tile with a value equal to their sum, and the resulting value is added to the running move, a new tile randomly appears in an empty spot on the board with a value of 2 or 4. The game ends when there are more moves available to the player that alter the board state.
of the game. After each
In our 2048 experiments we used a binary representation of the observation as an input to our model. Specifically, the 4 x 4 board was flattened into single vector of size 16, and binary representation
13

---

Published as a conference paper at ICLR 2022
of 31 bits for each number obtained, for a total size of 496 numbers. The representation, afterstate dynamics, dynamics and encoder functions were implemented using a 10 block ResNet v2 style pre-activation residual tower (He et al. 2016) coupled with Layer Normalisation (Ba et al. 2016) and Rectified Linear Unit (Rel.U) activations. Each linear layer has an output size of 256. We used a codebook of size 32 model the stochasticity in the environment.
In both the afterstate dynamics and dynamics network the parent states and codes or actions were com- bined by simple concatenation. The reward and value predictions used the categorical representation introduced in MuZero (Schrittieser et al. 2020). We used 601 bins for both the value and the reward predictions with both the value and the reward being able to represent values between [0, 600.]. Fur- thermore, similarly to MuZero, we used the invertible transform h(z) = sign(x)(Vr + 1-1100 ex) with e = 0.001 for scaling the targets. The value targets were computed using n-step TD(A) bootstrapping with
10, a discount of 0.999 and À = 0.5.
We trained the model using an Adam (Kingma Ba 2015) optimizer and learning rate of 0.0003 for 20M steps with batch size of 1024.
We used a prioritized replay buffer, as proposed in (Schaul et al. 2016) and used in the original MuZero, with priority P(i) = Σ where = ||e; =xll. v is the search value and 2 the observed
n-step return, and set a=1.. We correct for the sampling bias by scaling the total loss using the importance sampling weights w; = ( P()) )², where we set (3=1. Finally, the replay buffer holds the 125000 most recent sequences, where each sequence has a length of up to 200.
Following MuZero, we injected Dirichlet noise the prior policy at the root node of the search, to allow our agent to explore. We used hyperparameters alpha = 0.25 and fraction = 0.1 for the injected noise. Furthermore, the agent selected actions by sampling from the visit count distribution at the root node at the end of each search. We used a temperature scheduler with values [1.0, 0.5, 0.1] for the first [le5, 2e5, 3e5] training steps respectively, and a greedy selection thereafter. We used a budget of 100 simulations for each MCTS search.
D BACKGAMMON EXPERIMENTS
Backgammon is an ancient two player, zero-sum, perfect information, stochastic board game. The board consists of 24 squares (or points) and each player controls 15 checkers, which can move based on the outcome of dice roll. The two players move their checkers in opposite directions and their goal is to move all their checkers off the board first. In addition to a simple winning, a player can also score double ("gammon") when a player bears off all their checkers before their opponent manages to bear off any, while a "backgammon" when the opponent also has checkers left in the player's 's home quadrant (farthermost quadrant from the opponent's ' perspective). Each player can impede the progress of their opponent through "hitting" the opponent' checkers blocking their advancement. A "hit" is achieved when a player's checker advances to position with a single opponent's is checker. Then the opponent's checker needs to reenter the board in the player's home quadrant and no further moves are allowed to the opponent until that happens. A position is blocked to the opponent when it is occupied by at least two of the player's checkers. Each player makes moves based on the values yielded by rolling two dice. In the case of "doubles". aka the two dice have the same value, the player can play up to 4 moves.
triple (" ("backgammon") winning. A "gammon" is achieved
One of the challenges of computer Backgammon is the high branching ratio, since at each ply there are 21 chance outcomes, which yield positions with an average of 20 legal moves each, resulting in a branching ratio of several hundred per ply.
In our backgammon experiments, the board was represented using a vector of size 28, with the first 24 positions representing the number of chips for each player in the 24 possible points the board, and the last four representing the number of hit chips and born off chips for each of the two players. We used positive numbers for the current player's chips and negative ones for her opponent.
An action in our implementation consists of 4 micro-actions, the same as the maximum number of dice player can play at each turn. Each micro-action encodes the source position of a chip along with the value of the die used. We consider 26 possible source positions, with the th position corresponding to no-op, the 1st to retrieving a chip from the hit pile, and the remaining to selecting
14

---

Published as a conference paper at ICLR 2022
chip in one of the 24 possible points. Each micro-action is encoded as single integer with micro-action = src + die.
Similarly to the 2048 experiments, the representation, afterstate dynamics, dynamics and encoder functions were implemented using a 10 block ResNet v2 style pre-activation residual tower (He et al. 2016) coupled with Layer Normalisation (Ba al. 2016) and Rectified Linear Unit (ReL.) activations. Each linear layer has an output size of 256. The action was provided to the afterstate dynamics network vector which was the result of the concatenation of the one-hot representation of each micro-action. We used a codebook of size 32 to model the stochasticity in the environment.
Following the work of (Hubert et al., 202T), used autoregressive prediction head model the network policy, with each step corresponding to single micro-action. To generate a full action, the network was unrolled for 4 steps. In contrast to the 2048 experiments, the value was represented as scalar. Similarly to MuZero when applied to board games, we used Monte Carlo returns to compute the value targets 2,, and we assumed a discount of 1.
We trained the model using an Adam optimizer with weight decay (Loshchilov & Hutter 2017). with learning rate 0.0003 and a weight decay of 0.0001, with batch size of 1024 for a total 8M steps.
In all our experiments we used a replay buffer of 100000 games, and the training trajectories were sampled uniformly.
For exploration, we injected dirichlet noise to the prior policy at the root node. However, since the number of legal moves at each position can dramatically change in backgammon, we dynamically adapted the alpha parameter of the dirichlet noise based on the number of legal moves, with alpha- 1//num legal moves. We used budget of 1600 simulations for each MCTS search.
E Go EXPERIMENTS
In our Go experiments, used the same approach as the proposed in Hubert et al. (2021). The main differences between this setup and the one proposed the original MuZero Schrittwieser et al. (2020) is the use of n-step bootstrapping with target network to improve the data efficiency of the algorithm.
The MuZero and Stochastic MuZero players were evaluated during training by playing 100 matches with a search budget of 800 simulations for MuZero and 1600 for Stochastic MuZero. In order ensure that the methods are computationally equivalent, each of the chance and dynamics networks of Stochastic MuZero has half the depth of the dynamics network used by MuZero. The Elo scale was anchored that the performance of the final MuZero baseline corresponded an Elo of 2000.
F CHANCE ANALYSIS
We investigated the distribution of chance outcomes at each chance node for Stochastic MuZero. We collected a dataset for each game by storing the probability distribution over chance nodes, of ([(e=+===lus]]]) for all afterstate prediction network evaluations invoked throughout all searches in episodes. Subsequently, we sorted each chance node distribution and finally, we computed average distribution, as shown in figure 6 We observed that in the case of deterministic environment like Go, the chance distribution collapsed to a single code, while in stochastic environments the model used multiple codes. Furthermore, in Backgammon, the chance distribution had support of 21 codes with non-negligible probability, which corresponds to the number of distinct rolls of two dice.
G COMPUTATIONAL RESOURCES
All experiments were using second generation Google Cloud TPUs (Google, 2018).
For Backgammon, we used I TPU for training and 16 TPUs for acting, for approximately 27 hours equivalent to 10 days single V100 GPU.
15

---

Published as a conference paper at ICLR 2022
2048
Backgammon
0.200
1.0
0.25
0.175
Average Probability
Average Probability
0.150
Average Probability
0.8
0.20
0.6
0.15
0.100
0.10
0.075
0.4
0.050
0.05
0.2
0.025
0.00
0.000
0.0
0
10
20
30
10
20
30
10
20
30
Index
Index
Index
Figure 6: Average distribution of learned chance outcomes. The average distribution of learned chance outcomes over all chance nodes after running Stochastic MuZero each game for 5 episodes.
In 2048 used TPU for training and 4 TPUs for acting, for 80 hours per experiment; equivalent to roughly 8 days on a V100.
Finally, in Go we used the same setup in MuZero (Schrittwieser et al. 2020).
H IMPLEMENTATION
Stochastic MuZero was implemented described (Schrittwieser et al. 2020). We used the JAX (Bradbury et al. 2018) and Haiku (Hennigan et al. (020)) libraries to implement the neural networks and optimization methods described in this paper. Along with this work, we provide a detailed pseudocode of the Stochastic MuZero algorithm, along with all the hyperparameters used for each environment.
extension to the standard MuZero algorithm, it was
PSEUDOCODE
For completeness provide detailed pseudocode of the Stochastic MuZero algorithm, along with all the hyperparameters used by the agent.
"Pseudccode description the Stochastic MuZero algorithm.
This pscudocede was adapted from the original Mužero pscudocode.
pylint: dlsable-unnsdd-arrummnt
pylint: disable=missing--ooee pylint: import abc
import math
from typing import Any, Dict, Callable, List, NamedTuple, Tuple, Union, Optional, Sequence
import dataclasses import numpy as np
MAXIMM_FFLART
FLOUE_VALUE VALUE = float ('inf')
#* Environment interface АААДААВеее
An action to apply to the environment. It can single integer or 1ist of micro-actions for backgammon. Action Any
16

---

Published as a conference paper at ICLR 2022
The current player play. Player int
class Envientttn:
Immmmmmments the rules of the environment. "**
def asely action: Action):
"** "ppplees action chance outcome to the envirenment. .**
def observation (self):
***teus the observation of the environment to feed to the network.
def ter.inernnlllnmmm (self) -> bool: ***Rcturns true if the envirenment is in torminal stato. return False
def legal_actions (self) Sequence [Action]:
"**eturns the legal actions for the current state. return []
def (eward(sels, player: Player) -> float: **"Returns the last reward for the player."** return 0.0
def (sellay -> Player: ""*Returns the current player to play." """
return
Helpers AUEECOOEEE
class KnownBounds (Named/luple)
min: float
max: float
class MinMaxStats (object):
*** class that holds the min-max values of the tree. "..
def _nii__(se1_____________________ (self, thown_boundde Optional [KnoxnBounds]]:
self = known known_bounds.. max if known_bounds else
MAXIMOX_FI-CC _LCUUN__FFFAA____- FLCAT_VALUEE VALUE
self minimmm known known_bounds wownds.odn..o.o min if known_bounds else
def update (self, value: float):
self.maxime max cnelf.maximmn value) self.minimum min (self.minimum, value)
def (srm,l value: flcat) -> float:
if self. > self.minimum: We normalize only when we have set the maximum and minimum values
return (value - self.minimu / return value
sclf, fal.fminimmm
chance cutcome.
Outcome Any
An object that holds action chance outcome.
17

---

Published as a conference paper at ICLR 2022
ActionOrOutccome Union [Action, Outcome]
LatentState List [float]
AfterState List [float]
class NetworkOutput (NamedTuple):
value: float
probabilities: Dict [AstionorrOutccoee float]
rewara: [ftional[float] 0.0
class Network:
** instance of the network used by stochastic Mužero. "nn
def representation (self, observation) -> latentState:
"""cppresentation function maps from observation to latent stato. return []
def predictions (self, state: LatentState) NetworkOutput:
"**Returns the network predictions for a latent state. "nn return NetworkOutput (0, (h, 0)
def afterstate ditersttate_dnnaiics (self, state: LatentState, action: Action) -> AfterState:
*#*Implements the dynamics from latent state and action to afterstate enn
return (1
def afterstate inteessttteepeeiictoon (self, state: AfterState) -> NetworkOutput: *"*Rcturns the network predictiens for an afterstate, No reward for afterstate transitions.
return Networkoutput (0, (1)
def dynami cs (self, state: AfterState, action: Outcome) -> latentState:
*** *** Imm ementss the dynamics from afterstate and chance cutcome to state. mnn
return 11
def encoder(sel6, observation) -> Cutcome:
***An encoder maps an observation to outcome, "*n
class NetworkCacher:
***An object to share the netwerk between the self-play and training jobs. ann
def iniit_teelltt self. []
def save_network (self, step: int, network: Network): self._networks [step] network
def load netwowkwwokkkel -> Tuple/int, Network]:
training_step (self. new.o..ssttwrrk.kere return training_step, oolf.
Takes the training step and returns the temperature of the softmax policy. Visttsofttmaxeeeeprrruuunn Callable[[int] float]
Returns instance of the envircnment. EnviromeenFaacoory Callaplelll, Environment)
18

---

Published as a conference paper at ICLR 2022
The factory the network.
NetworkFactory Callablell], Network]
cestallasses.dataclle
class $tochaat:1KhuliercCoffii:
# factory for the environment.
saviromeent_factony EnvirommentFactory
network MetworkFactory
Self-play
num_actors: int
visitsiofimav/eepperatueee
num simmlatiiiinn int
discount: float
I Root prior exploration noise. cot__@rrichlt__lphh:: float
root_alrrichle__ffaaciios float root_cirichhee__aaptiiee bool
# UCB formula base: #loat 19652 pb init: :loat 1.25
If we already have some information about which values occur in the environment, we can use them to initialize the rescaling. This is not strictly necessary, but establishes identical behaviour to # AlphaZero board games.
known_bounds: Optional [KnownBounds] None
Replay buffer.
wum_traaeetories_in.buu
int (1e6)
batch_size: int int (128)
nun_unroll_steps: int -
td_steps: int
ta_lambda: float
# Alpha and beta parameters for prioritization.
# By default they are set to 0 which means uniform sampling.
float = 0.0
priority_beta: float 0.0
I Training
seaining_steps: int int (le6)
export_nntwork_eeerry int = int (1e3)
learning_rate: float 3e-4
weight_decay: float 1e-4
# The number of chance codes (codebook size)
# We use codebook of size 32 for all our experiments. codebook_size: int 32
******************************** ## Environment specific configs ##
def twentyfortyyiggttcoofff
-> StochesticMzerooonffi "**Returns the config for the game of 2048. .** def environment_factoryy :):
Returns an implementation of 2048. return Environment ()
def network fettory
19

---

Published as a conference paper at ICLR 2022
10 layer Fully connected Res V2 network with Layer normalization and size 256. return Network
def htt._ssfttma__tmperrtrreeee (train int) flcat: train_steps 1e5:
return 1.0
elit train_steps < 2e5:
return 0.5
elif train_steps 3e5:
return 0.1
else:
i Greedy selection.
return 0.0
return Stcchas:ieouuerocoonfi
neciockk_.cccrrrnnttooo.k.dattre nun_actors=1100,
oat.:.atta....eeettttettttttttte ciscount=0.999,
nun_simulations=1100,
pot__dizichheet.alphav...
peet_dirichee___aattiooo0...
roo__dittchhee__pdppiiveeeaaaa
wun.trafoettrreeeiinnoffferroree (125e3),
td_steps-10,
ta_lambda=0..,
priority._lphe-1.00,
priority_beta=1.0,
training_stops-int (20c6),
batch_size-1122,,
weight_decay=0.0)
def backgammon_conffi -> Stochast wtw.sttoocceeeeooooooooooo "**Returns the config for the game 2048. *** environment. enviromentt_factory ():
# Returns an backgammon. We consider single games without a doubling cube.
return Environment
def network_factory
10 layer fully connected Res V2 network with Layer normalization and size 256. return Network()
def bett__sottaxx_eempretttet (train_steps: int) -> float:
return 1.0
return Stochasti
environment_
network network_factoryy oatrre
num acm.atttos=0000
wat.::..c......aaaaa......aaaaaa mun_simmulatioon=1100,
discount=1.0,
# Unused, we use adaptive dirichlet for backgammon.
root_dirichhlt__shhaa111
rott_diziccheee__feeciioooo
goot_dirichlet_addaa atiriichhet_eadppiieeereue
t Max value is 3 for backgannen.
known Bownodnooooooonnuddd (min=-3, max=3),
20

---

Published as a conference paper at ICLR 2022
le5 full episodes stored. num_trajectrrieeinn..beeeee (1e5), $ We use monte carlo returns.
tö_steps=int (1c3), training_steps-intt (8e6), batch_size-1024, learning_rate=3e-41,
weight_decay=le-
НГКЯНААААИ Replay ІГІІІІэААААдИ
class Searchstats (NamedTuple):
search__polie Dict (Action, int] search_values float
class statc :NamedTuple):
"**aata for single state. """
observation: List [float]
reward: fleat
discount: float
player: Player
action: Action
search_stater SearchStats
Trajectory Sequence [State]
class ReplayBuffer:
"**A replay buffer to hold the experience generated by the selfplay."
init (self, config: Stochsst..MuuerrCCottii)) self.config config
self.data []
def save (self, seq: Trajectory);
self.config, num trajectories buffer: $ Remove the oldest sequence from the buffer. self. pop. self self.dataa.append
if len (self.
def *""Samples trajectory uniformly or using prioritization.
-> Trajectory:
return self.data [0]
semplle_iidexseeeee seg: Trajectory) int:
***Samples an index in the trajectory uniformly or using prioritization,"
return
def sample_elenteee elment (self) -> Trajectory:
*"*ampless single element from the buffer."" Sample trajectory.
trajectory self. trajeteytajjectorr(
state_idx self. t limit = max [seeff.cnfigggeec.....ccccc..cce # Returns trajectory of experiment.
solf.coofig..ddstopp))
return
+ limit]
def sample (self) -> Sequence [Trajectory]:
***Samples a training patch, return [self, le element () for in range
sample_element
21

---

Published as a conference paper at ICLR 2022
A Search
class ActicnoutcomeHisstoyy
""mSimple history container used inside the search.
Only used to keep track of the actions and chance outcomes executed. ***
def init init__it_____ (self, player: Player, history: Optional [List [AstionorOOatcome]]] None)
self.initial_playee = player
self.history = list (history or (])
clone (self)
return actionooutcometiistor (self. batp://ooy.......iiiiiiaaaaaaaa
self.history.
def ad (self,
ActionOrOutccome)
def paat_action_or_uttae (self) -> ActionOrOutccom::
return self. selt.historyl-2
def to play(se1E) -> Player:
Returns the next player to play based on the initial player and the history of actions and outcomes. For example for backgammon the two players alternate, while for 2048 it is always the same player. return
class Noác (object):
"**a Node in the MCTS search tree."" "".
def _nit___t (self, prior: float,
is_chance: bool False)
self.visit_count.
self.to_play -1
self.prior prior
self.value_ = 0
self.children ()
self.state None
self.is_chance is_chance
self.reward
def expanded (self) bool: return Lentself.childrer
def value (sels) -> float:
if self.visit_coont -- 0:
return 0 return self. value_sum self.visit_coum
Core Monte Carlo Tree Search algorithm.
To decide on an action, we run simulations, always starting the root of
the scarch tree and traversing the tree according to the UCB formula until we
reach leaf node.
def run_mcts (config: Stochasti Stochasticcu/eerccooo iMMueeroConfig, root: Node, network:
action_outcome : Network,
min stats: MinMaxStats)
22

---

Published as a conference paper at ICLR 2022
for in cange(confile num simulations) )
history action_outccomm_hstrryy. clone node = root
scarch_path = [node]
while expandeed
action_or_out.come node select_chie chacc_chhlllde (config, node, mri_max_sstati add_actddtnn auction_cr_outcoom
(conff
search_ appeae (node)
Inside the search tree use the dynamics function to obtain the next
hidden state given an acticn and the previous hidden state. parent search p
if parent chance:
The parent is a chance node, afterstate to latent state transition,
# The last action outcome is chance outcome. (parent.state, history.
child_state network_output.
last_act iattion aot.aattoooooottcoooo ()
network_outpur = netwerk_cutpue putput.prediction... (child_ state)
This child is decision nede.
is_child chance False
else:
The parent is decision nede, latent state to afterstate transition.
The last action or outcome is action.
child child_state statteeeee network wwt.kkotttttt...erttteee a dynamics parent. state, network_outpae = network_eutput. # The child is chance node.
ie_child_chance True
Expand the node.
expand_node (node, child_state, network_output, history.to_play(). is_child_chance)
Backpropagate the value up the tree.
backpropagate (search path, network network_output.ver output.value, history.to_plle play (),
config.discount min max sax_st
Select the child with the highest UCB score. def ccllcttchilldconffee
node:
Node,
the min_max_stats: MinMaxStats):
if node, Lode.is_chancer chance we sample from the prier.
outcomes, probs = zip ( (o, n.prob) for o, n in node mode.childee chw.dren. items ()
1)
outcome np. random. choi ce (cutcomes, p=probs)
return cutcome, node. chi ddren [cutcome]
For decision nodes we use the PUCT formula. action, child max (ucb_score (config, node, child, min max stats), action, child)
action, child in node. chi.dren. ()) return action, child
score for node is based on its value, plus exploration bonus based on
the prior.
def ucb_score (config: StochasticMuzern Stochastichhueeroooiiig parent: Node, child: Node,
min max stats: MinMaxStats) -> float:
23

---

Published as a conference paper at ICLR 2022
p math.log (parent visit, count config.pā base config. config.po_c_bar po_c_base) config.po_c_init
/
pb_c + math.sqrt
/ (child. 1)
prior_score pb_c child.prier
if thild.visit_come wtw.ssttttsss.mmllaaa.... nermalize (child. reward value_score min_max__taal stats
contig.discount discount * shild.value ()
else:
value_score
return prior_score + value_score
expand node using the value, reward and policy prediction obtained from the noural network.
def expand_nede (node: Node, state: Union [ntoonssattttttte, AfterState),
network_output: Networkoutput, player: Player, is_chance:
bool):
node.to_play player node. state - state
node. ichence = is_chance
redard network wwwpookkootputtreeaaaae
for acticn, prob network output. probabilities. ()
node.children (actionittinnnniion Node (prob)
the end of simulation, prepagate the evaluation all the way up the tree to the root.
def backpropagate (search_path, List [Node], value: float, to_play: Player,
discount: float, min_max_statss MinMaxStats); for noae (eversc
node. value_sum +- value if node.to_play to_play else -value rode.visit_caume +-
min min_max_stats se..a_ssttts.updteeeeeeeeeeee update (node.value()
value node, reward discount value
At the start of each search, we add dirichlet noise to the prior of the root
encourage the search explore new actions. def
(conff
StochasticMuzerrcoofff node: Nodo)
actions = list wwwlddee..elddd..keee.... dir_alpha config. hottiig.tot_dirichleet.eell
if config. seaffig.root.oee wttpgroot_ddiiiie pac.t__drrichhee_aaattieee
dir_alpha 1.0 np. (Lentaciioo)))
noise
dinncch.ddricchete.... (Ipir_alpha] len (actioons))
frac config. roteereoooolooraiooffffaaaa
for a, (actions noise)
node. chie.children [a] node. childdnnn [a].prior (1 frac) n frac
Solf-play
class Actor (metaclass--ab.
***An actor to interact with the envi ronment.
8abo. abstractnethod def reset (selš):
24

---

Published as a conference paper at ICLR 2022
"*"Resets the player for episode. "**
Babc. labc.abert abstracteethod
def selcct_action (sclf, env: Environment) -> Action:
"""Selects an action for the current state of the environment.
bbbe.abbstaccteeched def stats (self) -> SearchStats:
"""Reuurns the stats for the player after it has selected an action. unn
class
(Actor):
def init _iii__i__t___ (self,
config: StochastichhaferroCnffgg,
cacher: NetworkCacher)
self.config = config
self.cacher cacher
self.training_step
self.network None
def reset (selE):
Read network from the cacher for the episode.
self.training sttep self.network gelf.caches Load network
self.roct None
def _msk L__aak_iillggalatttiao (self,
Environment,
outputs: NetworkOutput) -> NetworkOutput:
"""Masks any actions which are illegal at the root. """ We mask out and keep only the legal actions. masked_policy ()
network_policy outputs.probbbilitiee norm
for action in env. lww.elggacccinnss] if action in network_policy:
masked_policy [action] network_pol icy [action] [acticioon else:
masked masked_policy polpcy laction] 0.0 norm += masked pacicylacticn]
Renormalize masked policy. maskcd_policy = (a: v / norm for a, v in ptek ()
return NetworkOutput (value-outputs..aalue
def _ee_eect_cciionneee (self, root: Nede):
***selects an action given the root n*d
Get the visit count distribution.
actions, visit_counts zip(+[ (action, sode.visit_coume
for action, node in mode.childee. chi.dren. items ()
1)
Temperature
temperature self val.:oooo.fy.aattattttaaaappapee training_stap)
Compute the search policy. search_pelicy [v norm sum sumsseacch_policy) search__polli Iv norm for v in search_policy
temperature) for in visit_counts]
return np. random. choice (acticns, p=search_policyy
25

---

Published as a conference paper at ICLR 2022
def select, select_action (self, env: Environment) Action: "**Selects an action. ***
New min max stats for the search tree, min_max san_max_atttts MinMaxStats (self. config. known snoon_bouudd)
the root of the search tree we use the representation function to
obtain a hidden state given the current observation. root (0)
Provide the history of observations to the representation network to
get the initial latent state.
latent_staatee state self. wel.f.newwor..eeeeeoeeee.. weetrort..erreeeetttiiio........ Compute the predictions.
outputs = www...nttorr.p.cdicticeo (latent_sttttt
Keep only the legal actions.
outputs melf__mak_illleglleettiie (env, outputs)
Expand the root node. expand_noce (root, latent_stare outputs, env.to_play (), is ss_chee False)
Backpropagate the value. backpropagate ( [[roo]], outputs. value, env anv.to_play(), self. wwt.:icoooiicioooooooooo.. add exploration noise to the root node.
min_max_stats)
add_ cet___eppoaatoonniiiiee noise (self.config, root)
We then run Monte Carlo Tree Search using only action sequences and the
model learned by the network.
run_mcts (self.config, root, Actionnut.comett (tory(eev.toopla) ))), self. min_max_stats)
Keep track of the reof root
return the stats.
Return action. return self (rect)
def stats (self) -> SearchStats:
"**Returns the stats of the latest search.
if self.reat is None:
raise ValueError No search was executed,")
return SearchStats(
search_policys
action: vodivvisit.oe for action, node in self. root. chiicchiidennn ittms
b,
search_ valee=sl root row.. value())
Self-play.
Each sclf-play job is independent of all others; it takes the latest network
snapshot, produces an episode and makes it available to the training job by
writing shared replay buffer.
def run_selfplay (config:
Stochasti ciw.aattiiiieeerrnffi... cacher: NetworkCacher,
petter: ReplayBufter):
actor #tocaatticcleeereetttr....... cacher)
26

---

Published as a conference paper at ICLR 2022
while True:
Create a new instance of the environment. env config. htt.gggaeniroommmtfffccco///////
Reset the actor. actor.reset
episode = I]
while net env. env.is is__terminal ():
acticn actor.ssdttt.cciic.err
state State(
dbarraciiooeen..obervvaaaoo)))))
discouts=corrig..iisout,
clayer=eny.to__laay),,
action=action,
scarch_stats=actor.co stats ())
www.applyyacc (action)
Send the spisode to the replay.
replay ay ay__bufeerr.ee wtw.bbuffrr.seeeeeeee. (episode)
F Training
class Learner (metaclass=abb ABCMett):
"**An learner to update the network weights based. ...
9abc.abstran abstractmethod
def learn
""*gingle training step of the learner. "**
labc. rbsttratthcttt def export (self) -> Network: """xxports the network. ***
def loss ""*Minimizes the KL-divergence of the predictions and labels. .*. return 0.0
pol (predictions, labels):
def value_ore maalue_o_rewaad or_reward__eeeeeeeeeeeeee loss (prediction, target) ""*Implmmens the value reward loss for Stochastic Mužero.
For backgammon this is implemented as an MSE loss of scalars.
For 2048, we use the two hot representation proposed in Muzero, and this loss is implemented as divergence between the value
and value target representations.
For 2048 also apply hyperbolic transformation to the target (see paper
for more information).
Args:
prediction: The reward value output of the network. target: The reward or value target.
Returns: The loss to minimize.
27

---

Published as a conference paper at ICLR 2022
return
class stochast.ccMo2errllaaaaaa (Learner):
" **Immlements the learning for Stechastic Mužero. """
def 1__niitttees___e (self,
config: StochasttiMMuuserCCaai...
replay_butfers ReplayBuffer):
self.config config
self.replay_ replay_buffer
Instantiate the network.
self.network = Beadiig.exxor__Facttoe
transpose_ (sel,,, batch)
***raansposs the data so the leading dimension is time instead of
batch.
return batch
def learn (self):
"**Aplles single training step. """ """
batch = self. peeplaa.ouffer.aampe
Transpose batch to make time the leading dimension.
batch self.crangpse_to_ttie (batch)
Compute the initial step loss.
laaent__ttaaa state wwt...eetooo...eiiccccoteee predictions self
www.::..ww....ooooookkookooo.... wtw.::okk..rrreeetttiiinoooooo (batch cbbeervation) (latent_state)
Computes the td target for the 0th position. value_target compute_td (self (self. td_steps,
self.config lambda, batch)
Train the network value towards the td target. tosal_loss value_or_reward_less (predictt value, value_target)
Train the network policy towards the MCTS policy.
+= policy loss (predictions. (oredditioo...p.oobblltt.... probabilities,
batch[0].searchin wtt.::seaachhhhheeeeeee stccchactotttsscccccccoo.....
Unroll the model for k steps.
for t in r1nge self.config.nm num_unroll_steep + 1):
$ Condition the afterstate the previous action. afterstate a.f..nettrrk _dnammics
latent_state, batchlt 11 action)
atterstate_predictionn = self, newwork.ww...... att.: atterstatee afterstate)
Call encoder the next observation. The encoder returns the chance code which is discrete one hot code.
The gradients flow to the encoder using straight through estinator.
chance_code self network. encoder (batch [t] observation)
I The aftorstate value is trained towards the previous value target & but conditioned on the selected action to obtain Q-estimate.
total_lost value_or__eeward_lss. afterstattepreeiiction value, value_target)
The afterstate distribution is trained to predict the chance code generated by the encoder.
28

---

Published as a conference paper at ICLR 2022
total loss policy_lose loss (afterstate p pradd
chance_code)
I Get the dynamic predicticns.
latent_state
dynamics (afterstäte, chance_code)
predictions self..etxerr..pediiite... (latent_state)
# Compute the new value target. value_targes = compute_td_tarpee
telf.contig.t batch[t:]]
The reward loss for the dynamics network. total_loss += value_or_rewwrddloss (predictions. batch[t].
reward)
total loss value car_oor___waae_ee (ps.prrtiiiiioooonosoooooo value_targen
total_ loss += policy_lose (predictions. ptr.:iiioos..roooiiiiiiiieee,
leamiing_rate-ssel...offf
learning_rate,
st__ddeayysel..com
weight_decay)
def export (selÉ) -> Network: return self.netword
def train StochsttiMuueeroCoff..g
cacher: NetworkCacher, replay_buffer: ReplayBuffer): replay_buffer)
learner =
Export the network so the actors can start generating experience.
cacher. catbhersave_neeworre (0,
for step in range trang traiiiing
Single learning step.
pearner.eearn/
if step > and step & config.expor _etworr every == 0: cacher.save www..rstoeeeeeookkeeee. (step, Learner.export. ())
#####****### RL loop #с** def launch_ wtt.__sscccstttcmeeeeeooooooooo (config: Stochastichh/eerconiii.
"**Full RL 1oop for stochastic Mužero. ...
replay_buffer ReplayBuffer (centig)) cacher = MetworkCacher()
launch 1earner job. launch_job (lambda: train_sttchaariiicmaarroooooff cacher, replay_buffer))
# Launch the actors.
for in range (config. acters): launch_job (lambda: run_sclfplay (config, cacher, roffey_buffer)
Stubs to make the typechecker happy. def doftmax_samplae (distribution, temperature: float) return 0, 0
def compute tdd
target (td ste_s, td lambda, trajectory):
29

---

Published as a conference paper at ICLR 2022
"**Computes the TD lambda targets given trajectory for the 0th element.
Args:
td_steps: The number of the n-step returns.
tá_lambda: The lambda in TD (lambda).
trajectory: A sequence of states.
Returns:
The n-step return.
**.
return 0.0
def minimize with sgd (loss, tearning_ratee rate): "**Minimizes the loss using SGD, nnn
def wtt. (loss, learning_rate, weight_decay
}: "*Minnimies the loss using Adam with weight decay."".
def (ounch_jobbEE)
"*"aanchhes job to run remotely." return f()
30


Reviews:
[{"review_id": "yPQZA9JkgnC", "reviewer": "Reviewer_TrM2", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "This work can largely be summarized as plugging the ideas from Ozair 2021 into a MuZero training pipeline. Nevertheless, I think it is a valuable contribution to the community, as it addresses a major shortcoming of MuZero and, unlike Ozair 2021, does so in self-play setting without a reconstruction loss.\n\nThe paper is nicely written (though see small comments below). My main comments are on the experimental side:\n\n1) Throughout the paper, the submission claims that Stochastic MuZero retains the MuZero’s performance in deterministic settings. However, the paper only demonstrates this in a case in which Stochastic MuZero is given twice the simulation budget of MuZero.  I think that making this claim would require showing that Stochastic MuZero matches the performance of MuZero for equal computation budgets. Maybe Stochastic MuZero’s network architecture could be tuned such that each simulation is only half as costly? Seems like this merits further investigation.\n\n2) One weakness of the paper is that the results on games with stochasticity are not on active (at least in a relative sense) research benchmarks. I think it would be interesting to see Stochastic MuZero benchmarked in Hanabi. While Hanabi is an N-player game, it can easily be turned into a POMDP by fixing N-1 players. This was the approach taken by single-agent SPARTA, for which code has been open sourced: https://github.com/facebookresearch/Hanabi_SPARTA. I think showing Stochastic MuZero’s performance against single-agent SPARTA would be a valuable contribution to the community. This experiment would also move toward addressing reviewer 4qYA's interest (which I also share) in seeing Stochastic MuZero in settings with a larger number of possible outcomes.\n\n3) A third, smaller point, regards the baseline AlphaZero implementation. The appendix gives a number of details about the Stochastic MuZero implementation, but few regarding the AlphaZero implementation. I think it would be beneficial to devote some space to this. For example, did the submission's AlphaZero implementation used Monte Carlo returns, as was done in the AlphaZero paper, or did the submission find it necessary to modify this aspect of AlphaZero to achieve good performance in stochastic settings?\n\n4) A fourth, again smaller point concerns the model learned by Stochastic MuZero. In particular, I would be interested to see some analysis on the extent to which Stochastic MuZero's model performs state abstractions. Ie, does Stochastic MuZero's model look more like the true transition model or more like MuZero's (transition determinization) model? Does Stochastic MuZero end up learning a deterministic model in deterministic environment? How do the shape of Stochastic MuZero's search trees look compared to those of AlphaZero and MuZero?\n\n### Small Comments\n\nIt would be more generous to also cite Libratus (not just DeepStack) for tree-based planning algorithms for card games.\n\n“However, learning a model in isolation from its use during planning has proven to be problematic in complex environments”\nWhat about Dreamer v2?\n\nVQ-VAE 3 — make clear this is referring to an equation\n\ncode ct+k+1 = one hot (arg maxi (e(o i ≤t+k+1))) produced \nThe = part of the <= on this line is hard to see because of its proximity to the capital V on the line below.\n\nIt took me a minute to figure out how the encoder was being used. Might be worth discussing it at a greater length in the training section.\n\nIn Figure 1, it took me a minute to figure out that the pictures of the 2048 game were not pictures of a VQ-VAE embedding. Would be better to use backgammon or something else less misinterpretable.", "summary_of_the_paper": "The submission proposes an algorithm (called Stochastic MuZero) that combines VQ-VAEs with MuZero. Unlike MuZero, Stochastic MuZero can handle settings with stochasticity in a principled way (in terms of value equivalence). The submission shows that Stochastic MuZero can perform comparably to AlphaZero in 2048 and Backgammon. Additionally, it shows that Stochastic MuZero can perform comparably to MuZero in Go in a setting in which its computation budget is twice as large.", "summary_of_the_review": "I think that, as a matter of correctness, my 1st point needs to be addressed for the submission to merit acceptance. Simply qualifying the claim would suffice but it would obviously be better if the submission were actually able to substantiate the claim that Stochastic MuZero can match MuZero's performance with an equal computation budget. I also think that the 3rd point should be addressed, toward making the submission's experiments more reproducible. While I think the paper would merit acceptance even without addressing my 2nd and 4th concerns, I will be disappointed if the authors choose not to perform these experiments, as I think they would provide significant value to the community.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "umyLGZY1-89", "reviewer": "Reviewer_qs8f", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Potential Areas of Improvements / Questions:\n- One thing I would have liked to see, possibly in the appendix, is confirmation that this model trained on Go-Zero learned to become fully, or almost-fully deterministic. It’s not essential since it clearly works, but I think it’d be interesting! I suspect it doesn’t completely do this as this is the only explanation I am able to come up for why you need a larger sampling budget for stochastic mu-zero?\n- To the former point, why do you need a larger sampling budget for stochastic mu-zero in the Go games?\n- Would this model also then be applicable to multi-player, imperfect information games? This also does not affect my score nor am I asking for additional experiments on this, I just think such an experiment would be interesting.\n- Is the notion of after-states strictly needed? Why is it advantageous to explicitly split the deterministic and stochastic components instead of having each node be stochastic?\n- It would be good to include the hyperparam grids that were searched over to get the results in the appendix, or if none was done, to say so. The robustness section is nice but it’s valuable to know how much tuning was involved.\n\nClarity:\n- In the “Model” section, I might have missed it, but it looks like l^p is not defined? Same for l^v.\n- Notational question: why is c_{t+k+1} the only thing where the k index is a subscript rather than a superscript?\n- In the section “chance outcomes” there seems to be something wrong with the sentence “By using a fixed codebook of one hot vectors, we can simplify the equations of the VQ-VAE 3”\n- I may have missed it but in case I didn’t, the appendix should include an exact description of the training details of VQ-VAE, including when different components are frozen, the size of the codebook, etc. Unless I missed something, the paper is not reproducible as is.\n- After equation (5), I am not totally clear what is meant by the expression $\\sigma_t^k, ,c_{t+k+1}\\sim \\sigma_t^k$. I assume you mean to say that chance outcomes are drawn from sigma but the construction of the sentence makes it read like you are drawing $\\sigma_t^k$ from $\\sigma_t^k$.\n- I think it would be worthwhile, at least in the appendix, to be clearer about the process by which the chance variables are actually sampled. The description in the “chance outcomes” section “The resulting encoder can also be viewed as a stochastic it \nfunction of the observation which makes use of the Gumbel softmax reparameterization trick (Jang et al., 2016) with zero temperature during the forward pass and a straight through estimator during the backward”  is a little unclear.\n- Given that the VQ-VAE model is a key contribution, I really do think clearer explanation of how the “novel” variant of VQ-VAE works could be given. I don’t have a substantive solution, I just want the authors to be aware that this section was slightly confusing to read. In particular, it might be worth expanding on how the deterministic codes let you still acquire stochasticity.\n- Is $l^\\sigma$ defined anywhere?\n- I think in Fig. 1B you want the sampling symbol rather than the approximately equal symbol?", "summary_of_the_paper": "The authors extend mu-zero to MDPs with stochasticity by adding after-states to the tree and using a VQ-VAE model. They show that this enables them to solve stochastic environments where mu-zero fails.", "summary_of_the_review": "A very good paper with a couple of areas that need some work in rewriting to make the method clearer. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "QoR0yH5X07l", "reviewer": "Reviewer_4qYA", "summary": "", "questions": "", "limitations": "", "rating": 10, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths: \n\n - This paper represents a significant contribution to reinforcement learning, by extending the state-of-the-art MuZero to stochastic or partially observed environments.\n- The paper is clearly written and clearly explains the methods it builds on.\n- The paper is evaluated on both stochastic and non-stochastic environments.\n\nWeakness:\n - I would have liked a discussion of optimality. Expectimax is known to be optimal, but computationally infeasible in many applications. How does this work compare to expectimax and what are the approximations?\n - I would also have liked a discussion of limitations. What are the limitations of the proposed method? For instance, modeling discrete chance outcomes seems to limit this to environments with discrete randomness, e.g. dice rolls, cars, etc. What about environments with continuous stochasticity? Also, how large are the random outcome spaces of the environments tried, and how well does the one-hot chance outcomes scale with larger random outcome spaces? How well does it scale to partially observed environments, like e.g. poker?", "summary_of_the_paper": "The paper proposes an extension of MuZero to stochastic environments.\n\nThe stochasticity of the environment is handled by using afterstates, as_t, and chance outcomes, c_t. This decomposes the modelling of the stochastic environment dynamics into a deterministic model s_{t+1}, r_{t+1} = M(as_t, c_t) and modelling chance outcomes p(c_t | as_t). The chance outcomes are modeled as a discrete categorical variable (1 of M) and learned using a VQ-VAE like setup. \n\nThe paper shows how the proposed model achieves ~SOTA on two stochastic environments: 2048 and backgammon, and retains SOTA performance on a single non-stochastic environment, Go, although, using twice the computational budget.", "summary_of_the_review": "The paper proposes an extension of MuZero to stochastic environments, and achieves ~SOTA results on 2 stochastic environments, 2048 and backgammon, while retaining SOTA performance on a hard non-stochastic environment, Go. The proposed method is a significant step towards a strong, universally applicable, reinforcement learning algorithm.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "7RstLsOEYVw", "reviewer": "Reviewer_1oV1", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "I generally liked the ideas in this paper. Though, I believe there is room for improving the related work and empirical results. \n\n## Related  Work\nThere are two references that were not mentioned as part of the related work, but which are important in the line of value-equivalent MBRL:\n* [Model-Based Reinforcement Learning with Value-Targeted Regression](https://arxiv.org/abs/2006.01107)\n* [Value Iteration Networks](https://arxiv.org/abs/1602.02867)\n\n## Experiments\nTo support their claims, the authors chose two environments that exhibit stochasticity: 2048 and Backgammon, and compared the performance of their algorithm against AlphaZero (where a perfect stochastic simulator is used) and MuZero (where a learned deterministic model is used). A similar methodology was used in an additional experiment, using the deterministic, perfect-information game of Go, and here the question was whether their algorithm's learned stochastic model could match the performance of a learned deterministic model.\n\nIn general, I see no issue with the choice of domains and baselines. My concerns are mostly with the way results are evaluated and reported. For instance, Figures 2, 3, and 4 all seem to represent a single trial. This is unacceptable for experiments with random outcomes, and it considerably weakens support for the paper's main claim of high performance in stochastic settings. I would imagine the authors expect their results to hold on average, given the randomness of a stochastic model. \n\nSome attempt was made to characterize the variability of results in Section 5.4. Each experiment was repeated---under unknown random conditions---three times for a small fraction of the data shown in figures 2, 3, and 4. There are two issues with this methodology:\n1. Three samples is insufficient to characterize a distribution's dispersion.\n2. The trials should be run with the same amount of data as the reported results. \n\nTo approximate the distribution of performance to a reasonable degree of accuracy, I suggest that a minimum of thirty trials are run for the full length of learning. In addition, the authors should describe exactly the sources of randomness their results represent. \n \nIt is also unclear if a hyperparameter search took place for the proposed algorithm. \n\nWhile the authors do reference relevant work that has used their domains, a concise description of the games and their sources of stochasticity is missing. This is important to communicate explicitly to the reader, since it is the primary environmental feature that experiments depend on. This information could be included in the appendix if the paper is tight on space.\n\n## Minor points\n* I didn't find the pseudocode in the appendix very informative. A simple table with all the considered parameters would be sufficient.\n", "summary_of_the_paper": "This paper aims to extend previous work on value-equivalent MBRL, such as MuZero, to stochastic environments. In contrast to conventional work in MBRL that fit transition models to be consistent with environmental observations, this line of work fits transition models to improve the accuracy / utility of a downstream value / policy. To this end authors consider the MuZero algorithm and advocate for learning a stochastic model with a VQ-VAE, and they modify MCTS so that it can be used with their stochastic model.\n\n​The authors propose Stochastic MuZero. Their algorithm makes use of a stochastic model to predict future values, policies and rewards. The authors suggest utilizing \"afterstates\": an imaginary state that is the result of taking an action but it is also before the environment responds with an actual state. In 2048, for example, an afterstate could be the state reached after applying a tile moving action but before a number \"2\" tile appears in a random place.\n\nAs illustrated in Figure 1, the stochastic model consists of 5 functions in contrast to 3 functions in MuZero. The notable addition is in incorporating afterstates in these functions which allows for incorporating chance outcomes. There is an afterstate dynamics function that predicts a latent after-state given a state and action. The typical dynamics function would then still predict a next actual state and reward but its input will be a latent afterstate and a chance outcome. There is also an afterstate prediction function for value and a distribution prediction, where the distribution is that of a chance outcome given an afterstate. That distribution could then be used for sampling chance outcomes in inference.\n\nTo adapt MCTS to this model, search starts from a state and then proceeds to alternate at every level between afterstates and states by using the corresponding dynamics function to reach each type of state. ​ ​", "summary_of_the_review": "Overall, this is an interesting idea, and I encourage the authors to continue working on this paper to provide sound empirical support.  If the claims were sufficiently supported, the results of this paper would be a significant contribution to MBRL. For now, I am leaning towards a rejection as the current empirical methodology is unsound.\n\n### Final \nThank you to the authors, for their response. I have read their replies and the other reviews. I will maintain my original decision to reject, because the biggest issue I have with the paper have not been addressed. This related to the empirical methodology.\n", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper extends MuZero to stochastic (but observable) MDPs. To represent stochastic dynamics, it splits transitions into two parts: a deterministic transition to an afterstate (incorporating all observations and actions up to the current time), followed by a stochastic outcome (accounting for new randomness that follows the last action). The transition to an afterstate is similar in spirit to ordinary MuZero's dynamics model; the stochastic outcome is learned by a VQ-VAE. At planning time, MuZero retains the MCTS lookahead from ordinary MuZero. Stochastic MuZero achieves impressive results: e.g., it maintains the original MuZero's strong performance on the deterministic game of Go, while improving on MuZero significantly (and achieving superhuman performance) on the stochastic game of backgammon.

This is a strong paper overall: it presents a convincing and successful extension of the already-influential MuZero work, along with large-scale computational experiments confirming the utility of the approach. There are nonetheless a few weaknesses: first, compared to the original AlphaZero and MuZero work, it is perhaps less surprising that the given approach is successful, since it is more closely related to prior work. Second, due to the large-scale computational infrastructure needed, it is only possible to run some of the experiments once. This is not in itself a problem, but care needs to be taken in interpreting the results of such single-run experiments: e.g., any figures that show results of single-run experiments should have a clear warning label, and any statements such as "stochastic MuZero performs better than original MuZero" should be tempered with a caveat about how reliable these conclusions are likely to be. Section 5.4 (which runs shorter experiments using three random seeds each) makes a start at evaluating reliability, but (a) the headline results in previous sections do not contain any caveats or pointers to 5.4, and (b) 5.4 should explicitly acknowledge that it cannot hope to detect even quite-common failure cases with so few seeds.
Decision: Accept (Spotlight)