On the role of planning in model-based deep reinforcement learning

Model-based reinforcement learning ( MBRL ) has seen much interest in recent years , with advances yielding impressive gains over model-free methods in data efficiency [ 12 , 15 , 25 , 76 ] , zero- and few-shot learning [ 16 , 37 , 60 ] , and strategic thinking [ 3 , 62 , 63 , 64 , 58 ] . These methods combine planning and learning in a variety of ways , with planning specifically referring to the process of using a learned or given model of the world to construct imagined future trajectories or plans . Some have suggested that models will play a key role in generally intelligent artificial agents [ 14 , 50 , 55 , 56 , 57 , 67 ] , with such arguments often appealing to model-based aspects of human cognition as proof of their importance [ 24 , 26 , 28 , 41 ] . While the recent successes of MBRL methods lend evidence to this hypothesis , there is huge variance in the algorithmic choices made to support such advances . For example , planning can be used to select actions at evaluation time [ e.g. , 12 ] and/or for policy learning [ e.g. , 34 ] ; models can be used within discrete search [ e.g. , 58 ] or gradient-based planning [ e.g. , 25 , 29 ] ; and models can be given [ e.g. , 45 ] or learned [ e.g. , 12 ] . Worryingly , some works even come to contradictory conclusions , such as that long rollouts can hurt performance due to compounding model errors in some settings [ e.g. , 34 ] , while performance continues to increase with search depth in others [ 58 ] . Given the inconsistencies and non-overlapping choices across the literature , it can be hard to get a clear picture of the full MBRL space . This in turn makes it difficult for practitioners to decide which form of MBRL is best for a given problem ( if any ) . The aim of this paper is to assess the strengths and weaknesses of recent advances in MBRL to help clarify the state of the field . We systematically study the role of planning and its algorithmic design choices in a recent state-of-the-art MBRL algorithm , MuZero [ 58 ] . Beyond its strong performance , MuZero ’ s use of multiple canonical MBRL components ( e.g. , search-based planning , a learned model , value estimation , and policy optimization ) make it a good candidate for building intuition about the roles of these components and other methods that use them . Moreover , as discussed in the ∗Correspondence addressed to : { jhamrick , theophane } @ google.com next section , MuZero has direct connections with many other MBRL methods , including Dyna [ 67 ] , MPC [ 11 ] , and policy iteration [ 33 ] . To study the role of planning , we evaluate overall reward obtained by MuZero across a wide range standard MBRL environments : the DeepMind Control Suite [ 70 ] , Atari [ 8 ] , Sokoban [ 51 ] , Minipacman [ 22 ] , and 9x9 Go [ 42 ] . Across these environments , we consider three questions . ( 1 ) For what purposes is planning most useful ? Our results show that planning—which can be used separately for policy improvement , generating the distribution of experience to learn from , and acting at testtime—is most useful in the learning process for computing learning targets and generating data . ( 2 ) What design choices in the search procedure contribute most to the learning process ? We show that deep , precise planning is often unnecessary to achieve high reward in many domains , with two-step planning exhibiting surprisingly strong performance even in Go . ( 3 ) Does planning assist in generalization across variations of the environment—a common motivation for model-based reasoning ? We find that while planning can help make up for small amounts of distribution shift given a good enough model , it is not capable of inducing strong zero-shot generalization on its own . 1 BACKGROUND AND RELATED WORK . Model-based reinforcement learning ( MBRL ) [ 9 , 26 , 47 , 49 , 74 ] involves both learning and planning . For our purposes , learning refers to deep learning of a model , policy , and/or value function . Planning refers to using a learned or given model to construct trajectories or plans . In most MBRL agents , learning and planning interact in complex ways , with better learning usually resulting in better planning , and better planning resulting in better learning . Here , we are interested in understanding how differences in planning affect both the learning process and test-time behavior . MBRL methods can be broadly classified into decisiontime planning , which use the model to select actions , and background planning , which use the model to update a policy [ 68 ] . For example , model-predictive control ( MPC ) [ 11 ] is a classic decision-time planning method that uses the model to optimize a sequence of actions starting from the current environment state . Decisiontime planning methods often feature robustness to uncertainty and fast adaptation to new scenarios [ e.g. , 76 ] , though may be insufficient in settings which require longterm reasoning such as in sparse reward tasks or strategic games like Go . Conversely , Dyna [ 67 ] is a classic background planning method which uses the the model to simulate data on which to train a policy via standard modelfree methods like Q-learning or policy gradient . Background planning methods often feature improved data efficiency over model-free methods [ e.g. , 34 ] , but exhibit the same drawbacks as model-free approaches such as brittleness to out-of-distribution experience at test time . A number of works have adopted hybrid approaches combining both decision-time and background planning . For example , Guo et al . [ 23 ] , Mordatch et al . [ 48 ] distill the results of a decision-time planner into a policy . Silver et al . [ 61 ] , Tesauro & Galperin [ 71 ] do the opposite , allowing a policy to guide the behavior of a decision-time planner . Other works do both , incorporating the distillation or imitation step into the learning loop by allowing the distilled policy from the previous iteration to guide planning on the next iteration , illustrated by the “ update ” arrow in Figure 1 . This results in a form of approximate policy iteration which can be implemented both using single-step [ 40 , 54 ] or multi-step [ 17 , 18 ] updates , the latter of which is also referred to as expert iteration [ 3 ] or dual policy iteration [ 66 ] . Such algorithms have succeeded in board games [ 3 , 4 , 63 , 64 ] , discrete-action MDPs [ 27 , 51 , 58 ] and continuous control [ 43 , 45 , 65 ] . In this paper , we focus our investigation on MuZero [ 58 ] , a state-of-the-art member of the approximate policy iteration family . MuZero is a useful testbed for our analysis not just because of its strong performance , but also because it exhibits important connections to many other works in the MBRL literature . For example , MuZero implements a form of approximate policy iteration [ 18 ] and has close ties to MPC in that it uses MPC as a subroutine for acting . There is also an interesting connection to recent Dyna-style MBRL algorithms via regularized policy optimization : specifically , these methods simulate data from the model and then update the policy on this data using TRPO [ 34 , 39 , 46 , 52 ] . Recent work by Grill et al . [ 20 ] showed that the MuZero policy update approximates TRPO , making the learning algorithm implemented by these Dyna-style algorithms quite similar to that implemented by MuZero . Finally , the use of value gradients [ 10 , 19 , 25 , 29 ] leverages model gradients to compute a policy gradient estimate ; this estimate is similar ( without regularization ) to the estimate that MCTS converges to when used as a policy iteration operator [ 20 , 75 ] . Thus , our results with MuZero have implications not just for its immediate family but to MBRL methods more broadly . In contrast to other work in MBRL , which focuses primarily on data efficiency [ e.g. , 34 , 36 ] or model learning [ e.g. , 12 ] , our primary concern in this paper is in characterizing the role of planning with respect to reward after a large but fixed number of learning steps , as well as to zero-shot generalization performance ( however , we have also found our results to hold when performing the same experiments and measuring approximate regret ) . We note that MuZero can also be used in a more data-efficient manner ( “ MuZero Reanalyze ” ) by using the search to re-compute targets on the same data multiple times [ 58 ] , but leave exploration of its behavior in this regime to future work . Our analysis joins a number of other recent works that seek to better understand the landscape of MBRL methods and the implications of their design choices . For example , Chua et al . [ 12 ] perform a careful analysis of methods for uncertainty quantification in learned models . Other research has investigated the effect of deep versus shallow planning [ 32 , 35 ] , the utility of parametric models in Dyna over replay [ 73 ] , and benchmark performance of a large number of popular MBRL algorithms in continuous control tasks [ 74 ] . Our work is complementary to these prior works and focuses instead on the different ways that planning may be used both during training and at evaluation . 2 PRELIMINARIES : OVERVIEW OF MUZERO . MuZero uses a learned policy and learned value , transition and reward models within Monte-Carlo tree search ( MCTS ) [ 13 , 38 ] both to select actions and to generate targets for policy learning ( Figure 1 ) . We provide a brief overview of MuZero here and refer readers to Appendix A and Schrittwieser et al . [ 58 ] for further details . Algorithm 1 and 2 present pseudocode for MuZero and MCTS , respectively . Model MuZero plans in a hidden state space using a learned model µθ parameterized by θ and comprised of three functions . At timestep t , the encoder embeds past observations into a hidden state , s0t = hθ ( o1 , . . . , ot ) . Given a hidden state and an action in the original action space , the ( deterministic ) recurrent dynamics function predicts rewards and next states , rkθ , t , s k t = gθ ( s k−1 t , a k−1 t ) , where k is the number of imagined steps into the future starting from a real observation at time t. In addition , the prior ( not used in the classic Bayesian sense of the term ) predicts a policy and value for a given hidden state , πkθ , t , v k θ , t = fθ ( s k t ) , and is used to guide the tree search . Search Beginning at the root node s0t , each simulation traverses the search tree according to a search policy until a previously unexplored action is reached . The search policy is a variation on the pUCT rule [ 53 , 38 ] that balances exploitation and exploration , and incorporates the prior to guide this ( see Equation 1 in Section A.3 ) . After selecting an unexplored action a ` at state s ` t , the tree is expanded by adding a new node with reward r ` +1θ , t , state s ` +1 t , policy π ` +1 θ , t , and value v ` +1 θ , t predicted by the model . The value and reward are used to form a bootstrapped estimate of the cumulative discounted reward , which is backed up to the root , updating the estimated return Q and visit count N of each node on the path . After B simulations , MCTS returns a value vMCTSt ( the average cumulative discounted reward at the root ) and policy πMCTSt ( a function of the normalized count of the actions taken at the root during search ) . Acting After search , an action is sampled from the MCTS policy , at ∼ πMCTSt , and is executed in the environment to obtain reward renvt . Data from the search and environment are then added to a replay buffer for use in learning : { ot , at , renvt , πMCTSt , vMCTSt } . Learning The model is jointly trained to predict the reward , policy , and value for each future timestep k = 0 . . .K . The reward target is the observed environment reward , renvt+k . The policy target is the MCTS-constructed policy πMCTSt+k . The value target is the n-step bootstrapped discounted return zt = renvt+1 + γr env t+2 + · · · + γn−1renvt+n + γnvMCTSt+n . For reward , value , and policy losses ` r , ` v , and ` p , respectively , the overall loss is then ` t ( θ ) = ∑K k=0 ` r ( rkθ , t , r env t+k ) + ` v ( vkθ , t , zt+k ) + ` p ( πkθ , t , π MCTS t+k ) . Note that MuZero is not trained to predict future observations or hidden states : the learning signal for the dynamics comes solely from predicting future rewards , values , and policies . We do , however , perform additional experiments in Section D.2 in which we train MuZero to also predict observations via an additional reconstruction loss ; our results indicate that the observationbased model is slightly more accurate but does not qualitatively change MuZero ’ s behavior .

Summary: This paper tries to disentangle the role of planning in model-based reinforcement learning with a number of different ablations and modifications to MuZero. Specifically, the authors analyze the overall contribution of planning by omitting planning from which it is originally used in MuZero, and investigate different planner settings that can drive performance. In addition, they check the generalization advantage of MBRL. Overall, the paper is well-written, and experiments are conducted appropriately. The results provide some insights that other researchers in the MBRL community can leverage for their future work. My major concern is the lack of direct ablation study that can clearly show the advantage of planning in providing a good learning signal. See the detailed comments below.