Online and Offline Reinforcement Learning by Planning with a Learned Model

1 Introduction . Offline reinforcement learning holds the promise of learning useful policies from many existing real-world datasets in a wide range of important problems such as robotics , healthcare or education ( Levine et al. , 2020 ) . Learning effectively from offline data is crucial for such tasks where interaction with the environment is costly or comes with safety concerns , but a large amount of logged and other offline data is often available . A wide variety of effective reinforcement learning ( RL ) algorithms for the online case have been described , achieving impressive results in video games ( Mnih et al. , 2015 ) , robotic control ( Akkaya et al. , 2019 ) and many other problems . However , applying these online RL algorithms to offline data often remains challenging due to off-policy issues , with the best results in offline RL so far obtained by specialised offline algorithms ( Kumar et al. , 2020 ; Wang et al. , 2020 ; Agarwal et al. , 2020 ) . At the same time , model-based reinforcement learning ( RL ) has long focused on learning efficiently from little data , even going as far as learning completely within a model of the environment ( Hafner et al. , 2018 ) - an approach ideally suited for offline RL . ∗Equal contribution 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . So far , these developments have been relatively independent , with no unified algorithm that could achieve state-of-the art results in both the online and offline settings . In this paper , we describe the Reanalyse algorithm , a simple yet effective technique for policy and value improvement at any data budget , including the fully offline case . A preliminary version of Reanalyse was briefly introduced in the context of MuZero ( Schrittwieser et al. , 2020 ) , but limited to data efficiency improvements in the discrete action case . Here , we delve deeper into the algorithm and push its capabilities much further – ultimately to the point where most or all of the data is reanalysed . Starting with the possible uses of Reanalyse , we show how it can be used for data efficient learning and offline RL , leading to MuZero Unplugged . We demonstrate its effectiveness for the online case through results on Atari and for the offline case through results on the RL Unplugged benchmark for Atari and DM Control . 2 Related Work . Recent work by Levine et al . ( 2020 ) provides a thorough review of offline RL literature and presents an excellent introduction to the subject . Much research has focused on regularising the value or policy learning to counteract off-policy issues and learn only from high quality data . Critic-Regularized Regression ( CRR ) uses a critic to filter out bad actions and uses only good actions to train the policy ( Wang et al. , 2020 ) . Random Ensemble Mixture ( REM ) regularises q-value estimation by using random convex combinations of ensemble members during training , and the ensemble mean during evaluation ( Agarwal et al. , 2020 ) . Conservative Q-Learning ( CQL ) learns a conservative Q-function , used to lower bound the value of the current policy ( Kumar et al. , 2020 ) . Pessimistic Offline Policy Optimization ( POPO ) also uses a pessimistic value function for policy learning ( He & Hou , 2021 ) . Existing work has also demonstrated the promise of model-based RL for offline learning ( Matsushima et al. , 2020 ; Argenson & Dulac-Arnold , 2020 ) , but has often been restricted to tasks with lowdimensional action or state spaces , and has not been applied to visually more complex tasks such as Atari ( Bellemare et al. , 2013 ) . Model-Based Offline Reinforcement Learning ( MOReL ) implements a two-step procedure , first learning a pessimistic MDP from offline data using Gaussian dynamics models , then a policy within this learned MDP ( Kidambi et al. , 2020 ) . Results are presented for state-based control tasks . Modelbased Offline Policy Optimization ( MOPO ) penalises rewards by the uncertainty of the model dynamics to avoid distributional shift issues ( Yu et al. , 2020 ) . Offline Reinforcement Learning from Images with Latent Space Models ( LOMPO ) extends MOPO to image based tasks ( Rafailov et al. , 2021 ) . Results are reported on newly introduced datasets with image observations , which the authors aim to open-source in the near future . Deep Averagers with Costs MDP ( DAC-MDP ) ( Shrestha et al. , 2021 ) builds non-parametric models from the offline data , solves these tabular MDPs using value iteration , then generalizes back to the original MDP . Most previous approaches primarily use the learned model for uncertainty estimation and to train a policy ; they do not directly use the learned model for planning over action sequences . In contrast , our method focuses on using the learned model directly for policy and value improvement through planning both offline ( when learning from data ) and online ( when interacting with an environment ) . It requires no regularisation of the value or policy function either in the online or offline case , works well even in very high dimensional state spaces and is equally applicable to both discrete and continuous action spaces . Reanalyse is also qualitatively different from Dyna ( Sutton , 1991 ) in several important regards : it uses both value and policy rather than value function alone ; and it also updates the state representation . In the specific case of MuZero Reanalyse it also performs a tree search rather than a single step lookahead used in Dyna . A combination of MuZero Unplugged with regularisation approaches such as introduced in the previous work discussed above ( Kidambi et al. , 2020 ; Yu et al. , 2020 ; Rafailov et al. , 2021 ) is possible ; we leave such investigations for future work . 3 Reanalyse . Reanalyse takes advantage of model-based value and policy improvement operators to generate new value and policy training targets for a given state ( Algorithm 1 ) . In this work , we will use MuZero ’ s Monte Carlo Tree Search ( MCTS ) planning algorithm combined with its learned model of the environment dynamics as the improvement operator.2 As the learned model and its predictions are updated and improved throughout training , Reanalyse can be repeatedly applied to the same state to generate better and better training targets . The improved training targets in turn are used to improve the model and predictions , leading to a virtuous cycle of improvement . Algorithm 1 The Reanalyse algorithm . MuZero Unplugged instantiates representation , predict , dynamics with the MuZero network architecture ; plan with MCTS ; loss with the MuZero loss in eqn ( 1 ) ; and optimise with Adam . for step← 0 ... N do t ∼ random ( 1 : T ) s0t = representation ( h1 : t , θ ) for i← 0 ... k do πit , ν i t = plan ( representation ( h1 : t+i , θ ) , θ ) pit , v i t = predict ( s i t , θ ) ri+1t , s i+1 t = dynamics ( s i t , at+i , θ ) end for l = loss ( ht : t+k , { r , p , v , u , π , ν } 0 : kt , θ ) ∆θ = optimise ( l , θ ) end for To run MCTS and compute new targets for a training point , the representation function of MuZero maps the history h1 : t of observations , actions and rewards up to timestep t into an agent state or embedding st . The search over possible future action sequences then takes place entirely in this embedding space , by rolling the dynamics forward and applying prediction functions at every step . These predictions output the key quantities required by planning : the policy , value function and reward . The resulting MCTS statistics at the root of the search tree - visit counts for the actions and value estimate averaged over the tree - are then used as new training targets . During reanalysis , no actions a are selected – instead the agent updates its model and prediction parameters based on the data it has already experienced . Specifically , MuZero Reanalyse jointly adjusts its parameters θ to repeatedly optimise the following loss at every time-step t , applied to a model that is unrolled 0 ... K steps into the future , lt ( θ ) = 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 ) where pkt , v k t , and r t k are respectively the policy , value and reward prediction produced by the k-step unrolled model . The respective targets for these predictions are drawn from the corresponding time-step t+ k of the real trajectory : πt+k , νt+k are the improved policy and value generated by the search , zt+k = ut+k+1 + ... + γn−1ut+k+n + γnνt+k+n is an n-step return , and ut+k is the true reward . The policy and value predictions are then updated towards the new training targets , in the same way they would be for targets computed based on environment interactions - through minimising losses lp , lv and lr . In other words , Reanalyse requires no changes on the part of the learner and can be implemented purely in terms of adapting the actors to generate improved targets based on stored data instead of environment interactions . Since the actual MCTS procedure used to Reanalyse a state is the same as the one used to choose an action when interacting with an environment , it is straightforward to perform a mix of both . We refer to this ratio between targets computed from direct interactions with the environment , and targets 2Other model-based algorithms can be used as improvement operators as well . computed by reanalysing existing data points as the Reanalyse fraction . A Reanalyse fraction of 0 % refers to training by only interacting with the environment , no Reanalyse of stored data , whereas a fraction of 100 % refers to the fully offline case with no environment interaction at all . Since Reanalyse only uses stored data points and the learned model to compute improved targets , it can be employed flexibly for many different purposes : • Data Efficiency . The simplest use of Reanalyse is to improve data efficiency by repeatedly computing updated targets on previously collected data throughout training . By scaling the Reanalyse fraction as described in Section 4 , learning can be optimised for any data budget . For this purpose , the data to be reanalysed is sampled from the N most recent environment interactions ; in the limit this includes all interactions throughout training . • Offline RL . When increasing the Reanalyse fraction to 100 % , learning takes place entirely from stored offline data as described in Section 5 , without any interaction with the environment . Offline data may be obtained from a variety of sources , such as other agents , logged data from a heuristic control system or human examples . • Demonstrations . Reanalyse can be used to quickly bootstrap learning from demonstrations containing desirable behaviour that might otherwise be hard to discover - collected for instance from humans - while still interacting with the environment , learning from both sources of data at the same time . This is useful to skip past what might otherwise be hard exploration problems while still improving beyond the quality of the initial demonstrations . • Exploitation of good episodes . When using Reanalyse to improve data efficiency , Reanalyse is applied to the most recently collected data . If instead data is ordered by some other metric , such as episode reward , Reanalyse can be used to quickly learn from rare events , such as rewards observed in hard-exploration tasks . This variant is most useful in deterministic environments , as it could otherwise bias the value estimates in stochastic environments . In this paper , we will focus on the data efficiency and offline RL cases . Remaining cases require no adjustments to the algorithm and only differ in the source of data to be reanalysed . Further combinations of the cases above are also possible , such as a mix of exploitation and data efficiency Reanalyse which we leave for future work . The Reanalyse algorithm has some similarities to experience replay ( Lin , 1992 ) . Whereas replay performs multiple gradient descent updates for the same data point and target , Reanalyse uses modelbased improvement operators to generate multiple training targets for the same data point . Reanalyse and replay have independent effects and can be combined to further improve data efficiency of learning ; in fact we do so for all experiments in this paper .

Summary: This work proposes the Reanalyse algorithm which uses model-based policy and value improvement operators to compute new improved training targets on existing data points, allowing efficient learning for data budgets varying by several orders of magnitude. The Reanalyse can also be used to learn entirely from demonstrations without any environment interactions, as in the case of offline Reinforcement Learning (offline RL). The MuZero Unplugged algorithm provides a single unified algorithm for any data budget, including offline RL.