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: The paper proposes the problem of fully offline meta-RL. Here, the idea is to leverage offline experience from multiple tasks to enable fast adaptation to new tasks. The paper distinguishes two settings of offline meta-RL, one where only the training data is collected offline and testing corresponds to sampling online trajectories, the other where both training and testing data are collected offline. The latter is termed as fully offline meta-RL and is the problem setting considered in this work.