Hindsight Task Relabelling: Experience Replay for Sparse Reward Meta-RL

1 Introduction . Reinforcement learning ( RL ) has seen tremendous success applied to challenging games ( Mnih et al. , 2015 ; Silver et al. , 2017 ) and robotic control ( Lillicrap et al. , 2015 ; Levine et al. , 2016 ) , driven by advances in compute and the use of deep neural networks as powerful function approximators in RL algorithms . However , agents trained using deep RL often struggle to meaningfully utilize past experience to learn new tasks , even if the new tasks differ only slightly from tasks seen during training ( Zhang et al. , 2018 ; Packer et al. , 2018 ; Cobbe et al. , 2019 ) . In contrast , humans are adept at utilizing prior experience to rapidly acquire new skills and adapt to unseen environments . Meta-reinforcement learning ( meta-RL ) aims to address this limitation by extending the RL framework to explicitly consider structured distributions of tasks ( Schmidhuber , 1987 ; Bengio et al. , 1990 ; Thrun & Pratt , 1998 ) . Whereas conventional RL is concerned with learning a single task , meta-RL is concerned with learning to learn , that is , learning how to quickly learn a new task by leveraging prior experience on related tasks . Meta-RL methods generally utilize a limited amount of experience in a new environment to estimate a latent task embedding which conditions the policy , or to compute a policy gradient which is used to directly update the parameters of the policy . A major challenge in both RL and meta-RL is learning with sparse rewards . When rewards are sparse or delayed , the adaptation stage in meta-RL becomes extremely difficult : inferring the task at hand requires receiving reward signal from the environment , which in the sparse reward setting only happens after successfully completing the task . Due to this inherent incompatibility between meta-RL and sparse rewards , existing meta-RL algorithms that consider the sparse reward setting either only work in simple environments that do not require temporally-extended exploration strategies for adaptation ( Duan et al. , 2016 ; Stadie et al. , 2018 ) , or train exclusively using dense reward functions ( Gupta et al. , 2018b ; Rakelly et al. , 2019 ) , which are designed to encourage an agent to learn adaptation strategies that can be directly applied to the original sparse reward setting . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . In this paper , we show that the concept of hindsight relabeling ( Andrychowicz et al. , 2017 ) from conventional RL can be applied in meta-RL to enable learning to learn in the sparse reward setting . Our key insight is that data collected on the true training tasks can be relabelled as pseudo-expert data for easier hindsight tasks , bootstrapping meta-training with the reward signal needed to train the agent . We introduce a new algorithm for off-policy meta-RL , which we call Hindsight Task Relabeling ( HTR ) , and demonstrate its effectiveness by achieving state-of-the-art performance on a collection of challenging sparse reward environments that previously required shaped reward to solve . Not only is our approach able to learn adaptation strategies only using sparse reward , but it also learns strategies with comparable performance to existing approaches that use shaped reward functions . 2 Related Work . In meta-RL , an agent learns an adaptation strategy by repeatedly adapting to training tasks in the inner loop of meta-training , with the hope that the learned adaptation procedure will generalize to new , unseen tasks during meta-testing . Context-based methods use recent experience ( i.e. , context , in the form of trajectories or transitions ) in a new task to estimate a latent task embedding , and differ mainly in how this embedding is computed : Duan et al . ( 2016 ) ; Wang et al . ( 2016 ) ; Mishra et al . ( 2018 ) propose aggregating context into the hidden state of the policy , whereas Rakelly et al . ( 2019 ) propose explicitly feeding the task embedding as input to the policy . Gradient-based methods use recent context to update differentiable hyperparameters ( Xu et al. , 2018b ) , loss functions ( Sung et al. , 2017 ; Houthooft et al. , 2018 ) , or to directly update policy parameters ( Finn et al. , 2017 ; Stadie et al. , 2018 ; Xu et al. , 2018a ; Zintgraf et al. , 2018 ; Gupta et al. , 2018b ; Rothfuss et al. , 2019 ) . Many of the aforementioned approaches struggle to learn effective exploration strategies for tasks with sparse or delayed rewards . Methods that directly update the policy either implicitly ( e.g. , with a hidden state ) or explicitly ( e.g. , with gradients ) are generally unable to learn a policy that meaningfully explores , since the primary source of randomness is action-space , and thus is time-invariant . Gupta et al . ( 2018b ) and Rakelly et al . ( 2019 ) propose using a probabilistic task variable that is sampled once per episode , which makes the primary source of variability task-dependent and enables temporallyextended exploration . While this approach enables effective adaptation in the sparse reward setting , both Gupta et al . ( 2018b ) and Rakelly et al . ( 2019 ) learn the actual adaptation strategies using shaped ( dense ) reward functions during meta-training . In the environments they consider , the adaptation strategies learned using dense rewards generalize to sparse rewards despite the difference in reward structure ; however , engineering a suitable reward function can be costly ( Ng et al. , 1999 ) and error prone ( Clark & Amodei , 2016 ) , and thus it is highly desirable to devise a mechanism for learning to learn in the sparse reward setting that does not require an auxiliary dense reward function . Our proposed method is closely related to prior work in unsupervised meta-RL and goal generation . Unsupervised meta-RL ( Jabri et al. , 2019 ; Gupta et al. , 2018a ) considers the meta-RL setting without access to a training task distribution ; training tasks are instead self-generated with the aim of learning skills useful for the test task distribution . Our proposed method also generates a curriculum of training tasks , but with the intent of learning adaptation strategies in the sparse reward setting , as opposed to learning transferable skills in a setting with no rewards at all . Several methods exist for curriculum generation in the context of goal-conditioned policies , including adversarial goal generation ( Florensa et al. , 2018 ) and goal relabeling ( Kaelbling , 1993 ; Andrychowicz et al. , 2017 ; Levy et al. , 2017 ; Nair et al. , 2018 ) . Recent work ( Li et al. , 2020 ; Eysenbach et al. , 2020 ) has proposed the use of inverse RL to generalize goal relabeling to broader families of tasks . Unlike prior work on goal and task relabeling that use goal- or task-conditioned policies , we focus on the meta-RL setting where the task is unknown to the policy and must be inferred through deliberate and coherent exploration . 3 Background . In this section we formalize the meta-RL problem setting and briefly describe two algorithms which our approach , Hindsight Task Relabeling , builds on : Probabilistic Embeddings for Actor-critic RL ( PEARL ) and Hindsight Experience Replay ( HER ) . 3.1 Meta-Reinforcement Learning ( Meta-RL ) . Conventional RL assumes environments are modeled by a Markov Decision Processes ( MDP ) M , defined by the tuple ( S , A , p , r , γ , ρ0 , T ) where S is the set of possible states , A is the set of actions , p : S×A×S → R≥0 is the transition probability density , r : S×A → R is the reward function , γ is the discount factor , ρ0 : S → R≥0 is the initial state distribution at the beginning of each episode , and T is the time horizon of an episode ( Sutton & Barto , 2017 ) . Let st and at be the state and action taken at time t. At the beginning of each episode , s0 ∼ ρ0 ( · ) . Under a stochastic policy π mapping a sequence of states to actions , at ∼ π ( at | st , · · · , s0 ) and st+1 ∼ p ( st+1 | st , at ) , generating a trajectory τ = { st , at , r ( st , at ) } , where t = 0 , 1 , · · · . Conventional RL algorithms , which assume the environment is fixed , learn π to maximize the expected reward per episode JM ( π ) = Eπ [ ∑T t=0 γ trt ] , where rt = r ( st , at ) . They often utilize the concepts of a value function V πM ( s ) , the expected reward conditional on s0 = s and a state-action value function QπM ( s , a ) , the expected reward conditional on s0 = s and a0 = a. Meta-RL algorithms assume that there is a distribution of tasks p ( T ) , and usually maximize the expected reward over the distribution , EπT ∼p ( T ) [ JT ( π ) ] . The distribution of tasks corresponds to a distribution of MDPs M ( T ) , where T can paramaterize an MDP ’ s dynamics or reward . In this work we focus on the latter case , where each T implies an MDP with a different reward function under the same dynamics . The agent is trained on a set of training tasks Ttrain ∼ p ( T ) during meta-training , and evaluated on a held-out set of testing tasks Ttest ∼ p ( T ) during meta-testing . An important distinction between meta-RL and multi-task or goal-conditioned RL is that in meta-RL , the task is never explicitly revealed to the agent or policy π through observations s. Instead , π must explore the environment to infer the task to maximize expected reward . This crucial difference is one reason why Hindsight Experience Replay can not be directly applied to the meta-RL setting . 3.2 Off-Policy Meta-Reinforcement Learning . We demonstrate the effectiveness of our approach by building on top of Probabilistic Embeddings for Actor-critic RL ( PEARL ) by Rakelly et al . ( 2019 ) , an off-policy meta-RL algorithm which itself is built on top of soft actor-critic ( SAC ) by Haarnoja et al . ( 2018 ) . SAC is an actor-critic algorithm based on the maximum entropy RL objective that optimizes a stochastic policy π with off-policy data . In addition to the policy ( actor ) π , SAC concurrently learns twin state-action value functions ( critics ) Q. Transitions collected using the policy π are added to the replay buffer B , and Q and π are optimized using loss functions regularized by the entropy of π. Rakelly et al . ( 2019 ) extend SAC to the meta-RL setting by adding a context encoder that uses recent history in an environment to estimate a task embedding , which is then utilized by the policy . Specifically , PEARL uses a network qφ ( z | c ) that takes as input recently collected data ( i.e. , context c ) to infer a probabilistic latent context variable Z . Samples from the latent context variable condition the actor , πθ ( a | s , z ) , and critic , Qθ ( s , a , z ) . Instead of a single replay buffer ( as in SAC ) , there are T separate buffers Bi=1 ... T , corresponding to the number of tasks sampled for Ttrain . The separate per-task replay buffers enable sampling off-policy ( but on-task ) context during the inner loop of meta-training , which significantly improves data efficiency . During meta-training , transitions are collected for each Bi by sampling z ∼ qφ ( z | c ) and acting according to πθ ( a | s , z ) . z is periodically sampled during data collection such that the context contains transitions collected using the policy conditioned on the prior , as well the same policy conditioned on the posterior . Gradients used to optimize qφ , πθ and Qθ are computed in task-specific batches across a random subset of training tasks Ti ∼ p ( T ) , each with an associated replay buffer Bi . During meta-testing , for each test task Ti ∼ p ( T ) , an empty buffer Bi is initialized , z is sampled from the prior , and the policy conditioned on z is rolled out . As additional context is added to the replay buffer , the belief over z narrows , enabling the policy to adapt to the task at hand .

Summary: This paper describes an extension of Hindsight Experience Replay (HER) to a meta-RL setting with task inference. Like HER, trajectories are relabelled with rewards as if they had achieved a goal. Unlike HER, the goal cannot be explicitly relabelled. Instead, a context is provided from which the agent attempts to learn a task embedding: this task inference is also trained using the reward-relabelled data.