Offline Meta-Reinforcement Learning with Online Self-Supervision

1 INTRODUCTION . 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 efficient 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 . Offline 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 offline dataset , rather than doing so in the inner loop of RL training , and the same offline multi-task data can be reused repeatedly for many training runs . While a few recent works have proposed offline meta-RL algorithms Dorfman & Tamar ( 2020 ) ; Mitchell et al . ( 2021 ) , we identify a specific problem when an agent trained with offline meta-RL is tested on a new task : the distributional shift between the behavior policy from the offline data and the meta-test time exploration policy means that adaptation procedures learned from offline 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 offline meta-RL never trains on data generated by the meta-learned exploration policy . In practice , we find 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 offline 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 offline data . Based on this principle , we propose semi-supervised meta actor-critic ( SMAC ) , which uses rewardlabeled offline data to bootstrap a semi-supervised meta-reinforcement learning procedure , in which an offline meta-RL agent collects additional online experience without any reward labels . SMAC uses the reward supervision from the offline 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 : first , it is the first method to combine the efficient PEARL ( Rakelly et al. , 2019 ) amortized inference method for meta-RL with AWAC ( Nair et al. , 2020 ) , an effective approach for offline RL with online finetuning . This already yields an effective new offline meta-RL procedure . Second , it is the first approach to perform offline meta-RL with self-supervised online finetuning , without ground truth rewards . We evaluate our method and prior offline 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 find 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 significantly better performance , on par with an online meta-RL method that receives fully labeled online interaction data . 2 RELATED WORKS . 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 offline 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 offline 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 , 1993 ; Schaul et al. , 2015 ; Andrychowicz et al. , 2017 ; Pong et al. , 2018 ; Colas et al. , 2018 ; Warde- Farley 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 defined by a single goal state or fixed basis function , and only requires reward labels for static offline data . Our method addresses a similar problem to prior offline meta-RL methods ( Mitchell et al. , 2021 ; Dorfman & Tamar , 2020 ) . In our comparisons , we find that after offline meta-training , our method is competitive to these prior approaches , but that with additional self-supervised online fine-tuning , our method significantly 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 offline dataset with reward labels that we can use to learn to generate similar rewards . 3 PRELIMINARIES . Meta-reinforcement learning . In meta-RL , we assume there is a distribution of tasks pT ( · ) . A task T is a Markov decision process ( MDP ) , defined by a tuple T = ( S , A , r , γ , p0 , pd ) , where S is the state space , A is the action space , r is a reward function , γ is a discount factor , p0 ( s0 ) is the initial state distribution , and pd ( st+1 | st , at ) is the environment dynamics distribution . A replay buffer D is a set of state , action , reward , next-states tuples , D = { si , ai , ri , s′i } Nsizei=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 τ = ( s1 , a1 , s2 , . . . ) without reward labels . A meta-episode consists of sampling a task T ∼ pT ( · ) , 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 specific 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 offline 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 = { si , ai , ri , s′i } Nenci=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 , Qw ( 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 . Offline reinforcement learning . In offline reinforcement learning , we assume that we have access to a dataset D collected by some behavior policy πβ . An RL agent must train on this fixed dataset and can not interact with the environment . One challenge that offline 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 offline 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 offline RL with meta-RL by training meta-RL on multiple datasets D1 , . . . , DNbuff ( Dorfman & Tamar , 2020 ; Mitchell et al. , 2021 ) , but in the next section we describe some limitations of this combination .

Summary: This work investigates the meta offline reinforcement learning problems. It highlights an issue called "distribution shift in z space" in the meta offline RL. In the meta-test phase, the agent needs to explore the test environments and the collected transitions are used to generate z representing the task feature. The policy is conditioned on the task feature z, so it is important that the transitions collected by the exploration policy are informative enough to extract the feature information. In the offline RL setting, the distritbution of transitions collected by the exploration policy deviates from the distribution of transitions in the static training dataset. Therefore, inferring task feature z becomes problematic.  This paper proposes an additional online training process (without the reward info) to fix the issue of "distribution shift in z space". It learns a reward function from the offline training dataset, and thus it is able to generate the reward for transitions collected during online training. These online training transitions with the estimated rewards are used to update the learned policy.