Planning in Stochastic Environments with a Learned Model

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 valueequivalent 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 superhuman performance of standard MuZero in the game of Go . 1 INTRODUCTION . Constructing plans and executing them is an important feature of human and animal behaviour . In the field of artificial intelligence there has been a great amount of research into adding planning capabilities to intelligent agents . Tree-based planning algorithms have shown a lot of success in a wide variety of environments such as card games ( Moravčík et al. , 2017 ) , board games ( Campbell et al. , 2002 ; Silver et al. , 2016 ) and more recently video games ( Schrittwieser et al. , 2020 ) and continuous control tasks ( Hubert et al. , 2021 ) . 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 estimates the dynamics of the environment , with a planning component , using the learned model as a simulator . However , learning a model in isolation from its use during planning has proven to be problematic in complex environments ( van Hasselt et al. , 2019 ) . Instead , value-equivalent modellearning 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 a 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 state-of-the-art performance in large and continuous action spaces ( Hubert et al. , 2021 ) and offline reinforcement learning ( Schrittwieser et al. , 2021 ) . 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 and may be poorly approximated by a deterministic model . Partially observed environments may also be perceived by the agent as stochastic , whenever aliased states can not be disambiguated . Similarly , large and complex environments may appear stochastic to a small agent with finite capacity . In this paper we introduce the first empirically effective approach for handling stochasticity in value equivalent model-learning and planning . The model is factored to first transition deterministically 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 to 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 to MuZero , using a discrete generative network to 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 is 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 . In our experiments , Stochastic MuZero 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 MuZero extends MuZero without sacrificing performance . 2 RELATED WORK . Observation models ( Oh et al. , 2015 ; Chiappa et al. , 2017 ; Łukasz Kaiser et al. , 2020 ) explicitly learn the dynamics of an environment by fitting a model of observations and rewards to observed transitions . Subsequently , these models can be combined with a model-free learning rule in a Dyna fashion ( Sutton , 1991 ) . 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-based methods are unlikely to outperform model-free approaches that use a replay buffer . Latent models ( Schrittwieser 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 actions and is unrolled for k steps . Subsequently , it is trained to make predictions about rewards , values , policies or observations at each timestep based on the current latent state . The model can then be combined with a tree-based planning algorithm or used to generate synthetic trajectories . Recently , MuZero has shown that it is possible to use this approach to achieve state-of-the-art performance in many challenging domains ( Hubert et al. , 2021 ) while using less data ( Schrittwieser et al. , 2021 ) . 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. , 2021 ) the authors propose a recurrent state-space model which consists of three main modules , a recurrent module which generates the deterministic recurrent state ht , a representation model which combines ht with the current observation xt to generate a distribution over stochastic states st and plays the role of the posterior , and a transition predictor which depends only on ht and acts as the prior of the model . By combining the deterministic and stochastic states ht and st the model is trained to predict the current observation ot , the transition reward rt and the discount dt . The next deterministic recurrent state is generated using ht , st and action at . The stochastic states st are modeled as multidimensional multinomial variables . The learned model is then used to generate synthetic data which are used to 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 . In ( Ozair et al. , 2021 ) the authors learn a stochastic transition model using a VQ-VAE generative network ( van den Oord et al. , 2017 ) and subsequently combine it with MCTS . They show that their method can match the performance of MuZero in chess , while viewing the problem as a singleplayer task and implicitly learning to model the behaviour of the opponent . Despite its promise their approach was only applied in a 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 to explicitly predict the observation at every step , which can be a limiting factor in terms of computation and model efficiency when dealing with high dimensional observations . Finally , the authors used a 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 . 3 BACKGROUND 3.1 MuZero MuZero is a model-based general reinforcement learning agent which combines a learned model of the environment dynamics with a Monte Carlo tree search planning algorithm . The model is conditioned on the history of observations o≤t at timestep t and a sequence of future actions at : t+K , and it is trained to predict the search policies πt : t+K , values vπt : t+K and intermediate rewards rt : t+K 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 learned model consists of 3 functions : a representation function h , a dynamics function g and a prediction function f . The representation function maps the current history of observations o≤t into a latent state s0t . The dynamics function g receives the previous latent state s k t and combines it with an action at+k to produce the next latent state sk+1t and the reward r k t . Finally , the prediction function f receives each latent state skt as an input and computes the policy p k t and value vkt . Given a sequence of policy πt : T , value zt : T , and reward ut : T targets , the model is trained to minimize the loss shown in 1 . LMuZero = K∑ k=0 lp ( πt+k , p k t ) + K∑ k=0 lv ( zt+k , v k t ) + K∑ k=1 lr ( ut+k , r k t ) ( 1 ) The policy targets πt+k correspond to the MCTS policy that was generated when searching from observation o≤t+k . The value targets zt+k are computed using n-step returns ( Sutton & Barto , 2018 ) . Finally , the reward targets ut+k correspond to the real instantaneous rewards observed when this sequence was generated . Search MuZero uses a 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 3 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 a which maximizes the upper confidence bound proposed in ( Silver et al. , 2016 ) and shown in equation 2. a = argmax a [ Q ( s , a ) + P ( a | s ) · √∑ bN ( s , b ) 1 +N ( s , a ) ( α1 + log ( ∑ bN ( s , b ) + α2 + 1 α2 ) ) ] ( 2 ) Here , Q ( s , a ) is the value estimate for action a , N ( s , a ) the visit count , P ( a | s ) the prior probability of selecting action a , and α1 , α2 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 a 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 .

Summary: 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.  ​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.  As 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.  To 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. ​ ​