Consistent Meta-Reinforcement Learning via Model Identification and Experience Relabeling

1 INTRODUCTION . Meta-reinforcement learning algorithms can enable acquisition of new tasks from just a small number of samples by leveraging experience from previous related tasks ( Duan et al. , 2016 ; Wang et al. , 2016 ; Finn et al. , 2017 ) . However , the performance of these methods on new tasks depends crucially on how close the tasks are to the meta-training task distribution . Meta-learned models can adapt quickly to tasks that are similar to those seen during training , but can lose much of their benefit when adapting to tasks that are too far away from the meta-training distribution . This places a significant burden on the user to carefully construct meta-training task distributions that sufficiently cover the kinds of tasks that may be encountered at test time . Many meta-RL methods either utilize a variant of model-agnostic meta-learning ( MAML ) ( Finn et al. , 2017 ; Rothfuss et al. , 2018 ; Nagabandi et al. , 2018 ) , or an inference-based formulation with recurrent ( Duan et al. , 2016 ; Wang et al. , 2016 ) , attentional ( Mishra et al. , 2017 ) , or variational ( Rakelly et al. , 2019 ) methods . The latter class of methods generally fails to handle out-of-distribution tasks , because the model can not adequately deal with out-of-distribution inputs corresponding to the new task . Most of the former class of methods , where gradient-based adaptation corresponds to a well-defined and consistent learning process , require on-policy samples , resulting in high sample complexity for meta-training in order to produce efficient adaptation . In this paper , we make use of a simple insight to develop model identification and experience relabeling ( MIER ) , a meta RL algorithm that is both efficient and which extrapolates effectively when faced with out-of-distribution tasks at meta-test time : we recognize that dynamics and reward models can be adapted consistently , using MAML-style update rules with off-policy data , even if policies and value functions can not . These models can then be used to train new policies for out-of-distribution tasks without using meta-reinforcement learning at all , by generating synthetic experience for the new task . To be able to quickly learn dynamics models , we reformulate the meta-reinforcement learning problem as one of MDP identification . We use this approach since identifying a task involves determining its transition dynamics and reward function , which is exactly the information that a model of the task needs to represent . Specifically , we use a supervised meta-learning method that optimizes for a dynamics model initialization such that conditioned on a MDP context descriptor , prediction error on a validation batch of data sampled from the task is minimized , after updating only the context descriptor with a few gradient steps . Effective model training requires the validation batch to contain data corresponding to optimal behavior for the tasks , which we obtain by training a universal policy conditioned on the context descriptor . Note that since our formulation ensures that the context descriptor contains sufficient information about the task , the policy does not need to be meta-trained or adapted , and can hence be learned with simple and efficient off-policy RL algorithms , without needing to handle the complexity of meta-reinforcement learning . At test time , given out-of-distribution tasks , the adapted context descriptor would indeed be out of distribution and thus our context-conditioned model and policy might not perform well for the test task . However , since our method uses gradient descent which is a consistent learning method for adaptation , we can continue to improve our model using more gradient updates . We then leverage all of the experience collected from other tasks during meta-training , by using the learned model to relabel the next state and reward information , thus obtaining synthetic data to continue training the policy . Our main contribution is an off-policy meta-RL algorithm that is sample efficient , stable and which extrapolates well to out-of-distribution tasks . By formulating the meta-adaptation as MDP identification , we are able to transform the meta-RL problem into a supervised meta-learning problem and thus benefit from the stability and consistency of supervised learning methods . The consistency of our model also enables us to continue improving our policy without collecting extra data by relabeling data collected from other tasks , thus allowing us to efficiently adapt to out-of-distribution tasks . 2 RELATED WORK . Meta-reinforcement learning algorithms extend the framework of meta-learning ( Schmidhuber , 1987 ; Thrun & Pratt , 1998 ; Naik & Mammone , 1992 ) to the reinforcement learning setting . Current metareinforcement learning algorithms can be roughly categorized into 3 large categories , model-free context based methods , model-free gradient based methods and model based methods . Model-free , context based methods such as Duan et al . ( 2016 ) ; Wang et al . ( 2016 ) ; Mishra et al . ( 2017 ) ; Rakelly et al . ( 2019 ) ; Humplik et al . ( 2019 ) often encode the experience during adaptation into a latent context variable , and the policy is conditioned on the task specific context to adapt to a given task . The context encoding process is often done via a recurrent network ( Duan et al. , 2016 ; Wang et al. , 2016 ) , via an attention mechanism ( Mishra et al. , 2017 ) or via amortized probabilistic inference ( Rakelly et al. , 2019 ; Humplik et al. , 2019 ) . While many of these methods can efficiently summarize experiences during adaptation , it is often hard for them to encode any universal learning algorithm such as gradient descent . This is because given a fixed distribution of training tasks , the most efficient way of adapting to a particular task is to directly infer it from the task distribution , rather than applying any universal learning algorithm . While efficient in handling in distribution tasks , the lack of universal learning algorithm in these methods implies that they would be able to adapt to out of distribution tasks . Our method , on the other hand , is able to handle out of distribution tasks through continual adaptation with gradient descent . Model-free gradient based meta-RL methods , such as Finn et al . ( 2017 ) ; Rothfuss et al . ( 2018 ) ; Zintgraf et al . ( 2018 ) ; Rusu et al . ( 2018 ) ; Liu et al . ( 2019 ) , implement gradient descent as the adaptation process . The parameters are often optimized such that the model achieves good performance after only a few steps of gradient descent . Since gradient descent is a universal learning algorithm , with a large capacity model , these methods are consistent in the way that they would guarantee to improve even on out-of-distribution tasks . However , most of these methods are based on on-policy RL algorithm , which means they are rather sample inefficient at training time and requires more data to be collected when adapting to new tasks . It is worth noting that there are also works that combine gradient based and context based methods such as Lan et al . ( 2019 ) . However such methods still suffer from the same sample efficiency problem as other gradient based methods because of the use of on-policy gradient descent . Our method mitigate this problem by applying an off-policy RL algorithm on top of a gradient based meta-learning algorithm to achieve better sample efficiency at training time . At test time , our method reuses the experiences collected during training to enable fast adaptation without collecting a lot more data . Model based meta-RL methods such as Nagabandi et al . ( 2018 ) ; Sæmundsson et al . ( 2018 ) , a rapidly adapting model instead of a policy is learned . Since learning a fast adapting model is a supervised meta-learning problem , a broader range of meta-learning algorithms can be applied . If an universal learning algorithm is used to represent the adaptation process , the model learning can then be made consistent such as in Nagabandi et al . ( 2018 ) . At test time , when the model is adapted to a particular task , standard planning techniques such as model predictive control is often applied to select actions . While possible to be consistent , the model based meta-RL methods often suffer in long horizon tasks due to the accumulated error of model while planning long time ahead into the future . Our method does not suffer from this problem since we only use one step prediction of our model to train a model-free policy , and only execute the policy during test time . This approach is similar to methods presented in Sutton ( 1991 ) ; Janner et al . ( 2019 ) . 3 PRELIMINARIES . 3.1 REINFORCEMENT LEARNING . Formally , a reinforcement learning problem is defined by a Markov decision process ( MDP ) . We adopt the standard definition of a MDP T = ( S , A , p , µ0 , r , γ ) , where S is the state space , A is the action space , p ( s′|s , a ) is the unknown transition probability of landing on next state s′ at the next time step when an agent takes action a at state s , µ0 ( s ) is the initial state distribution , r ( s , a ) is the reward function , and γ ∈ ( 0 , 1 ) is the discount factor . An agent acts according to some policy π ( a|s ) and the objective of learning is to maximize the expected rewards Est , at∼π [ ∑ t γ tr ( st , at ) ] . 3.2 MODEL BASED REINFORCEMENT LEARNING . In model based reinforcement learning , a model p̂ ( s′ , r|s , a ) that predicts the reward and next state from current state and action is trained using standard supervised learning approaches . The model is then used to generate data to train a policy . Specifically , given a model , we perform the following optimization to obtain the policy : arg maxπ Est , at∼π , p̂ [ ∑ t r ( st , at ) ] . Note here the expectation is taken with respect to the model distribution of states instead of the true environment distribution . 3.3 META-REINFORCEMENT LEARNING . In meta-reinforcement learning , we represent various of similar tasks with a distribution of MDPs ρ ( T ) , where each specific task is a sample drawn from the distribution . Given a specific task T , the agent is allowed to collect little amount of data D ( T ) adapt , adapt the policy to form πT according to the data . The objective of this problem is to maximize the expected rewards of the adapted policy ET ∼ρ ( T ) , st , at∼πT [ ∑ t γ tr ( st , at ) ] . 3.4 SUPERVISED META-LEARNING AND MODEL AGNOSTIC META LEARNING . We briefly introduce the supervised meta-learning problem and the model agnostic meta-learning approach , which is an important foundation of our work . In supervised meta-learning , we also have a distribution of tasks ρ ( T ) similar to the meta-RL setup , except that the task T is now a pair of input and output random variables ( XT , YT ) . Given a small dataset D ( T ) adapt sampled from a specific task T , the objective is to built a model that performs well on the evaluation data D ( T ) eval sampled from the same task . If we denote our model as f ( X ; θ ) , the adaptation process as Ξ ( θ , D ( T ) adapt ) and our loss function as L , the objective can be written as : min f , Ξ ET ∼ρ ( T ) [ L ( f ( XT ; Ξ ( θ , D ( T ) adapt ) ) , YT ) ] Model agnostic meta-learning ( Finn et al. , 2017 ) is an approach to solve the supervised meta-learning problem . Specifically , the model f ( X ; θ ) is represented as a neural network , and the adaptation process is represented as few steps of gradient descent . For simplicity of notation , we only write out the one step of gradient descent : ΞMAML ( θ , D ( T ) adapt ) ) = θ − α∇θEX , Y∼D ( T ) adapt [ L ( f ( X ; θ ) , Y ) ] Note that because ΞMAML is the standard gradient descent , it does not contain any trainable parameters and hence the trainable parameters are all contained in the network weights θ . Therefore , the training process of MAML can be summarized as optimizing the loss of the model after few steps of gradient descent on data from the new task .

Summary: The authors propose an algorithm for meta-rl which reduces the problem to one of model identification. The main idea is to meta-train a fast-adapting model of the environment and a shared policy, both conditioned on task-specific context variables. At meta-testing, only the model is adapted using environment data, while the policy simply requires simulated experience. Finally, the authors show experimentally that this procedure better generalizes to out-of-distribution tasks than similar methods.