Hindsight Foresight Relabeling for Meta-Reinforcement Learning

1 INTRODUCTION . Deep Reinforcement Learning ( RL ) has achieved success on a wide variety of tasks , ranging from computer games to robotics . However , RL agents are typically trained on a single task and are extremely sample-inefficient , often requiring millions of samples to learn a good policy for just that one task . Ideally , RL agents should be able to utilize their prior knowledge and adapt to tasks quickly , just as humans do . Meta-learning , or learning to learn , has achieved promising results in this regard , allowing agents to exploit the shared structure between tasks in order to adapt to new tasks quickly during meta-test time . Although meta-learned policies can adapt quickly during meta-test time , training these meta-learned policies could still require a large amount of data . Several popular meta-RL methods ( Duan et al. , 2016 ; Finn et al. , 2017 ; Mishra et al. , 2017 ; Rothfuss et al. , 2018 ; Wang et al. , 2016 ) utilize on-policy data during meta-training to better align with the setup at meta-test time , where the agent must generate on-policy data for an unseen task and use it for adapting to the task . Recent works ( Fakoor et al. , 2019 ; Rakelly et al. , 2019 ) have sought to incorporate off-policy RL ( Fujimoto et al. , 2018 ; Haarnoja et al. , 2018 ) into meta-RL to improve sample efficiency . The combination of off-policy RL and relabeling , in which experience is shared across tasks , has been utilized in the multi-task RL setting , in which an agent learns to achieve multiple different yet related tasks , for both goal-reaching tasks ( Andrychowicz et al. , 2017 ) and more general multi-task settings ( Eysenbach et al. , 2020 ) . Experience collected for one task may be completely useless for training a policy to learn that task , but could be extremely informative in training a policy to learn a different task . For example , an agent trying to shoot a hockey puck into a net might miss to the right . This experience could easily be used to train an agent to shoot a puck into a net positioned further to the right ( Andrychowicz et al. , 2017 ) . Both meta-RL and multi-task RL involve training on a distribution of tasks , so it follows that we can also combine relabeling techniques with meta-RL algorithms in order to boost both sample efficiency and asymptotic performance . In meta-RL , an agent learns to explore sufficiently to identify the task it is supposed to be solving , and then uses that knowledge to achieve high task returns . The agent collects exploratory pre-adaptation data , then undergoes some adaptation process using that pre-adaptation data . Finally , after adaptation , the agent attempts to solve the task . Meta-RL algorithms typically have a meta-training phase followed by a meta-test phase . The goal during meta-training is to train the meta-parameters such that they could be quickly adapted to solve any task from the meta-train task distribution , given a small amount of data from that task . At meta-test time , given a new unseen task , the goal is to rapidly adapt the learned meta-parameters for this task , using a small amount of task-specific data . The focus in this paper is to improve the sample efficiency of the meta-training phase via data sharing . Using concepts from maximum entropy RL ( MaxEnt RL ) , we introduce a relabeling scheme for the meta-RL setting . Prior relabeling methods for multi-task RL have used the total reward of the trajectory under different tasks to guide the relabeling ( Eysenbach et al. , 2020 ; Li et al. , 2020 ) . Direct application of this type of relabeling to the meta-RL setting is potentially sub-optimal since the multi-task RL and meta-RL objectives are distinct ( learning to perform many tasks vs. learning to learn a new task ) . Towards developing an approach more suited to meta-RL , we define the notion of the utility of a trajectory under the different tasks , where the utility captures the usefulness of the trajectory for efficient adaptation under those tasks . We call our method Hindsight Foresight Relabeling ( HFR ) – we use hindsight in replaying the experience using reward functions from different tasks , and we use foresight in computing the utility of trajectories under different tasks and constructing a relabeling distribution over tasks using these utilities . We demonstrate the efficacy of our method on a variety of robotic manipulation and locomotion tasks . Notably , we show that our method , as the first meta-RL relabeling technique ( applied during meta-training ) that we are aware of , leads to improved performance compared to prior relabeling schemes designed for multi-task RL . 2 RELATED WORK . Meta-learning , or learning to learn ( Baxter , 1998 ; Naik et al. , 1992 ; Schmidhuber , 1987 ; Thrun & Pratt , 1998 ) , has been a topic of interest since the 1980s . Various approaches have been developed in recent years . Prior works have attempted to represent the RL process using a recurrent neural network ( RNN ) ( Duan et al. , 2016 ; Miconi et al. , 2018 ; Wang et al. , 2016 ) - the hidden state is maintained across episode boundaries and informs the policy as to what task it is currently solving . Similarly , Mishra et al . ( 2017 ) also maintain the internal state across episode boundaries while incorporating temporal convolution and attention into a recursive architecture . Gradient-based metalearning methods have also been explored ( Finn et al. , 2017 ; Nichol & Schulman , 2018 ; Xu et al. , 2018 ; Zheng et al. , 2020 ) . MAML ( Finn et al. , 2017 ) seek to learn a good policy initialization so that only a few gradient steps are needed to achieve good performance on unseen meta-test tasks . Stadie et al . ( 2018 ) build on this gradient-based approach but explicitly consider the effect of the original sampling distribution on final performance . Similar to these works , our relabeling method also considers the impact of pre-adaptation data on the post-adaptation performance . Another body of work focuses on designing strategies for structured exploration in meta-RL such that task-relevant information could be efficiently recovered ( Liu et al. , 2020 ; Rakelly et al. , 2019 ; Zintgraf et al. , 2019 ) . Rakelly et al . ( 2019 ) devise an off-policy meta-RL method called PEARL that trains an encoder to generate a latent context vector on which the meta-RL agent is conditioned . Although we use PEARL as our base algorithm in this work , our relabeling scheme is general enough to be integrated into any off-policy meta-RL algorithm . Experience Relabeling in Meta-RL . Recent work has studied the scope of sharing experience among tasks in the meta-RL paradigm . Mendonca et al . ( 2020 ) propose to tackle meta-RL via a model identification process , where context-dependent neural networks parameterize the transition dynamics and the rewards function . Their method performs experience relabeling only at the meta-test time , with the purpose of consistent adaptation to the out-of-distribution tasks . Crucially , there is no relabeling or sharing of data amongst tasks during the meta-train time . In contrast , the goal of the relabeling in HFR is to improve the sample efficiency of the meta-training phase . Dorfman et al . ( 2020 ) study the offline meta-RL problem . They propose reward relabeling as a mechanism to mitigate the “ MDP ambiguity ” issue , which the authors note is specific to the offline meta-RL setting . Their relabeling is based on random task selection . HFR , on the other hand , operates in the online meta-RL setting and provides a principled approach to compute a relabeling distribution that suggests tasks for relabeling . We compare with the random relabeling method used in Dorfman et al . ( 2020 ) in our experiments . A more detailed comparison to these two prior works is included in Appendix A.11 . In the context of multi-task RL ( Caruana , 1997 ; Kaelbling , 1993 ; Schaul et al. , 2015 ) , recent methods have proposed relabeling to improve the sample-efficiency ( Andrychowicz et al. , 2017 ; Eysenbach et al. , 2020 ; Li et al. , 2020 ) . HER ( Andrychowicz et al. , 2017 ) relabels transitions using goals that the agent actually achieves . Doing so allows for learning even with a sparse binary reward signal . However , HER is only applicable to goal-reaching tasks and can not be incorporated into meta-RL algorithms because the meta-RL agent is trained on a batch of tasks sampled from a fixed task distribution . Similar to our work , Eysenbach et al . ( 2020 ) use MaxEnt RL to construct an optimal relabeling distribution for multi-task RL and apply this to both goal-reaching tasks and tasks with arbitrary reward functions . 3 BACKGROUND . 3.1 REINFORCEMENT LEARNING . In reinforcement learning ( RL ) , the environment is modeled as a Markov Decision Process ( MDP ) M = ( S , A , r , p , γ , p1 ) , where S is the state-space , A is the action-space , r is the reward function , p is the transition dynamics , γ ∈ [ 0 , 1 ) is the discount factor , and p1 is the initial state distribution . At timestep t , the agent πθ , parameterized by parameters θ , observes the state st ∈ S , takes an action at ∼ πθ ( at|st ) , and observes the next state st+1 ∼ p ( st+1|st , at ) and the reward r ( st , at ) . The goal is to maximize the expected cumulative discounted rewards : maxθ Est , at∼πθ [ ∑∞ t=1 γ t−1r ( st , at ) ] . 3.2 META-REINFORCEMENT LEARNING . In the general meta-reinforcement learning ( meta-RL ) setting , there is a family of tasks that is characterized by a distribution p ( ψ ) , where each task ψ is represented by an MDP Mψ = ( S , A , rψ , pψ , γ , p1 ) . The tasks share the components ( S , A , γ , p1 ) , but can differ in the reward function rψ ( e.g . navigating to different goal locations ) and/or the transition dynamics pψ ( e.g . locomotion on different terrains ) . In this work , we consider the setting where the tasks share the same transition dynamics ( i.e. , pψ = p ) , but differ in the reward function . The goal in meta-learning is to learn a set of meta-parameters such that given a new task from p ( ψ ) and small amount of data for the new task , the meta-parameters can be efficiently adapted to solve the new task . In the context of meta-RL , given new task ψ , the agent collects some initial trajectories { τpre } , each being a sequence { s1 , a1 , s2 , a2 , . . . } , and then undergoes some adaptation procedure fϕ ( πθ , τpre , rψ ) ( e.g. , a gradient update ( Finn et al. , 2017 ) or a forward pass through an RNN ( Duan et al. , 2016 ) ) . The adaptation procedure returns a new policy π′ . Using this post-adaptation policy , the agent should seek to maximize the cumulative discounted rewards it achieves . Overall , the meta-RL objective is : max θ , ϕ Eψ∼p ( ψ ) , ( st , at ) ∼π′ ( θ , ϕ ) [ ∞∑ t=1 γt−1rψ ( st , at ) ] ; π′ ( θ , ϕ ) = fϕ ( πθ , τpre , rψ ) ( 1 ) where θ , ϕ are the meta-parameters that are learned in the meta-training phase . A meta-RL agent must learn a good adaptation procedure fϕ that is proficient in extracting salient information about the task at hand , using few pre-adaptation trajectories τpre . At the same time , it should learn the policy meta-parameters θ such that it can achieve high returns after the adaptation process , i.e. , while following the policy π′ = fϕ ( πθ , τpre , rψ ) .

Summary: This paper studies task relabelling in hindsight to increase the efficiency of meta-reinforcement-learning. The authors propose a strategy for calculating a distribution of tasks for which a particular batch of data would be useful for adaptation, and sample from this distribution to construct a relabelled batch which augments the training data. The authors show empirically that this improves sample efficiency over more naive relabelling schemes, particularly for sparse reward tasks. A series of ablations further justifies several design decisions or investigates robustness to hyperparameters.