Offline Meta-Reinforcement Learning with Advantage Weighting

1 INTRODUCTION . Meta-reinforcement learning ( meta-RL ) has emerged as a promising strategy for tackling the high sample complexity of reinforcement learning algorithms , when the goal is to ultimately learn many tasks . Meta-RL algorithms exploit shared structure among tasks during meta-training , amortizing the cost of learning across tasks and enabling rapid adaptation to new tasks during meta-testing from only a small amount of experience . Yet unlike in supervised learning , where large amounts of pre-collected data can be pooled from many sources to train a single model , existing meta-RL algorithms assume the ability to collect millions of environment interactions online during meta-training . Developing offline meta-RL methods would enable such methods , in principle , to leverage existing data from any source , making them easier to scale to real-world problems where large amounts of data might be necessary to generalize broadly . To this end , we propose the offline meta-RL problem setting and a corresponding algorithm that uses only offline ( or batch ) experience from a set of training tasks to enable efficient transfer to new tasks without any further interaction with either the training or testing environments . See Figure 1 for a comparison of offline meta-RL and standard meta-RL . Because the offline setting does not allow additional data collection during training , it highlights the desirability of a consistent meta-RL algorithm . A meta-RL algorithm is consistent if , given enough diverse data on the test task , adaptation can find a good policy for the task regardless of the training task distribution . Such an algorithm would provide a ) rapid adaptation to new tasks from the same distribution as the train tasks while b ) allowing for improvement even for out of distribution test tasks . However , designing a consistent meta-RL algorithm in the offline setting is difficult : the consistency requirement suggests we might aim to extend the model-agnostic meta-learning ( MAML ) algorithm ( Finn et al. , 2017a ) , since it directly corresponds to fine-tuning at meta-test time . However , existing MAML approaches use online policy gradients , and only value-based approaches have proven effective in the offline setting . Yet combining MAML with value-based RL subroutines is not straightforward : the higher-order optimization in MAML-like methods demands stable and efficient gradient-descent updates , while TD backups are both somewhat unstable and require a large number of steps to propagate reward information across long time horizons . To address these challenges , one might combine MAML with a supervised , bootstrap-free RL subroutine , such as advantage-weighted regression ( AWR ) ( Peters and Schaal , 2007 ; Peng et al. , 2019 ) , for both for the inner and outer loop of a gradient-based meta-learning algorithm , yielding a ‘ MAML+AWR ’ algorithm . However , as we will discuss in Section 4 and find empirically in Section 5 , naïvely combining MAML and AWR in this way does not provide satisfactory performance because the AWR policy update is not sufficiently expressive . Motivated by prior work that studies the expressive power of MAML ( Finn and Levine , 2018 ) , we increase the expressive power of the meta-learner by introducing a carefully chosen policy update in the inner loop . We theoretically prove that this change increases the richness of the policy ’ s update and find empirically that this policy update dramatically improves adaptation performance and stability in some settings . We further observe that standard feedforward neural network architectures used in reinforcement learning are not well-suited to optimization-based meta-learning and suggest an alternative that proves critical for good performance across four different environments . We call the resulting meta-RL algorithm and architecture Meta-Actor Critic with Advantage Weighting , or MACAW . Our main contributions are the offline meta-RL problem setting itself and MACAW , an offline meta-reinforcement learning algorithm that possesses three key properties : sample efficiency , offline meta-training , and consistency at meta-test time . To our knowledge , MACAW is the first algorithm to successfully combine gradient-based meta-learning and off-policy value-based RL . Our evaluations include experiments on offline variants of standard continuous control meta-RL benchmarks as well as settings specifically designed to test the robustness of an offline meta-learner when training tasks are scarce . In all of these settings , MACAW significantly outperforms fully offline variants state-of-the-art off-policy RL and meta-RL baselines . 2 PRELIMINARIES . In reinforcement learning , an agent interacts with a Markov Decision Process ( MDP ) to maximize its cumulative reward . An MDP is a tuple ( S , A , T , r ) consisting of a state space S , an action space A , stochastic transition dynamics T : S ×A× S → [ 0 , 1 ] , and a reward function r. At each time step , the agent receives reward rt = r ( st , at , st+1 ) . The agent ’ s objective is to maximize the expected return ( i.e . discounted sum of rewards ) R = ∑ t γ trt , where γ ∈ [ 0 , 1 ] is a discount factor . To extend this setting to meta-RL , we consider tasks drawn from a distribution Ti ∼ p ( T ) , where each task Ti = ( S , A , pi , ri ) represents a different MDP . Both the dynamics and reward function may vary across tasks . The tasks are generally assumed to exhibit some ( unknown ) shared structure . During meta-training , the agent is presented with tasks sampled from p ( T ) ; at meta-test time , an agent ’ s objective is to rapidly find a high-performing policy for a ( potentially unseen ) task T ′ ∼ p ( T ) . That is , with only a small amount of experience on T ′ , the agent should find a policy that achieves high expected return on that task . During meta-training , the agent meta-learns parameters or update rules that enables such rapid adaptation at test-time . Model-agnostic meta-learning One class of algorithms for addressing the meta-RL problem ( as well as meta-supervised learning ) are variants of the MAML algorithm ( Finn et al. , 2017a ) , which involves a bi-level optimization that aims to achieve fast adaptation via a few gradient updates . Specifically , MAML optimizes a set of initial policy parameters θ such that a few gradient-descent steps from θ leads to policy parameters that achieve good task performance . At each meta-training step , the inner loop adapts θ to a task T by computing θ′ = θ − α∇θLT ( θ ) , where L is the loss function for task T and α is the step size ( in general , θ′ might be computed from multiple gradient steps , rather than just one as is written here ) . The outer loop updates the initial parameters as θ ← θ − β∇θL′T ( θ′ ) , where L′T is a loss function for task T , which may or may not be the same as the inner-loop loss function LT , and β is the step size . MAML has been previously instantiated with policy gradient updates in the inner and outer loops ( Finn et al. , 2017a ; Rothfuss et al. , 2018 ) , which can only be applied to on-policy meta-RL settings ; we address this shortcoming in this work . Advantage-weighted regression . To develop an offline meta-RL algorithm , we build upon advantage-weighted regression ( AWR ) ( Peng et al. , 2019 ) , a simple offline RL method . The AWR policy objective is given by LAWR ( ϑ , ϕ , B ) = Es , a∼B [ −log πϑ ( a|s ) exp ( 1 T ( RB ( s , a ) − Vϕ ( s ) ) ) ] , ( 1 ) where B = { sj , aj , s′j , rj } can be an arbitrary dataset of transition tuples sampled from some behavior policy , andRB ( s , a ) is the return recorded in the dataset for performing action a in state s , Vϕ ( s ) is the learned value function for the behavior policy evaluated at state s , and T > 0 is a temperature parameter . The termRB ( s , a ) − Vϕ ( s ) represents the advantage of a particular action . The objective can be interpreted as a weighted regression problem , where actions that lead to higher advantages are assigned larger weights . The value function parameters ϕ are typically trained using simple regression onto Monte Carlo returns , and the policy parameters ϑ are trained using LAWR . Next , we discuss the offline meta-RL problem and some of the challenges it poses . 3 THE OFFLINE META-RL PROBLEM . In the offline meta-RL problem setting , we aim to leverage offline multi-task experience to enable fast adaptation to new downstream tasks . Each task Ti is drawn from a task distribution p ( T ) . In the offline setting , the meta-training algorithm is not permitted to directly interact with the meta-training tasks Ti , but instead is provided with a fixed dataset of transition tuples Bi = { si , j , ai , j , s′i , j , ri , j } for each task . Each Bi is populated with trajectories sampled from a corresponding behavior policy µi . Each µi might be an expert policy , sub-optimal demonstrations , other RL agents , or some mixture thereof . Regardless of the behavior policies µi , the objective of offline meta-RL is to maximize return after adaptation on the test tasks . However , depending on the quality of the behavior policies , the maximum attainable return may vary . We observe such a phenomenon in a offline data quality ablation experiment in Section 5 . Sampling data from a fixed dataset at both meta-training and meta-testing time , rather than from the learned policy itself , distinguishes offline meta-RL from the standard meta-RL setting . This constraint is significant , because most algorithms for meta-RL require a large amount of on-policy experience from the environment during meta-training ; these algorithms are generally unable to fully make use of data collected by external sources . During meta-testing , a ( generally unseen ) test task Ttest is drawn from p ( T ) , and the meta-trained agent is presented with a new batch of experience D sampled from a distribution Btest . The agent ’ s objective is to use this batch of data to find the highest-performing policy for the test task . We consider the case where only Bi is fixed during meta-training and Btest corresponds to sampling online trajectories to be the offline meta-RL problem . The case where both Bi and Btest are fixed data buffers is called the fully offline meta-RL problem , which is especially applicable in situations when allowing online exploration might be difficult or dangerous . In the fully offline case , we might also consider the setting where we perform additional online rollouts with our adapted policy and fine-tune with this online data after the initial offline adaptation step . We call this the fully offline meta-RL problem with online fine-tuning . The experiments performed in this paper mostly correspond to the fully offline setting . In Appendix C.2 we also conduct an experiment in the setting of fully offline meta-RL with online fine-tuning . Prior meta-RL methods require interaction with the MDP for each of the meta-training tasks ( Finn et al. , 2017a ) , and though some prior methods build on off-policy RL algorithms ( Rakelly et al. , 2019 ) , these algorithms are known to perform poorly in the fully offline setting ( Levine et al. , 2020 ) . Both of the offline meta-RL settings described above inherit the distributional difficulties of offline RL , which means that addressing this problem setting requires a new type of meta-RL method that is capable of meta-training from offline data . 4 MACAW : META ACTOR-CRITIC WITH ADVANTAGE WEIGHTING Algorithm 1 MACAW Meta-Training 1 : Input : Tasks { Ti } , offline buffers { Di } 2 : Hyperparameters : learning rates α1 , α2 , η1 , η2 , training iterations n , temperature T 3 : Randomly initialize meta-parameters θ , φ 4 : for n steps do 5 : for task Ti ∈ { Ti } do 6 : Sample disjoint batches Dtri , D ts i ∼ Di 7 : φ′i ← φ− η1∇φLV ( φ , Dtri ) 8 : θ′i ← θ − α1∇θLπ ( θ , φ′i , Dtri ) 9 : φ← φ− η2 ∑ i [ ∇φLV ( φ′i , Dtsi ) ] 10 : θ ← θ − α2 ∑ i [ ∇θLAWR ( θ′i , φ′i , Dtsi ) ] Algorithm 2 MACAW Meta-Testing 1 : Input : Test task Tj , offline experience D , meta-policy πθ , meta-value function Vφ 2 : Hyperparameters : learning rates α1 , η , adaptation iterations n , temperature T 3 : Initialize θ0 ← θ , φ0 ← φ . 4 : for n steps do 5 : φt+1 ← φt − η1∇φtLV ( φt , D ) 6 : θt+1 ← θt − α1∇θtLπ ( θt , φt+1 , D ) In addition to satisfying the demands of the offline setting , an ideal method for offline meta-RL should not be limited to the distribution of tasks observed at training time . This is especially important in the offline meta-RL setting , in which the sampling of the training data is out of the control of the agent . In other words , it is critical that an offline meta-RL algorithm be consistent , in the sense that given enough , sufficiently diverse adaptation data at meta-test time , the algorithm can find a good solution to that task , regardless of the meta-training tasks . To address the numerous challenges posed by offline meta-RL , we propose meta actor-critic with advantage weighting ( MACAW ) . MACAW is an offline meta-RL algorithm that learns initializations φ and θ for a value function Vφ and policy πθ , respectively , that can rapidly adapt to a new task seen at meta-test time via gradient descent . Both the value function and the policy objectives correspond to simple regression losses in both the inner and outer loop , leading to a stable and consistent inner-loop adaptation process and outer-loop meta-training signal . While these objectives build upon AWR , we show that the naive application of an AWR update in the inner loop leads to unsatisfactory performance , motivating the enriched policy update that we describe in Section 4.1 . In Sections 4.2 and 4.3 , we detail the full meta-training procedure and an important architectural component of the policy and value networks .

Summary: This paper introduces a new problem setting in meta reinforcement learning, namely metaRL. Here the agent is trained on a fixed offline dataset, which distinguishes it from most metaRL algorithms that interact with the environment during meta training. The authors propose a gradient-based meta learning method (MACAW) to approach this problem, which uses an actor-critic method combined with advantage weighting which is an offline RL method. In experiments they show that this outperforms the offline metaRL method PEARL, and combining multi-task offline RL with AWR in a naive way.