O FFLINE M ETA -R EINFORCEMENT L EARNING WITH O NLINE S ELF -S UPERVISION
A BSTRACT
Meta-reinforcement learning (RL) methods can meta-train policies that adapt to new tasks with orders of magnitude less data than standard RL, but meta-training itself is costly and time-consuming.
If we can meta-train on ofﬂine data, then we can reuse the same static dataset, labeled once with rewards for different tasks, to meta-train policies that adapt to a variety of new tasks at meta-test time.
Although this capability would make meta-RL a practical tool for real-world use, ofﬂine meta-RL presents additional challenges beyond online meta-RL or standard ofﬂine
RL settings.
Meta-RL learns an exploration strategy that collects data for adapting, and also meta-trains a policy that quickly adapts to data from a new task.
Since this policy was meta-trained on a ﬁxed, ofﬂine dataset, it might behave unpredictably when adapting to data collected by the learned exploration strategy, which differs systematically from the ofﬂine data and thus induces distributional shift.
We propose a hybrid ofﬂine meta-RL algorithm, which uses ofﬂine data with rewards to meta-train an adaptive policy, and then collects additional unsupervised online data, without any reward labels to bridge this distribution shift.
By not requiring reward labels for online collection, this data can be much cheaper to collect.
We compare our method to prior work on ofﬂine meta-RL on simulated robot locomotion and manipulation tasks and ﬁnd that using additional unsupervised online data collection leads to a dramatic improvement in the adaptive capabilities of the meta-trained policies, matching the performance of fully online meta-RL on a range of challenging domains that require generalization to new tasks.
I NTRODUCTION
Reinforcement learning (RL) agents are often described as learning from reward and punishment analogously to animals: in the same way that a person might train a dog by providing treats, we might train RL agents by providing rewards.
However, in reality, modern deep RL agents require so many trials to learn a task that providing rewards by hand is often impractical.
Meta-reinforcement learning in principle can mitigate this, by learning to learn using a set of meta-training tasks, and then acquiring new behaviors in just a few trials at meta-test time.
Current meta-RL methods are so efﬁcient that meta-trained policies require only a handful of trajectories (Rakelly et al., 2019), which is reasonable for a human to provide by hand.
However, the meta-training phase in these algorithms still requires a large number of online samples, often even more than standard RL, due to the multi-task nature of the meta-learning problem.
Ofﬂine reinforcement learning methods, which use only prior experience without active data collection, provide a potential solution to this issue, because a user must only annotate multi-task data with rewards once in the ofﬂine dataset, rather than doing so in the inner loop of RL training, and the same ofﬂine multi-task data can be reused repeatedly for many training runs.
While a few recent works have proposed ofﬂine meta-RL algorithms Dorfman & Tamar (2020); Mitchell et al.
(2021), we identify a speciﬁc problem when an agent trained with ofﬂine meta-RL is tested on a new task: the distributional shift between the behavior policy from the ofﬂine data and the meta-test time exploration policy means that adaptation procedures learned from ofﬂine data might not perform well on the (differently distributed) data collected by the exploration policy at meta-test time.
This mismatch in training distribution occurs because ofﬂine meta-RL never trains on data generated by the meta-learned exploration policy.
In practice, we ﬁnd that this mismatch leads to a large degradation in performance when adapting to new tasks.
Moreover, we do not want to remove this distributional shift by simply adopting a conservative exploration strategy, because learning an exploration strategy enables an agent to collect better data for faster adaptation.
We propose to address this challenge by collecting additional online data without any reward supervision, leading to a semi-supervised ofﬂine meta-RL algorithm, as illustrated in Figure 1.
Online data can be relatively cheap to collect when it does not require reward labels, but it can still make it possible to bridge the distributional shift issue.
To make it feasible to use this data for meta-training, we can generate synthetic reward labels for it based on the labeled ofﬂine data.
Based on this principle, we propose semi-supervised meta actor-critic (SMAC), which uses rewardlabeled ofﬂine data to bootstrap a semi-supervised meta-reinforcement learning procedure, in which an ofﬂine meta-RL agent collects additional online experience without any reward labels.
SMAC uses the reward supervision from the ofﬂine dataset to learn to generate new reward functions, which it uses to autonomously annotate rewards in these rewardless interactions and meta-train on this new data. Our method contains two novel contributions:
ﬁrst, it is the ﬁrst method to combine the efﬁcient PEARL (Rakelly et al., 2019) amortized inference method for meta-RL with AWAC (Nair et al., 2020), an effective approach for ofﬂine RL with online ﬁnetuning. This already yields an effective new ofﬂine meta-RL procedure. Second, it is the ﬁrst approach to perform ofﬂine meta-RL with self-supervised online ﬁnetuning, without ground truth rewards. We evaluate our method and prior ofﬂine meta-RL methods on a number of benchmarks (Dorfman & Tamar, 2020; Mitchell et al., 2021), as well as a challenging robot manipulation domain that requires generalization to new tasks, with just a few reward-labeled trials at meta-test time.
We ﬁnd that, while standard meta-RL methods perform well at adapting to training tasks, they suffer from data-distribution shifts when adapting to new tasks.
In contrast, our method attains signiﬁcantly better performance, on par with an online meta-RL method that receives fully labeled online interaction data.
R ELATED W ORKS
Many prior meta-RL algorithms assume that reward labels are provided with each episode of online interaction (Duan et al., 2016; Finn et al., 2017; Gupta et al., 2018b; Xu et al., 2018; Hausman et al., 2018;
Rakelly et al., 2019; Humplik et al., 2019; Kirsch et al., 2019; Zintgraf et al., 2020; Xu et al., 2020; Zhao et al., 2020; Kamienny et al., 2020).
In contrast to these prior methods, our method only requires ofﬂine prior data with rewards, and additional online interaction does not require any ground truth reward signal.
Prior works have also studied other formulations that combine unlabeled and labeled trials.
For example, imitation and inverse reinforcement learning methods use ofﬂine demonstrations to either learn a reward function (Abbeel & Ng, 2004; Finn et al., 2016; Ho &
Ermon, 2016; Fu et al., 2017) or to directly learn a policy (Schaal, 1999; Ross & Bagnell, 2010; Ho &
Ermon, 2016; Reddy et al., 2019; Peng et al., 2020).
Semi-supervised and positive-unlabeled reward learning (Xu & Denil, 2019; Zolna et al., 2020; Konyushkova et al., 2020)
methods use reward labels provided for some interactions to train a reward function for RL.
However, all of these methods have been studied in the context of a single task.
In contrast, we focus on meta-learning an RL procedure that can adapt to new reward functions.
In other words, we do not focus on recovering a single reward function, because there is no single test time reward or task.
SMAC uses a context-based adaptation procedure similar to that proposed by Rakelly et al.
(2019), which is related to contextual policies, such as goal-conditioned reinforcement learning (Kaelbling,
; Schaul et al., 2015; Andrychowicz et al., 2017; Pong et al., 2018; Colas et al., 2018; WardeFarley et al., 2018; Péré et al., 2018; Nair et al., 2018) or successor features (Kulkarni et al., 2016;
Barreto et al., 2017; 2019; Grimm et al., 2019).
In contrast, our meta-learning procedure applies to any RL problem, does not assume that the reward is deﬁned by a single goal state or ﬁxed basis function, and only requires reward labels for static ofﬂine data.
Our method addresses a similar problem to prior ofﬂine meta-RL methods (Mitchell et al., 2021;
Dorfman & Tamar, 2020).
In our comparisons, we ﬁnd that after ofﬂine meta-training, our method is competitive to these prior approaches, but that with additional self-supervised online ﬁne-tuning, our method signiﬁcantly outperforms these methods by mitigating the aforementioned distributional shift issue.
Our method addresses the distribution shift problem by using online interactions without reward supervision.
In our experiments, we found that SMAC greatly improves performance on both training and held-out tasks.
Lastly, SMAC is also related to unsupervised meta-learning methods (Gupta et al., 2018a; Jabri et al., 2019), which annotate data with their own rewards.
In contrast to these methods, we assume that there exists an ofﬂine dataset with reward labels that we can use to learn to generate similar rewards.
P RELIMINARIES
Meta-reinforcement learning.
In meta-RL, we assume there is a distribution of tasks p T ( · ) .
A task T is a Markov decision process (MDP), deﬁned by a tuple T =
( S , A , r, γ, p 0 , p d ) , where S is the state space,
A is the action space, r is a reward function, γ is a discount factor, p 0 ( s 0 ) is the initial state distribution, and p d ( s t +1 | s t , a t ) is the environment dynamics distribution.
A replay buffer D is a set of state, action, reward, next-states tuples, D
= { s i , a i , r i , s (cid:48)
i } N size i =1 , where all the rewards come from the same task.
We will use the letter h to denote a mini-batch or “history” and the notation h ∼ D to denote that a mini-batch h is sampled from a replay buffer D .
We will use the letter τ to represent a trajectory τ =
( s 1 , a 1 , s 2 , . . . )
without reward labels.
A meta-episode consists of sampling a task T ∼ p T ( · ) , collecting T trajectories with a policy π θ with parameters θ , adapting the policy to the task between trajectories, and measuring the performance on the last trajectory.
Between trajectories, the adaptation procedure transforms the history of states and actions h from the current meta-episode into a context z
=
A φ ( h ) , which is then given to the policy π θ ( a , | s , z ) for adaptation.
The exact representation of π θ , A φ , and z depends on the speciﬁc meta-RL method used.
For example, the context z can be weights of a neural network (Finn et al., 2017) outputted by a gradient update, hidden activations ouputted by a recurrent neural network (Duan et al., 2016), or latent variables outputted by a stochastic encoder (Rakelly et al., 2019).
Using this notation, the objective in meta-RL is to learn the adaptation parameters φ and policy parameters θ to maximize performance on a meta-episode given a new task T sampled from p ( T ) .
PEARL.
Since we require an off-policy meta-RL procedure for ofﬂine meta-training, we build on probabilistic embeddings for actor-critic RL (PEARL) (Rakelly et al., 2019), an online off-policy meta-RL algorithm.
In PEARL, z is a vector and the adaptation procedure A φ consists of sampling z from a distribution z ∼ q φ e ( z | h ) .
The distribution q φ e is generated by an encoder with parameters φ e .
This encoder is a set-based neural network that processes the tuples in h
= { s i , a i , r i , s (cid:48) i } N enc
i
=1 in a permutation-invariant manner to produce the mean and variance of a diagonal multivariate Gaussian.
The policy is a contextual policy π θ ( a | s , z ) conditioned on z by concatenating z to the state s .
The policy parameter θ is trained using soft-actor critic (Haarnoja et al., 2018) which involves learning a Q -function, Q w ( s , a , z ) , with parameter w that estimates the sum of future discounted rewards conditioned on the current state, action, and context.
The encoder parameters are trained by back-propagating the critic loss into the encoder.
The actor, critic, and encoder losses are minimized via gradient descent with mini-batches sampled from separate replay buffers for each task.
Ofﬂine reinforcement learning.
In ofﬂine reinforcement learning, we assume that we have access to a dataset D collected by some behavior policy π β .
An RL agent must train on this ﬁxed dataset and cannot interact with the environment.
One challenge that ofﬂine RL poses is that the distribution of states and actions that an agent will see when deployed will likely be different from those seen in the ofﬂine dataset as they are generated by the agent, and a number of recent methods have tackled this distribution shift issue (Fujimoto et al., 2019b;a; Kumar et al., 2019;
Wu et al., 2019; Nair et al., 2020; Levine et al., 2020).
Moreover, one can combine ofﬂine RL with meta-RL by training meta-RL
[Figure]
[Figure]
|
[Figure]
) |
[Figure]
P o s tA
d a p t a t i o n
R e t u r n s
Ant Direction: Adapting from Online vs Oﬄine Data
[Figure]
on multiple datasets D 1 , . . .
, D N buff (Dorfman & Tamar, 2020; Mitchell et al., 2021), but in the next section we describe some limitations of this combination.
T HE P ROBLEM WITH N AÏVE O FFLINE M ETA -R EINFORCEMENT L EARNING
Ofﬂine meta-RL is the composition of meta-RL and ofﬂine RL: the objective is to maximize the standard meta-RL objective using only a ﬁxed set of replay buffers, D
=
{D i } N buff i
=1 , where each buffer corresponds to data for one task.
Ofﬂine meta-RL methods can in principle utilize the same constraint-based approaches that standard ofﬂine RL algorithms have used to mitigate distributional shift.
However, they must also contend with an additional distribution shift challenge that is speciﬁc to the meta-RL scenario: distribution shift in z -space.
Distribution shift in z -space occurs because meta-learning requires learning an exploration policy π θ that generates data for adaptation.
However, ofﬂine meta-RL only trains the adaptation procedure
A φ ( h ) using ofﬂine data generated by a previous behavior policy, which we denote as π β .
After ofﬂine training, there will be a mismatch between this learned exploration policy π θ and the behavior policy π β , leading to a difference in the history h and in turn, in the context variables z
= A φ ( h ) .
For example, in a robot manipulation setting, the ofﬂine dataset may contain smooth trajectories that were collected by a human teleoperator (Kofman et al., 2005). In contrast, the learned exploration may contain jittering due learning artifacts or the use of a stochastic policy. This jittering may not impede the robot from exploring the environment, but may result in a trajectory distribution shift that degrades the adaptation process, which only learned to adapt to smooth, human-generated trajectories in the ofﬂine dataset. More formally, if p ( z | h ofﬂine ) and p ( z | h online ) denote the marginal distribution of z given histories h sampled using ofﬂine and online data, respectively, the differences between π θ and π β will lead to differences between p ( z | h ofﬂine ) during ofﬂine training and p ( z | h online ) at meta-test time.
To illustrate this difference, we compare p ( z | h ofﬂine ) and p ( z | h online ) on the Ant Direction task (see Section 6).
We approximate these distributions by using the PEARL encoder discussed in Section 3, with p ( z | h ) ≈
q φ e ( z | h ) , where h is sampled either from the ofﬂine data set or using the learned policy.
We measure the KL-divergence observed at the end of ofﬂine training between the posterior p ( z | h ) and a ﬁxed prior p z ( z ) , for different samples of h .
If the two distributions were the same, then we would expect the distribution of KL divergences to also be similar.
However, Figureshows that the two distributions are markedly different after the ofﬂine training phase of SMAC.
We also observe that this distribution shift negatively impacts the resulting policy.
In Figure 2, we plot the performance of the learned policy when conditioned on z sampled from q φ e ( z | h ofﬂine ) compared to q φ e ( z | h online ) .
We see that the policy that uses ofﬂine data, h ofﬂine , leads to improvement, while the same policy that uses data from the learned policy, h online drops in performance.
Since we evaluate the same policy π θ and only change how z is sampled, this degradation in performance suggests that the policy suffers from distributional shift between p ( z | h ofﬂine ) and p ( z | h online ) .
In other words, the encoder produces z vectors that are too unfamiliar to the policy when conditioned on these exploration trajectories.
We note that this issue arises in any method that trains non-Markovian policies with ofﬂine data.
For example, recurrent policies for partially observed MDPs (Jaakkola et al., 1995) depend both on the current observation o and a history h .
When deployed, these policies must also contend with potential distributional shifts between the training and test-time history distributions, in addition to the change in observation distribution o .
This additional distribution shift may explain why many memory-based recurrent policies are often trained online (Duan et al., 2016; Heess et al., 2017; Espeholt et al., 2018) or have beneﬁted from refreshing the memory states (Kapturowski et al., 2018).
In this paper, we focus on addressing this issue speciﬁcally in the ofﬂine meta-RL setting.
Ofﬂine meta-RL with self-supervised online training.
In complex environments where many behaviors are possible, the distribution shift in z -space will likely be inevitable, since the learned policy is likely to deviate from the behavior policy.
To address this issue, we introduce an additional assumption: in addition to the ofﬂine dataset, we assume that the agent can autonomously interact with the environment without observing additional reward supervision .
This problem statement is useful for scenarios where autonomously interacting with the world is relatively easy, but online reward supervision is more expensive to obtain.
Formally, we assume that the agent can generate additional rollouts in an MDP without a reward function, T \ r = ( S , A , γ, p 0 , p d ) .
These additional interactions enable the agent to explore using the learned exploration policy.
Since the resulting states and actions are from the learned exploration policy, we can use them to construct h online and enable an agent to train on the online context distribution, p ( z | h online ) , thus mitigating the distribution shift issue described above.
However, meta-training requires that the history h online contain not just states and actions, but also rewards.
In the next section, we describe a method for autonomously labeling these rollouts with synthetic reward labels to enable an agent to meta-train on this additional data.
S EMI -S UPERVISED M ETA A CTOR -C RITIC
In this section, we present our method, semi-supervised meta actor-critic (SMAC). SMAC consists of ofﬂine meta-training followed by self-supervised online meta-training to mitigate the distribution shift in z -space.
The SMAC adaptation procedure consists of passing history through the encoder described in Section 3, resulting in a posterior q φ e ( z | h ) .
Below, we describe both phases.
5. O FFLINE M ETA -T RAINING
To learn from the user-provided ofﬂine data, we adapt the PEARL (Rakelly et al., 2019) to the ofﬂine setting.
Similar to PEARL, we update the critic by minimizing the Bellman error:
[Equation]
(cid:48) , a (cid:48) , z )))
2 (cid:3) , (1) where ¯ w are target network weights updated using soft target updates (Lillicrap et al., 2016).
PEARL has a policy updated based on soft actor critic (SAC) (Haarnoja et al., 2018), but when naïvely applied to the ofﬂine setting, these policy updates suffer from off-policy bootstrapping error accumulation (Kumar et al., 2019), which occurs when the target Q-function for bootstrapping
Q ( s (cid:48) , a (cid:48) ) is evaluated at actions a (cid:48) outside of the training data.
To avoid error accumulation during ofﬂine training, we implicitly constrain the policy to stay close to the actions observed in the replay buffer, following the approach in a single-task ofﬂine RL algorithm called AWAC (Nair et al., 2020).
AWAC uses the following loss to approximate a constrained optimization problem, where the policy is constrained to stay close to the data observed in D :
[Equation]
We estimate the value function V ( s , z ) =
E a ∼ π [Equation] a | s , z ) Q ( s , a , z ) with a single sample, and λ is the resulting Lagrange multiplier for the optimization problem.
See Nair et al.
(2020) for a full derivation.
With this modiﬁed actor update, we can train the encoder, actor, and critic on the ofﬂine data without the overestimation issues that afﬂict conventional actor-critic algorithms (Kumar et al., 2019).
However, it does not address the z -space distributional shift issue discussed in Section 4, because the exploration policy learned via this ofﬂine procedure will still deviate signiﬁcantly from the behavior policy π β .
As discussed previously, we will aim to address this issue by collecting additional online data without reward labels and learning to generate reward labels if self-supervised meta-training.
Learning to generate rewards.
To continue meta-training online without ground truth reward labels, we propose to use the ofﬂine dataset to learn a generative model over meta-training task reward functions that we can use to label the transitions collected online.
Recall that during ofﬂine learning, we learn an encoder q φ e that maps experience h to a latent context z that encodes the task.
In the same way that we train our policy π θ ( a | s , z ) that conditionally decodes z into actions, we additionally train a reward decoder r φ d ( s , a , z ) with parameter φ d 1 that conditionally decodes z into rewards.
We train the reward decoder r φ d to reconstruct the observed reward in the ofﬂine dataset through a mean squared error loss.
Because we use the latent space z for reward-decoding, we back-propagate the reward decoder loss into q φ e .
As visualized in Figure 3, we also regularize the posteriors q φ e ( z | h ) against a prior p z ( z ) to provide an information bottleneck in that latent space z and ensure that samples from p z ( z ) represent meaningful latent variables.
We found it beneﬁcial to not back-propagate the critic loss into the encoder, in contrast to prior work such as PEARL.
To summarize, we train the encoder and reward decoder by minimizing the following loss, in which we assume that z ∼ q φ e ( h ) : [Equation]
[Equation]
5. S ELF -S UPERVISED O NLINE M ETA -T RAINING
We now describe the self-supervised online training procedure, during which we use the reward decoder to provide supervision.
First, we collect a trajectory τ by rolling out our exploration policy π θ conditioned on a context sampled from the prior p ( z ) .
To emulate the ofﬂine meta-training supervision, we would like to label τ with rewards that are in the distribution of meta-training tasks.
As such, we sample a replay buffer D i uniformly from D to get a history h ofﬂine ∼ D i from the ofﬂine data.
We then sample from the posterior z ∼ q φ e ( z | h ofﬂine ) and label the reward r generated of a new state and action, ( s , a ) , using the reward decoder r generated = r φ d ( s , a , z ) , where z ∼ q φ e ( z | h ) .
We then add the labeled trajectory to the buffer and perform actor and critic updates as in ofﬂine metatraining.
Lastly, since we do not observe additional ground-truth rewards, during the self-supervised phase we do not update the reward decoder r φ d or encoder q φ e and instead only train the policy and Q-function.
We visualize this procedure in Figure 3.
We note that the distribution shift discussed in Section 4 is not an issue for the reward decoder.
The distribution shift only occurs when we sample z from the encoder using online data, i.e. z ∼ q φ e ( z | h online ) , but, we only sample from this reward decoder using z sampled from the encoder using ofﬂine data, i.e. z ∼ q φ e ( z | h ofﬂine ) .
The reward decoder does need to generalize to new states and actions, and we hypothesize that reward prediction generalization is easier than policy state, action, and z generalization.
If true, then we would expect that using the reward decoder to label rewards and train on those labels, as in SMAC, will outperform directly using the policy on new tasks (as in Figure 2).
5. A LGORITHM S UMMARY AND D ETAILS
We visualize SMAC in Figure 3.
For ofﬂine training, we assume access to ofﬂine datasets D
=
{D i } N buff i
=1 , where each buffer corresponds to data generated for one task.
Each iteration, we sample a buffer
D i ∼ D and a history from this buffer h ∼ D i .
We condition the stochastic encoder q φ e on this history to obtain a sample z ∼ q φ e ( z | h ) .
We then use this sample z and a second history sample h (cid:48) ∼ D i to update the Q -function, the policy, encoder, and decoder by minimizing Equation (1),
Equation (2), and Equation (3) respectively.
During the self-supervised phase, we found it beneﬁcial to train the actor with a combination of the loss in Equation (2) and the original PEARL actor loss, weighted by hyperparameter λ pearl .
We provide pseudo-code for SMAC in Appendix A.
E XPERIMENTS
We presented a method that uses self-supervised data to mitigate the z -space distribution shift that occurs in ofﬂine meta-RL.
In this section, we evaluate how well the self-supervised phase of SMAC mitigates the resulting drop in performance, and we compare SMAC to prior ofﬂine meta-RL methods on a range of meta-RL benchmark tasks that require generalization to unseen tasks at meta-test time.
Meta-RL tasks.
We ﬁrst evaluate our method on multiple simulated MuJoCo (Todorov et al., 2012) meta-learning tasks that have been used in past online and ofﬂine meta-RL papers (Finn et al., 2017;
Rakelly et al., 2019; Dorfman & Tamar, 2020; Mitchell et al., 2021) (see Figure 9).
Although there are standardized benchmarks for ofﬂine RL (Fu et al., 2020) and meta-RL (Yu et al., 2020), there are no standard benchmarks for ofﬂine meta-RL particular, so we use a combination of tasks in prior work (Dorfman & Tamar, 2020; Mitchell et al., 2021), and introduce a more complex robotic manipulation domain meant to push task generalization to the limit.
We ﬁrst evaluate ﬁve different domains, Cheetah Velocity, Ant Direction, Humanoid, Walker Param, and Hopper Param, which have been used in prior work (Rothfuss et al., 2018; Rakelly et al., 2019; Dorfman & Tamar, 2020). In the ﬁrst three domains, the agent is rewarded for moving in an unobserved target velocity and a meta-episode consists of sampling a desired velocity.
The last task, Humanoid, is particularly challenging due to the high space dimension of 376 .
The last two domains, Walker Param and Hopper Param, require adapting to different dynamics parameters, such as friction, joint mass, and inertia.
The agent must adapt to each task within T = 3 episodes, each of length 200 .
We also evaluated SMAC on a signiﬁcantly more diverse robot manipulation meta-learning task called Sawyer Manipulation, based on the goal-conditioned environment introduced by Khazatsky et al.
This is a simulated PyBullet environment Coumans & Bai (2016–2021) in which a Sawyer robot arm can manipulate various objects.
Sampling a task T ∼ p ( T ) involves sampling both a new conﬁguration of the environment and the desired behavior to achieve, such as pushing a button, opening a drawer, or lifting an object (see Figure 4).
The sparse reward is 0 when the desired behavior is achieved and − 1 otherwise.
The action space consists of 3 dimensions to control the end-effector in Euclidean space and one dimension to control the gripper.
The state is a 13-dimensional vector, which we detail in Appendix C.
The sparse reward, precise manipulation requirement, and diversity of object conﬁguration between each task sample make this environment difﬁcult.
In all of the environments, we test the meta-RL procedure’s ability to generalize to new tasks by evaluating the policies on held-out tasks sampled from the same distribution as in the ofﬂine datasets.
We give a complete description of environments and task distribution in Appendix C.
Ofﬂine data collection.
For the MuJoCo tasks, we use the replay buffer from a single PEARL run that uses the ground-truth reward.
We limit the data collection to 1200 transitions (i.e.,trajectories) per task and terminate the PEARL run early, forcing the meta-RL agent to learn from highly suboptimal data.
For Sawyer Manipulation, we collect data using a scripted policy that randomly performs one of many possible tasks in the environment.
We used 50 training tasks andtrajectories of length 75 each.
Note that the performed task (e.g. open drawer) may differ from the task that is rewarded (e.g. push a button).
As a result, in the ofﬂine dataset, the robot succeeds on the task in 46% of the transitions, indicating that this data is highly suboptimal.
In contrast to prior work (Dorfman & Tamar, 2020; Mitchell et al., 2021), we use this limited data because it enables us to test how well the different methods can improve over suboptimal, ofﬂine data.
See Appendix C for further discussion and details.
Comparisons and ablations.
As an upper bound, we include the performance of PEARL with online training using oracle ground-truth rewards rather than self-generated rewards, which we label
Online Oracle .
To understand the importance of using the actor loss in Equation (2), we include an ablation that replaces Equation (2) with the actor loss from PEARL, which we label SMAC (actor ablation) .
We also include a meta-imitation baseline, labeled meta behavior cloning , which replaces the actor update in Equation (2) with simply maximizing log π θ ( a | s , z ) .
A gap between SMAC and this imitation method would help us understand whether or not our method improves over the (possibly sub-optimal) behavior policy.
For comparisons to prior work, we include the two previously proposed ofﬂine meta-RL methods: meta-actor critic with advantage weighting (labelled MACAW )
(Mitchell et al., 2021) and Bayesian ofﬂine RL (labelled BOReL ) (Dorfman & Tamar, 2020).
MACAW and BOReL assume that rewards are provided during training and cannot make use of the self-supervised online interaction, and so we cannot compare directly to these past works. Therefore, we report their performance only after ofﬂine training.
For these prior work, we used the open-sourced code directly from each paper (PEARL, MACAW, and BOReL), and and trained these methods using the same ofﬂine dataset.
To ensure a fair comparison, we ran the original hyperparameters and matched SMAC hyperparameters (matching network size, learning rate, and batch size), taking the better of the two as the result for each prior method.
We include a list of all hyperparameters and network architectures in Appendix C.3.
Comparison results.
We plot the mean post-adaptation returns and standard deviation acrossseeds in Figure 5.
We see that across all three environments, SMAC consistently improves during the self-supervised phase, and often achieves a similar performance to the oracle that uses ground-truth reward during the online phase of learning.
SMAC also signiﬁcantly improves over meta behavior cloning, which conﬁrms that the data in the ofﬂine dataset is far from optimal.
We also observed even before the online phase, our method is competitive with BOReL and MACAW as a stand-alone ofﬂine meta-RL method.
When self-supervised training is included, we see that
SMAC signiﬁcantly improves over the performance of BOReL on all three tasks and signiﬁcantly improves over MACAW on two of the three tasks, including the Sawyer Manipulation environment, which is by far the most challenging and exhibits the most variability between tasks.
In this domain, we also see the largest gains from the AWAC actor update, in contrast to the actor ablation in blue, which corresponds to a PEARL-style update, indicating that properly handling the ofﬂine phase is also important for good performance.
In conclusion, our method is the only one attains generalization performance close to the Online
Oracle upper bound baseline on all three tasks, indicating that self-supervised online ﬁne-tuning is highly effective for mitigating distributional shift in meta-RL, whereas without it ofﬂine meta-RL generally does not exceed the performance of a meta-imitation learning baseline.
Visualizing the distribution shift.
Lastly, we further investigate if the self-supervised training helps speciﬁcally because it mitigates a distribution shift caused by the exploration policy.
To investigate this, we visualize the trajectories of the learned policy both before and after the selfsupervised phase for the Ant Direction task in Figure 6.
For each plot, we show trajectories from the [Figure] policy π θ ( a | s , z ) when the encoder q φ e ( z | h ) is conditioned on histories from either the ofﬂine dataset ( h ofﬂine ) or from the learned exploration policy ( h online ).
Since the same policy is evaluated, differences between the resulting trajectories represent the distribution shift caused by using history from the learned exploration policy rather than from the ofﬂine dataset.
We see that before the self-supervised phase, there is a large difference between the two modes that can only be attributed to the difference in h .
When using h online , the post-adaptation policy only explores one mode, but when using h ofﬂine , the policy moves in all directions.
This qualitative difference explains the large performance gap observed in Figure 2 and highlights that the adaptation procedure is sensitive to the history h used to adapt.
In contrast, after the self-supervised phase, the policy moves in all directions regardless of where the history came from.
In Appendix B, we also visualize the exploration trajectories and found that the exploration trajectories are qualitatively similar both before and after the self-supervised phase.
Together, these results illustrate the SMAC policy learns to adapt to the exploration trajectories by using the self-supervised phase to mitigate the distribution shift that occurs with naïve ofﬂine meta RL.
C ONCLUSION
In this paper, we studied a problem speciﬁc to ofﬂine meta-RL: distribution shift in the context parameters z .
This distribution shift occurs because the data collected by the meta-learned exploration policy differs from the ofﬂine dataset.
To address this problem, we assumed that an agent can sample new trajectories without additional reward labels and presented SMAC, a method that uses these additional interactions together with state-of-the-art ofﬂine RL techniques to provide an effective meta-RL method. Experimentally, we found that SMAC signiﬁcantly improves over the performance of prior ofﬂine meta-RL methods, after the self-supervised online phase.
A limitation of SMAC is that it assumes that the agent can gather additional unlabeled online samples. However, this assumption may not be satisﬁed due to safety concerns and engineering needed to ensure that automatic resets are feasible, and therefore enabling safe, autonomous environment interactions is an important direction for future research.
Lastly, using self-supervised interaction to mitigate distribution shift in ofﬂine meta RL is orthogonal to the underlying meta RL algorithm. In this paper, we built off of PEARL (Rakelly et al., 2019) due to its simplicity, but this insight could easily be used to improve other existing (Mitchell et al., 2021; Dorfman & Tamar, 2020) or future ofﬂine meta RL algorithms.
To ensure reproducibility, we have included the source code for our method and environments in the supplementary ﬁle, which includes a README ﬁle that describes how to install our code and run the experiments. We describe in detail the environments and hyperparameters used for the experiment in Appendix C. For clarity, we have also included pseudo-code of SMAC on the ﬁrst page of the supplementary material.
Supplementary Material A M ETHOD P SEUDO - CODE
We present the pseudo-code for SMAC in Algorithm 1.
Algorithm 1 Semi-Supervised Meta Actor-Critic
1: Input decoder do (cid:46)
Sample buffer D
∼ D two histories from buffer h h (cid:48) D i . 4: Use the ﬁrst history sample to h to infer z encode it z ∼ q φ e ( h ) . Update π θ , Q w , q φ e , r φ d by minimizing L actor , L critic , L reward with samples z , h (cid:48) .
6: for iteration n = 1 ,
. . , N online do (cid:46) self phase 7:
Collect trajectory τ with π a s with z
∼ p
Sample buffer D i ∼ D and ofﬂine history h
∼ D
Use h ofﬂine to label the rewards in τ , as in Equation (4), and add the resulting data to D i .
Sample buffer D i ∼ D and two histories from buffer h , h (cid:48) ∼ D i .
Encode ﬁrst history z
= q φ e ( h ) .
Update minimizing
B A DDITIONAL E XPERIMENTAL R ESULTS
Exploration and ofﬂine dataset visualization In Figure 7, we visualize the post-adaption trajectories generated when conditioning the encoder the online exploration trajectories h online and the ofﬂine trajectories h ofﬂine .
Similar to Figure 6, and also visualize the online and ofﬂine trajectories themselves.
We see that the exploration trajectories h online and the ofﬂine trajectories h ofﬂine are very different (green vs red, respectively), but the self-supervised phase mitigates the negative impact that this distribution shift has on ofﬂine meta RL.
In particular, the post-adaptation trajectories conditioned on these two data sources (blue and orange) are similar after the self-supervised training, whereas before the self-supervised training, only the trajectories conditioned on the ofﬂine data (blue) move in multiple directions.
Addressing state-space distribution shift by self-supervised meta-training on test tasks.
Another source of distribution shift that can negatively impact a meta-policy is a distribution shift in state space.
While this distribution shift occurs in standard ofﬂine RL, we expect this issue to be
[Figure] [Figure] [Figure]
t - A d a p t a t i o n
R e t u r n s
[Figure]
[Table]
more prominent in meta RL, where there is a focus on generalizing to completely novel tasks.
In many real-world scenarios, experiencing the state distribution of a novel task is possible, but it is the supervision (ie. reward signal) that is expensive to obtain.
Can we mitigate state distribution shift by allow the agent to meta-train in the test task environments, but without rewards?
In this experiment, we evaluate our method, SMAC, when training online on the test tasks instead of on the meta-training tasks as in the experiments in Section 6.
Prior work has explored this idea of self-supervision with test tasks in supervised learning Sun et al.
and goal-conditioned
RL Khazatsky et al.
We use the Sawyer Manipulation environment to study how selfsupervised training can mitigate state distribution shifts, as these environments contain signiﬁcant variation between tasks.
To further increase the complexity of the environment, we use a version of the environment which samples from a set of eight potential desired behaviors instead of three.
We compare self-supervised training on test tasks to self-supervised training on the set of metatraining tasks, which are also the tasks contained in the ofﬂine dataset.
A large gap in performance indicates that interacting with the test tasks can mitigate the resulting distribution shift even when no reward labels are provided.
We show the results in Figure 8 and ﬁnd that there is indeed a large performance gap between the two training modes, with self-supervision on test tasks improving post-adaptation returns while self-supervision on meta-training tasks does not improve post-adaptation returns.
We also compare to an oracle method that performs online training with the test tasks and the ground-truth reward signal.
We see that SMAC is competitive with the oracle, demonstrating that we do not need access to rewards in order to improve on test tasks.
Instead, the entire performance gain comes from experiencing the new state distribution of test tasks.
Overall, these results suggest that SMAC is effective for mitigate distribution shifts in both z -space and state space, even when an agent can interact in the environment without reward supervision.
Reward Accuracy Dependence Our method involves training a reward decoder, and so a natural question is: How good must the reward decoder be for SMAC to work?
We found that the reward loss does not need to be particularly low for our method to work.
Speciﬁcally, in the Sawyer Manipulation environment, the reward scale is either − 1 or 0 , meaning that the maximum reward scale is 1 .
We observed that the reward decoder loss is around 0 .
2 on the training task and around 0 . 25 on the test tasks, indicating that the method does not need a relatively low reward decoder loss to perform well. C E XPERIMENTAL D ETAILS C. D ATA C OLLECTION D IFFERENCE FROM P RIOR W ORK BOReL and MACAW were both developed assuming several orders of magnitude more data than the regime that we tested.
For example, in the BOReL paper (Dorfman & Tamar, 2020), the Cheetah
Velocity was trained with an ofﬂine dataset using 400 million transitions and performs additional reward relabeling using ground-truth information about the transitions.
In contrast, our ofﬂine dataset contains only 240 thousand transitions, roughly three orders of magnitude fewer transitions.
Similarly,
MACAW uses 100M transitions for Cheetah Velocity, over 40 times more transitions than used in our experiments.
These prior methods also collect ofﬂine datasets by training task-speciﬁc policies , which converge to near-optimal policies within the ﬁrst million time step (Haarnoja et al., 2018), meaning that they utilize very high-quality data.
In contrast, our experiments focused on the scenario with many fewer and much lower quality trajectories, which is likely the cause of the relatively worse performance of BOReL and MACAW than in the original papers.
C. E NVIRONMENT D ETAILS
In this section, we describe the state and action space of each environment.
We also describe how reward functions were generated and how the ofﬂine data was generated.
Ant Direction The Ant Direction task consists of controlling a quadruped “ant” robot that can move in a plane.
Following prior work (Rakelly et al., 2019; Dorfman & Tamar, 2020), the reward function is the dot product between the agent’s velocity and a direction uniformly sampled from the unit circle.
The state space is R 20 , comprising the orientation of the ant (in quaternion) as well as the angle and angular velocity of all 8 joints.
The action space is [ − 1 , 1] 8 , with each dimension corresponding to the torque applied to a respective joint.
The reward function is the negative absolute difference between the agent’s x-velocity and a target velocity uniformly sampled from [0 , 3] .
The ofﬂine data is collected by running PEARL (Rakelly et al., 2019) on this meta RL task with 100 pre-sampled 2 target velocities.
We terminate PEARL after 100 iterations, with each iteration consisting of collecting trajectories until at least 1000 new transitions have been observed.
As discussed in Appendix C.1, this results in highly suboptimal data, enabling us to test how well the different methods can improve over the ofﬂine data.
In PEARL, there are two replay buffers saved for each task, one for sampling data for training the encoder and another for training the policy and
Q-function.
We will call the former replay buffer the encoder replay buffer and the latter the RL replay buffer.
The encoder replay buffer contains data generated by only the exploration policy, in which z ∼ p
z ( z ) .
The RL replay buffer contains all data generated, including both exploration and post-adaptation, in which z ∼ q φ e ( z | h ) .
To make the ofﬂine dataset, we load the last 1200 samples of the RL replay buffer and the last 400 transitions from the encoder replay buffer into corresponding
RL and encoder replay buffers for SMAC.
Cheetah Velocity The Cheetah Velocity task consists of controlling a two-legged “half cheetah” that can move forwards or backwards along the x-axis.
Following prior work (Rakelly et al., 2019; Dorfman & Tamar, 2020), the reward function is the absolute difference product between the agent’s x-velocity and a velocity uniformly sampled from [0 , 3] .
The state space is R 20 , comprising the z -position; the cheetah’s x- and z- velocity; the angle and angular velocity of each joint and the half-cheetah’s y-angle; and the XYZ position of the center of mass.
The action space is [ − 1 , 1] 6 , with each dimension corresponding to the torque applied to a respective joint.
The ofﬂine data is collected in the same way as in the Ant Direction task, using a run from PEARL with 100 pre-sampled target velocities.
For the ofﬂine dataset, we use the ﬁrst 1200 samples from the RL replay buffer and last 400 samples from the encoder replay buffer after 50 PEARL iterations, with each iteration containing at least 1000 new transitions.
For only this environment, we found that it was beneﬁcial to freeze the encoder buffer during the self-supervised phase.
Sawyer Manipulation The state space, action space, and reward is described in Section 6.
Tasks are generated by sampling the initial conﬁguration, and then the desired behavior.
There are ﬁve objects: a drawer opened by handle, a drawer opened by button, a button, a tray, and a graspable object.
The state is a 13-dimensional vector comprising of the 3D position of the end-effector and positions of objects in the scene: 3 dimensions for the graspable object, and 1 dimension for each articulated joint in the scene including the robot grippers.
If an object is not present, it takes on position 0 in the corresponding element of the state space.
First, the presence or absence of each of the ﬁve is randomized.
Next, the position of the drawers (from 2 sides), initial position of the tray (from 4 positions), and the object (from 4 positions) is randomized.
Finally, the desired behavior is randomly chosen from the following list, but only including the ones that are possible in the scene: "move hand", "open top drawer with handle", or "open bottom drawer with button".
The ofﬂine data is collected using a scripted controller that does not know the desired behavior and randomly performs potential tasks in the scene, choosing another task if it ﬁnishes one task before the trajectory ends.
This data is loaded into a single replay buffer used for both the encoder and RL. Walker Param
This environment involves controlling a bi-pedal robot that can move along the Zand Y-axis.
The source code is taken from the rand_param_envs repository 3 , which has been used in prior work (Rakelly et al., 2019).
A task involves randomly sampling the body mass, the joint damping coefﬁcients, the body inertia, and the friction parameters from a log-uniform distribution on the range [1 . 5 − 3 , 1 . 5 3 ] .
The reward is the velocity plus a bonus of 1 for staying alive and a control penalty of 10 − 3 × (cid:107) a (cid:107) 2 , where (cid:107) a (cid:107) is the (cid:96) 2 norm of the action.
The state space is 17 dimensions and the action space is 6 dimensions (one for each joint). Hopper Param This environment involves controlling a one-legged robot that can move along the Z- and Y-axis.
This environment is also taken from the rand_param_envs repository.
The reward is the same as in Walker Param, and tasks are sampled the same as for the Walker Param.
The state space is 11 dimensions and the action space is 3 dimensions (one for each joint).
Humanoid This environment is based on the Humanoid environment from OpenAI Gym (Brockman et al., 2016).
We reuse the standard reward function but replace the forward-velocity reward with a velocity reward based on target direction.
Each task consists of sampling a target direction
[Table]
[Table] horizon (max # of transitions per trajectory) [Table] ofﬂine dataset
[Table] uniformly at random. The state space is 376 dimensions and the action space is 17 dimensions (one for each joint).
C.
We list the hyperparameters for training the policy, encoder, decoder, and Q-network in Table 1.
If hyperparameters were different across environments, they are listed in Table 2.
For pretraining, we use the same hyperparameters and train for 50000 gradient steps.
Below, we give details on non-standard hyperparameters and architectures.
Batch sizes.
The RL batch size is the batch size per task when sampling ( s , a , r, s (cid:48) )
tuples to update the policy and Q-network.
The encoder batch size is the size of the history h per task used to conditioned the encoder q φ e ( z | h ) .
The meta batch size is how many tasks batches were sampled and concatenated for both the RL and encoder batches.
In other words, for each gradient update, the policy and Q-network observe ( RL batch size ) × ( meta batch size ) transitions and the encoder observes ( RL batch size ) ×
( encoder batch size ) transitions.
Encoder architecture.
The encoder uses the same architecture as in Rakelly et al.
The posterior is given as the product of independent factors
[Equation]
where each factor is a multi-variate Gaussian over R d z with learned mean and diagonal variance.
In other words,
[Equation]
The mean and standard deviation is the output of a single MLP network with output dimensionality 2 × d z .
The output of the MLP network is split into two halves.
The ﬁrst half is the mean and the second half is passed through the softplus activation to get the standard deviation.
Self-supervised actor update.
The parameter λ pearl controls the actor loss during the selfsupervised phase, which is
[Equation]
where L PEARLactor is the actor loss from PEARL (Rakelly et al., 2019).
For reference, the PEARL actor loss is
[Equation]
When the parameter λ pearl is zero, the actor update is equivalent to the actor update in AWAC (Nair et al., 2020).
Comparisons As discussed in Section 6, we used the authors’ code for PEARL (Rakelly et al., 2019), 4 BOReL (Dorfman & Tamar, 2020), 5 and MACAW (Mitchell et al., 2021).
6 To ensure a fair comparison, we ran the original hyperparameters and matched SMAC hyperparameters (matching network size, learning rate, and batch size), taking the better of the two as the result for each prior method.
For all comparisons, we evaluated the ﬁnal policy using the same evaluation method as
SMAC, i.e. collecting exploration trajectories with the learned policy and perform adaptation using this newly collected data.