H INDSIGHT F ORESIGHT R ELABELING FOR M ETA -R EINFORCEMENT L EARNING
A BSTRACT
Meta-reinforcement learning (meta-RL) algorithms allow for agents to learn new behaviors from small amounts of experience, mitigating the sample inefficiency problem in RL.
However, while meta-RL agents can adapt quickly to new tasks at test time after experiencing only a few trajectories, the meta-training process is still sample-inefficient.
Prior works have found that in the multi-task RL setting, relabeling past transitions and thus sharing experience among tasks can improve sample efficiency and asymptotic performance.
We apply this idea to the meta-RL setting and devise a new relabeling method called Hindsight Foresight Relabeling (HFR).
We construct a relabeling distribution using the combination of hindsight , which is used to relabel trajectories using reward functions from the training task distribution, and foresight , which takes the relabeled trajectories and computes the utility of each trajectory for each task.
HFR is easy to implement and readily compatible with existing meta-RL algorithms.
We find that HFR improves performance when compared to other relabeling methods on a variety of meta-RL tasks.
I NTRODUCTION
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
. R ELATED W ORK
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 cannot 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.
B ACKGROUND 3. R EINFORCEMENT L EARNING
In reinforcement learning (RL), the environment is modeled as a Markov Decision Process (MDP)
M = ( S , A , r, p , γ, p 1 ) , 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 p 1 is the initial state distribution.
At timestep t , the agent π θ , parameterized by parameters θ , observes the state s t ∈ S , takes an action a t ∼ π θ ( a t | s t ) , and observes the next state s t +1 ∼ p
( s t +1 | s t , a t ) and the reward r ( s t , a t ) .
The goal is to maximize the expected cumulative discounted rewards: max θ E s [Equation] ∼ [Equation]
(cid:2)(cid:80) ∞ [Equation] =1 γ t − 1 r ( s t , a t ) (cid:3) . M ETA -R EINFORCEMENT L EARNING
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 ψ , γ, p 1 ) .
The tasks share the components ( S , A , γ, p 1 ) , 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 { s 1 , a 1 , s 2 , a 2 , . . . } , 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:
[Equation]
[Equation]
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 ψ ) .
3. PEARL
In this work, we use PEARL (Rakelly et al., 2019) as our base meta-RL algorithm since it uses offpolicy RL and provides structured exploration via posterior sampling.
PEARL is built on top of Soft
Actor-Critic (Haarnoja et al., 2018) and trains an encoder network q ϕ ( z | c ) that takes in the “context” c , which consists of a batch of ( s t , a t , r t , s t +1 ) transitions, and produces the latent embedding z .
The intent is to learn the encoder such that embedding z encodes some salient information about the task.
The adaptation step f ϕ in PEARL corresponds to generating this latent z and then conditioning the policy and the value function networks on it.
The policy π θ ( a | s, z ) is trained using loss L actor :
[Equation]
where B is the replay buffer.
The critic Q θ ( s, a, z ) and the encoder q ϕ ( z | c ) are trained with temporal difference learning:
[Equation]
[Equation]
where ¯ V is the target state value and ¯ z denotes that the gradient does not flow back through the latent.
H INDSIGHT F ORESIGHT R ELABELING
The objective in this section is to derive a formalism for data-sharing amongst the tasks during the meta-training phase.
This is achieved via trajectory-relabeling, wherein a trajectory collected for a training task ψ i is reused or re-purposed for training a different task ψ j .
Reward-based trajectoryrelabeling has received a lot of attention in recent works on multi-task RL and goal-conditioned
RL (Andrychowicz et al., 2017; Eysenbach et al., 2020; Li et al., 2020).
The intuition is that if a trajectory τ collected while solving for the task ψ i achieves high returns under the reward definition for another task ψ j ( i.e., (cid:80) t r ψ j ( s t , a t ) is large), then τ can be readily used for policy-optimization for the task ψ
j as well.
The meta-RL setting presents the following subtlety – for any given task, the meta-RL agent generates trajectories with the aim of utilizing them in the adaptation procedure and subsequently seeks to maximize the post-adaptation returns ( cf. §3.2).
To improve the efficiency of the meta-training stage, we would like to share these pre-adaptation trajectories amongst the different tasks, accounting for the fact that the metric of interest with these trajectories is their usefulness for task-identification, rather than the returns (as in multi-task RL).
This difference is illustrated in
Figure 3.
Hence, when deciding if a trajectory τ collected for task ψ i is appropriate to be reused for task ψ j , it is sub-optimal to consider the return value of this trajectory under ψ j .
Instead, we argue that this reuse compatibility should be determined based on the performance on the task ψ j , after the agent has undergone adaptation using τ .
Concretely, we define a function to measure the utility of the trajectory τ for a task ψ j :
[Equation]
where π ′ =
f ϕ (cid:0) π θ , τ, r ψ j (cid:1) denotes the policy after using τ for adaptation.
The trajectory-relabeling mechanism during meta-training now incorporates this function U ψ j , which we refer to as the utility function , rather than the return R ψ j .
Broadly, a trajectory τ collected for task ψ i can be relabeled for use in another task ψ j if U ψ j ( τ ) is high.
Subsection §4.1 makes this more precise by deriving a relabeling distribution q ( ψ | τ ) that informs us of the tasks for which τ should be reused.
Figure 4, and the caption therein, describe a high-level overview of our approach, HFR.
Comparison to HIPI (Eysenbach et al., 2020) with a didactic example.
We consider a toy environment to further motivate that return-value based data sharing and trajectory relabeling (as proposed by
HIPI) is potentially sub-optimal for meta-RL.
The Four-Corners environment consists of a point robot placed at the center of a square where each corner of the square represents a goal location, as shown in Figure 1.
For each goal (task), there is a section of the space in the corresponding quadrant in which the robot receives a large negative reward.
Consider a trajectory τ that hovers over the blue square in top-right quadrant.
Note that τ could have been generated by the agent while collecting data for any of the four tasks.
We examine if τ can be reused for the blue task.
Since R blue ( τ ) is highly negative, the relabeling strategy in HIPI does not reuse τ for meta-training on the blue task.
It is clearly evident, however, that τ carries a significant amount of signal pertaining to task-identification on the blue task, making it a useful pre-adaptation trajectory.
HFR reuses τ for the blue task since the utility U blue ( τ ) is high.
To quantify this effect, we include the numerical data on the returns and the utility values for a sampled trajectory that hovers over the blue square in top-right quadrant.
The values for the returns { R purple ( τ ) , R blue ( τ ) , R orange ( τ ) , R green ( τ ) } are {− 20 , − 58 , − 20 , − 20 } , while the utility values { U purple ( τ ) ,
U blue ( τ ) , U orange ( τ ) ,
U green ( τ ) } are {− 1015 , − 756 , − 935 , − 931 } .
These (unnormalized) numbers show that the probability of relabeling this trajectory with the blue task is low under HIPI, but high under
HFR.
Further analysis is included in Appendix A.6.
Figure 2 compares HFR and HIPI in terms of the success-rate in the Four-Corners environment, and shows the performance benefit of using the utility function for trajectory relabeling.
4. D ERIVING A M ETA -RL R ELABELING D ISTRIBUTION
Our derivation in this subsection largely follows HIPI (Eysenbach et al., 2020), but differs in that we adapt it to the meta-RL setting to promote sharing of pre-adaptation trajectories amongst tasks, using the concept of trajectory utility.
Assume a dataset D of trajectories gathered by the meta-RL agent when solving the different tasks in the meta-train task distribution.
We wish to learn a trajectory relabeling distribution q ( ψ | τ ) such that, given any trajectory τ ∼ D , we could reuse τ for tasks with high density under this posterior distribution.
To that end, we start by defining a variational distribution q ( τ | ψ ) to designate the trajectories used for the adaptation process f ϕ , for a given task ψ .
Using the definition of the utility function (Eq. 4), the meta-RL objective from Eq.
1 could be written as: max [Equation] E [Equation]
|
[Equation] [ U ψ
[Equation] τ )] . For fixed meta-parameters ( θ, ϕ ) , a natural approach to optimize the variational distribution q ( τ | ψ ) is to use this same objective since it facilitates alignment with the goals of the meta-learner.
Thus, the combined objective for the variational distributions for all the tasks, augmented with entropy regularization, is:
[Equation]
(5) where H q ( τ | ψ ) denotes the causal entropy of the policy associated with q ( τ | ψ ) .
Now, we note that the above optimization is equivalent to a reverse-KL divergence minimization objective: min q ( τ,ψ ) D KL [ q ( τ, ψ )
|| p ( τ, ψ )]
( cf.
Appendix A.10).
Here, the joint distributions over the tasks and the trajectories are defined as q ( τ, ψ ) =
q ( τ | ψ ) p ( ψ ) and p ( τ, ψ ) = p ( τ | ψ ) p ( ψ ) , where
[Equation]
Our goal is to formulate the trajectory relabeling distribution q ( ψ | τ ) .
To make this explicit in our objective, we use the trick proposed in HIPI (Eysenbach et al., 2020) and factor q ( τ, ψ ) as q ( ψ | τ ) q ( τ ) , thereby rewriting the reverse-KL divergence minimization objective as:
[Equation]
Ignoring the terms independent of ψ , we can analytically solve (by differentiating and setting to zero) for the optimal trajectory relabeling distribution for the meta-RL setting:
[Equation]
Given a trajectory τ ∼ D , the HFR algorithm uses this relabeling distribution to sample tasks for which τ should be reused.
Concretely, we compute the utility function under τ for all the tasks, construct the distribution q ( ψ | τ ) using these utilities (Eq.
8), and sample tasks from it.
Please see
Figure 4 for details.
We assume a uniform prior p ( ψ ) over the tasks in our experiments.
4. A LGORITHM AND I MPLEMENTATION DETAILS
Our relabeling algorithm is summarized in Algorithm 1 and fits seamlessly into the meta-training process of any of the base meta-RL algorithms.
Once the meta-RL agent generates a trajectory τ for a training task, τ is fed as input to HFR, and it returns another task that can reuse this experience τ .
Input : Trajectory to be relabeled ( τ ) Output: Task to relabel the trajectory with ( ψ ) for each training task ψ i do
U ψ i ( τ ) ← ComputeUtility ( τ, ψ i ) log Z (cid:0) ψ i (cid:1) ← GetLogPartition ( ψ i ) end
Return ψ ∼ softmax { U ψ i ( τ ) − log Z ( ψ i ) } (Eq. 8)
Function GetLogPartition( ψ ) :
Sample batch of trajectories (cid:8) τ i (cid:9) N i =1 ∼ B ψ for each trajectory τ
i do
U ψ (cid:0) τ i (cid:1) ← ComputeUtility ( τ i , ψ ) end
Return log Z ( ψ ) ≈ log (cid:16) 1 N (cid:80) Ni =1 e U ψ ( τ i ) (cid:17)
Function ComputeUtility( τ, ψ ) : for each ( s t , a t , r t ) ∈ τ do Replace r t with r ψ ( s t , a t ) end
Sample embedding using encoder z ∼ q ϕ ( z | τ )
Sample a batch of initial states (cid:8) s i 1 (cid:9) N U
i =1 ∼ B ψ Sample actions for these states the
[Equation] [Equation] Return [Equation]
[Equation]
We compute the utility of the input trajectory for every training task, along with an empirical estimate of the log-partition function of the tasks.
The task to relabel the trajectory with is then sampled from a categorical distribution.
For our experiments, we build on top of the PEARL algorithm (Rakelly et al., 2019), which is a data-efficient off-policy meta-RL method.
PEARL maintains task-specific replay buffers B ψ .
If HFR returns the task ψ ′ , then τ is relabeled using the reward function r ψ ′ and added to B ψ ′ for meta-training on the task ψ ′ .
The adaptation procedure π ′ =
f ϕ ( π θ , τ, r ψ ) for a task ψ corresponds to a sequence of steps: 1.)
augment τ by marking each transition with a reward value computed using r ψ ( s t , a t ) , 2.) condition the encoder on τ to sample an embedding, z ∼ q ϕ
( z | τ ) ; and 3.) condition the policy on z to obtain the post-adaptation policy, π ′ =
π θ ( ·| s, z ) .
The calculation of the utility function (Eq. 4) requires generation of post-adaptation trajectories, which could be computationally inefficient, especially if the number of tasks is large.
To avoid this cost, for each task, we sample a batch of initial states s 1 ∼ p 1 ( s 1 ) and the corresponding actions from the post-adaptation policy, and compute the utility based on an estimate of the state-action value function Q π ′ ψ ( s 1 , a 1 ) as:
[Equation]
Since we use PEARL, we can avoid training separate task-specific value functions Q ψ , and instead get the required estimates from the task-conditioned critic Q ( s, a, z ) already used by PEARL (Eq. 3).
We
[Figure]
Visual- highlight facilitates efficient data-sharing among the
E XPERIMENTS
The goal in this section is to quantitatively evaluate the benefit of sharing experience among tasks using HFR, during the meta-train stage.
We evaluate on a set of both sparse and dense reward MuJoCo environments (Todorov et al., 2012) modeled in OpenAI Gym (Brockman et al., 2016).
Please refer to the Appendix for environment details.
We compare HFR with two relabeling methods: Random , in which each trajectory is relabeled with a randomly chosen task, and HIPI (Eysenbach et al., 2020), which utilizes MaxEnt RL to devise an algorithm that relabels each transition using a distribution over tasks involving the soft Q values for that transition.
In contrast, HFR proposes the concept of the utility of a trajectory for a given task.
Using the utility aligns the relabeling methodology with the objective of the meta-RL agent ( cf.
section 4).
All methods are built on top of PEARL.
Finally, we compare to PEARL with no relabeling at all, which we refer to as None .
5. R ESULTS
Sparse Reward Environments.
We first evaluate HFR on a set of sparse reward robotic manipulation and locomotion tasks.
We use five environments: a goal-reaching task involving a quadruped Ant robot, a pushing task on the Sawyer robot, a reaching task on the Sawyer robot, a velocity-matching task involving a bipedal Cheetah robot, and a reaching task involving the MuJoCo Reacher where the agent learns directly from images.
The environments are described in detail in Appendix A. and shown pictorially in Figure 5.
Figure 6 plots the performance (average returns or success-rate) on the held-out meta-test tasks on the y -axis, with the total timesteps of environment interaction for meta-training on the x -axis.
We note that HFR tends to be more sample-efficient than the baselines and achieves a higher asymptotic score.
Meta-RL in sparse reward tasks is hard due to the challenges of task-identification and efficient exploration.
HFR is especially useful for these tasks as the datasharing afforded by the trajectory relabeling algorithm mitigates the need for an elaborate exploration strategy during meta-training.
This leads to the sample-efficiency gains exhibited in Figure 6.
Dense Reward Environments.
We next evaluate HFR on dense reward environments (Figure 7).
We experiment with the Cheetah-Highdim environment, in which the bipedal robot is required to best match its state vector to a set of predefined vectors, as well as the Cheetah-Vel and quadruped
Ant-Vel environments, in which the robots are required to run at various velocities.
Note that the impact of HFR is much less pronounced for these environments, with the exception of
Cheetah-Highdim.
We believe this is because exploration is not as critical for these environments as it was for the sparse reward tasks.
This hypothesis is supported by the fact that, in these environments,
PEARL with no relabeling is competitive with the various relabeling methods, which all share similar performance, whereas in the sparse reward environments HFR is the relabeling method that performs best, with the two other relabeling methods also vastly outperforming vanilla PEARL.
In the case of the Cheetah-Vel and Ant-Vel environments, agents are provided with an informative dense reward that immediately informs them as to which task they’re supposed to be solving.
Although the agent in Cheetah-Highdim is also provided with an informative dense reward, the reward function in this task is a linear combination of an 18-dimensional state vector.
Thus, reasonably good exploration is needed to determine optimal values for each of these 18 dimensions.
HFR relabeling provides improvement over the baselines for this task.
5. A BLATION S TUDIES
Batch Size.
We investigate the impact of the batch size N U used in Algorithm 1 for computing an empirical estimate of the state-action value: 1 N U (cid:80) N U i =
1 Q θ (cid:0) s i 1 , a
i 1 , z (cid:1) .
We compare the effect of batch size across batch sizes { 16 , 32 , 64 , 128 , 256 } .
We expect lower values of N U to lead to a higher variance estimate of the post-adaptation cumulative discounted rewards and thus potentially a worse approximation of the optimal meta-RL relabeling distribution.
In Figure 8a, we see some evidence of this in the comparatively worse performance when using N U = 16 and N U = 32 .
However, we note that even with these small batch sizes, our method still performs well, and in general achieves high returns across all choices of N U .
Partition Function.
We investigate the impact of the log-partition function log Z ( ψ ) in the optimal meta-RL relabeling distribution (Eq. 8).
Prior work has noted the importance of the partition function in the multi-task RL setting when tasks may have different reward scales (Eysenbach et al., 2020).
We believe that in the meta-RL setting, the partition function may be crucial even if the tasks share the reward scale, since some tasks in the meta-training distribution may be easier to solve than others.
We speculate that with the omission of the partition function from our relabeling distribution, trajectories would find high utility and thus be disproportionately labeled with the easily-solved tasks, causing a degradation in overall performance.
In our experiments, we find that the partition function serves as an essential normalization factor; Figure 8b shows an example.
Reward Function.
One assumption our method assumes is access to the true reward function r ψ ( s, a ) , which we can query to get the reward for an individual transition under any training task ψ .
This availability has also been utilized in existing works on relabeling for multi-task RL (Eysenbach et al., 2020; Li et al., 2020).
In many real-world applications of meta-RL, e.g. a distribution over robotic tasks, it is reasonable to assume that the task-designer outlines a rough template for the rewards corresponding to the different tasks from the distribution.
Furthermore, several of the rewards used in our experiments are success/failure indicators, which are simple to specify.
Nevertheless, we consider the scenario where we cannot query the true reward function for individual transitions.
Figures 8c and 8d show good performance even when we use a learned reward function rather than the true reward function to relabel trajectories.
[Figure]
Soft-max vs Hard-max Relabeling Distribution.
Given a trajectory τ , the relabeling distribution derived in Section §4.1 samples tasks for which τ should be reused.
Specifically, tasks are sampled as: ψ ∼ softmax { U ψ i ( τ ) − log Z ( ψ i ) } .
This raises the following question: is it crucial to have stochasticity in the relabeling distribution, or could we deterministically select the task for which the normalized utility value is the highest, i.e., ψ = argmax { U ψ i ( τ ) − log Z ( ψ i ) } ?
A minor modification to the equations in Section §4.1 yields the hard-max relabeling distribution.
Concretely, we can add a term to the starting objective for the variational distribution q that explicitly minimizes the entropy of the relabeling distribution q ( ψ | τ ) :
[Equation]
where ϵ is a value less than 1.
Proceeding with the derivation in the exact same manner as in
Section §4.1, we obtain the adjusted relabeling distribution:
[Equation]
In the limit when ϵ → 0 , q ϵ ( ψ | τ ) is the hard-max relabeling distribution.
Figures 8e, 8f, 8g, and 8h compare the performance of HFR when sampling tasks from a soft-max relabeling distribution (orange) vs a hard-max distribution (blue).
The results indicate that stochasticity is an important factor.
C ONCLUSION
In this paper, we introduced HFR, a trajectory relabeling method for meta-RL that enables data sharing between tasks during meta-train.
We argue that unlike the multi-task RL setting, where the appropriateness of a trajectory for a task could be measured by the returns under the task, for meta-RL, it is preferable to consider the future (expected) task-returns of an agent adapted using that trajectory.
We capture this notion by defining the utility function for a trajectory-task pair and incorporate these utilities in our relabeling mechanism.
Inspired by prior work on multi-task RL, an optimal relabeling distribution is then derived that informs us of the tasks for which a generated trajectory should be reused.
Hindsight is used to relabel trajectories with different reward functions, while foresight is used in computing the utility of each trajectory under different tasks and constructing a relabeling distribution.
HFR is easy to implement, can be integrated into any existing meta-RL algorithm, and yields improvement on a variety of meta-RL tasks, especially those with sparse rewards.
To the best of our knowledge, HFR is the first relabeling method designed explicitly for the meta-RL paradigm.
R EPRODUCIBILITY
To ensure reproducibility, we have created an anonymized repository at https://anonymous.
A A PPENDIX A. F AILED E XPERIMENTS (A N ALTERNATE UTILITY FUNCTION )
While the expected post-adaptation return is a direct indicator of the usefulness of a trajectory for a task, we now mention an indirect measure that is pertinent to the meta-RL baseline algorithm we employ (PEARL).
The idea behind PEARL is that the encoder network q ϕ , conditioned on the given trajectory (referred to as “context”), produces an embedding z that captures some salient information about the task.
The encoder is trained to minimize the Bellman error ( L critic , Eq. 3).
A lower value of
L critic thus indicates that the trajectory is valuable for identifying the task ψ .
Therefore, we experiment with using the negative of L critic as a utility function:
[Equation]
We can efficiently compute this utility function with a single forward pass through the encoder, followed by computation of the Bellman error using a batch of ( s, a, r, s ′ ) tuples.
Algorithm 2 includes the details.
We note that this utility function explicitly gauges the viability of task-identification through the PEARL encoder, rather than measure the post-adaptation returns.
1 − ( Q θ ( s i , a i , z ) − ( r i + V ( s ′ i , z )))
Our results when using this new utility function can be seen in Figure 9.
Overall, relabeling based on the Bellman error (HFR-Bellman) does not perform as well as relabeling based on the expected post-adaptation returns, likely because the Bellman error is simply an indirect measure of the true metric of interest (post-adaptation returns).
A. T RANSITION D YNAMICS
Although this work focuses on the setting where the tasks share the same transition dynamics, but differ in the reward function , in this section we show results on environments where the tasks differ in both the transition dynamics and reward function .
We consider two environments: Ant-Goal-Rand-Params and CheetahVel-Sparse-Rand-Params , which are identical to Ant-Goal and Cheetah-VelSparse , respectively, except for the fact that the transition dynamics now also differ across tasks.
We do not modify any of the methods (HFR, HIPI, None, Random) to take into consideration the different transition dynamics across tasks.
Our results can be seen in Figure 10.
We find that (a) Relabeling can surprisngly still provide some benefit in this setting and (b) HFR still outperforms baselines when transition dynamics differ.
A. H YPERPARAMETERS
Table 1 lists the hyperparameters that were shared across all environments.
Hyperparameters for the Sawyer environments were taken from Yu et al.
(2020), while hyperparameters for the other environments were taken from the open-source
PEARL implementation (Rakelly et al., 2019).
[Figure]
[Table]
[Figure]
. E NVIRONMENTS
The environments we evaluate HFR on are as follows:
Four -Corners: We create a 2D navigation task in which a point robot must navigate to one of four goal locations.
A reward of 0 is given when the robot is within a distance of 0.2 from the goal, with the episode ending.
Otherwise, a reward of − 1 is given, with a reward of − 3 being given if the robot is within a certain section of space in the same quadrant as the goal (Figure 1).
Ant-Goal: We use the Ant-Goal task from (Gupta et al., 2018).
Tasks correspond to goal locations sampled uniformly from a half-circle of radius 2.
The reward function is 4 − || x
ant − x goal
||
if || x ant − x goal || 2 ≤ 0 .
8 and 4 − c otherwise, where c is a large uninformative constant.
[Table]
Ant-Vel: We use the Ant environment from OpenAI gym.
Tasks correspond to goal velocities sampled uniformly in [0 , 3] .
The reward is the negative absolute value of the difference between the agent’s velocity and the goal velocity.
We take this task from (Finn et al., 2017).
Cheetah-Highdim: We take the Cheetah-Highdim task from (Lin et al., 2020).
Each task corresponds to an 18-dimensional vector ψ sampled uniformly in [ − 1 , 1] 18 .
The reward is linear in the posttransition state s ′ and is given as r ( s, a, s ′ )
=
ψ T s ′ .
Cheetah-Vel: The Cheetah-Vel task is taken from (Rakelly et al., 2019).
Velocities are sampled uniformly in [0 , 3] and the reward is the negative absolute value of the difference between the agent’s velocity and the goal velocity.
Cheetah-Vel-Sparse: Velocities are sampled uniformly in [0 , 3] , with positive reward being given if the absolute value of the difference between the agent’s velocity and the goal velocity is less than 0 . 3 .
Sawyer-Push: An agent must control a simulated Sawyer robot and push a block to a specified goal location.
A reward of 0 is given when the block is within a distance of 0 .
07 of the goal and a reward of − 1 is given otherwise.
The episode ends when either the goal is reached or the agent has taken 150 steps.
Both the Sawyer Reach and Sawyer Push environment were taken from (Yu et al., 2020).
Sawyer-Reach: An agent must control a simulated Sawyer robot and reach a specified 3D goal location.
A reward of 0 is given when the end effector is within a distance of 0 .
05 of the goal and a reward of − 1 is given otherwise.
The episode ends when either the goal is reached or the agent has taken 150 steps.
Visual-Reacher: We use the standard MuJoCo Reacher environment where an agent must reach various 2D goal locations, with positive reward being given if the end effector is within a distance of 0 .
03 of the goal.
The agent must meta-learn directly from 64 × 64 grayscale images of the environment as the input observations.
Table 2 includes the details for each environment and Figure 5 shows them pictorially.
The average return shown in our plots is the average return obtained after N initial exploration steps, where N is the value under Number of Exploration Steps, after which the agent should attempt to solve the task
. A. T IME /S PACE C OMPLEXITY
HFR incurs an O ( NC ) cost every time a trajectory is collected, where N is the number of tasks in the training task distribution and C is the cost of computing the utility of a relabeled trajectory.
Computing the utility of a relabeled trajectory involves passing this trajectory through the encoder to generate a latent embedding, sampling a batch of initial states from the replay buffer, passing these initial states and embedding to the policy network to generate a batch of actions, and evaluating the Q function on these state-action pairs (Algorithm 1).
HFR takes O ( N ) space to store the relabeling distribution.
A. A NALYSIS OF T RAJECTORIES IN THE F OUR -C ORNERS ENVIRONMENT
In this section, we analyze sample trajectories in the Four-Corners environment, as shown in Figure 11.
Each trajectory’s original goal is different from the goal located in the quadrant it explores.
In Table 3,
[Figure] [Figure]
HFR
HIPI
Top Left − 727 .− Top Right − 957 .−
Left − 985 .−
Right − 937 .−
Task
HFR
HIPI
Top Left − 1015 .− Top Right − 756 .− Bottom Left − 935 .− Right
[Table] [Table] [Table]
[Table]
trajectory even though τ may be extremely useful for meta-training on ψ .
By considering post-adaptation returns , HFR does not suffer from this issue
A. L EARNING A R EWARD F UNCTION
In this section we describe how we learn a reward function in the setting where we do not assume we can query the true reward function r ψ ( s t , a t ) for individual transitions, as explored in Subsection §5.2.
We can use a learned reward function r ωψ ( s t , a t ) parameterized by a neural network with parameters ω to relabel trajectories.
This network takes in task ψ and state and action s t , a t as input.
We train our reward function using the mean-squared error:
[Equation]
This update rule is applied during every training update during the meta-training process.
A. A DAPTATION AT M ETA -T EST T IME
In Algorithm 3, we include the meta-test time procedure of PEARL used in HFR.
In PEARL, adaptation procedure f ϕ corresponds to a forward pass through encoder q ϕ .
Algorithm 3:
Input : Test task( ψ )
Number exploration trajectories (
Initialize context c ψ = {} for k = 1 to K do
Sample embedding z ∼ q ϕ
( z | c ψ )
Collect trajectory τ k
= { ( s t , a t , r t , s t +1 ) }
Nt =1 using policy π θ ( ·| s, z ) Append trajectory τ k to context: c ψ = c ψ ∪ τ k end
Sample embedding z ∼ q ϕ
( z | c ψ ) Compute adapted policy π ′ =
f ϕ ( π θ , c ψ , r ψ ) =
π θ ( ·| s, z )
Roll out adapted policy π to returns on ψ
A. I S THE LEARNED Q FUNCTION CORRECTLY CAPTURING TASK INFORMATION ?
In this section, we examine whether it makes sense to use the learned Q function as the source of relabeling signal.
To do this, we train HFR to convergence, freeze the weights of the HFR context encoder q ϕ , and take a set of trajectories T N = (cid:8)
τ Nk (cid:9) K k =1 from the training process that have not been relabeled.
We split T N into two collections of trajectories, which we denote as T N train and T N test .
We construct
T R by taking the trajectories from T N test and relabeling them using HFR (Algorithm 1).
We then use our trained HFR context encoder q ϕ to generate three collections of context embeddings: Z N train =
{ z ∼ q ϕ ( ·| τ ) : τ ∈ T N train } , Z N test =
{ z ∼ q ϕ ( ·| τ ) : τ ∈ T N test } , Z R = { z ∼ q ϕ ( ·| τ ) : τ ∈ T R } .
We then train a N -way classifier f ( ψ | z ) on Z N train .
Given context embedding z , f should correctly identify the task ψ that z was generated from.
If z was generated from a relabeled trajectory, then f should classify z as coming from the task that the trajectory was relabeled with.
If z was generated from a trajectory that was not relabeled, then f should classify z as coming from the task that the trajectory was originally generated for.
We evaluate the accuracy of our classifier on Z N test and Z R .
If
HFR is correctly capturing task information during relabeling, we would expect the accuracy on Z R to be greater than or equal to the accuracy on Z N test .
For clarity, we illustrate this entire procedure in
Algorithm 4.
[Table]
We conduct this experiment on the Four-Corners environment, as well as the Cheetah-Vel-Sparse environment with 10 training tasks.
We parameterize f as a feed-forward neural network with two hidden layers, each with 200 units, and ReLU activations.
We train f using cross-entropy loss forepochs on a training set of non-relabeled trajectories and report the accuracy on relabeled trajectories and a test set of non-relabeled trajectories at the end of training.
Our results are shown in Table 4.
We find that the accuracy on relabeled trajectories is higher than the accuracy on non-relabeled
[Table]
trajectories.
This indicates that HFR is indeed correctly capturing task information when relabeling – our relabeled trajectories are easier to identify as coming from a specific task than the non-relabeled trajectories and are thus more useful for adaptation.
A. O BJECTIVE E QUIVALENCE
In this section, we show that the objective in Equation 5 is equivalent to the reverse-KL divergence minimization objective: min q ( τ,ψ ) D KL [ q ( τ, ψ )
|| p ( τ, ψ )] , where q ( τ, ψ ) =
q ( τ | ψ ) p ( ψ ) , and p ( τ, ψ ) = p ( τ | ψ ) p ( ψ ) .
Further, p ( τ | ψ ) is defined in Equation 6.
[Equation]
where H q ( τ | ψ ) denotes the causal entropy of the policy associated with q ( τ | ψ ) , that is ˜ q
( a | s ; ψ )
. A. C OMPARISON TO P RIOR W ORK IN E XPERIENCE R ELABELING A.11. C OMPARISON WITH MIER (M ENDONCA
ET AL ., 2020)
Mendonca et al.
(2020) propose to tackle meta-RL via a model identification process, where contextdependent neural networks parameterize the transition dynamics and the rewards function.
These models are trained with supervised meta-learning.
For a new task (at meta-train/test time), the context is adapted to encapsulate the task information.
Although MIER uses experience relabeling, the scope and the purpose of this relabeling is very different from HFR:
• Relabeling during meta-train vs meta-test.
In MIER, experience relabeling is performed only at 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 motivation of the relabeling in HFR is to improve the sample-efficiency of the meta-training phase.
HFR does not modify the meta-test time behavior of the underlying meta-RL algorithm.
• Problem setup.
MIER considers a setup where the meta-test tasks are “out-of-distribution” w.r.t.
the meta-train task distribution.
The authors, therefore, introduce relabeling to handle this discrepancy.
HFR operates in the standard meta-RL setting where the meta-test tasks are “in-distribution”.
MIER does not perform relabeling in this case (please see Figure 2 in their paper).
• Relabeling strategy.
Lastly, the relabeling strategy is very different for the two approaches.
While MIER relabels all of the meta-training data for the new meta-test time task, HFR selectively relabels the trajectory for a new task
during
meta- train, basing the selection on the idea of the utility function and a relabeling distribution q ( ψ | τ ) .
A.11. C OMPARISON WITH BOR E L (D ORFMAN ET AL ., 2020)
Dorfman et al.
(2020) propose a method for offline meta-RL where the meta-learning agent is provided with fixed trajectory data collected in different environments and the goal of the agent is to learn to maximize returns in an unseen environment from the same task distribution.
The authors highlight the issue of “MDP ambiguity”, which refers to the difficulty in MDP identification using the belief from a VAE encoder trained with offline data.
To mitigate this challenge, the authors propose a reward relabeling scheme that replaces the reward in a trajectory from an MDP- i in the offline data, with rewards from a randomly chosen MDP- j .
There are two critical differences from our work:
• Motivation for relabeling.
Relabeling in BOReL is done to alleviate the “MDP ambiguity” issue, which the authors mention is specific to the offline meta-RL setting.
HFR operates in the online meta-RL setting and our motivation for relabeling is to improve the sampleefficiency of the online meta-training phase via data sharing.
• Relabeling strategy.
Perhaps the more important difference is in the strategy used to select the task for which a given trajectory should be relabeled.
While BOReL chooses the task at random, HFR provides a principled procedure of computing a relabeling distribution q ( ψ | τ ) over the tasks using the concept of the utility function and sampling a task from this distribution for the relabeling.
Empirically, we observe that this approach is more performant than random task relabeling.
Figure 6 in the paper includes the learning curves for both the methods – HFR outperforms Random in all those tasks.