Relay Policy Learning: Solving Long-Horizon Tasks via Imitation and Reinforcement Learning

Abhishek Gupta, Vikash Kumar, Corey Lynch, Sergey Levine, Karol Hausman

Introduction

Recent years have seen reinforcement learning (RL) successfully applied to a number of robotics tasks such as in-hand manipulation , grasping and door opening . However, these applications have been largely constrained to relatively simple short-horizon skills. Hierarchical reinforcement learning (HRL) has been proposed as a potential solution that should scale to challenging long-horizon problems, by explicitly introducing temporal abstraction. However, HRL methods have traditionally struggled due to various practical challenges such as exploration , skill segmentation and reward definition . We can simplify the above-mentioned problems by utilizing extra supervision in the form of unstructured human demonstrations, in which case the question becomes: how should we best use this kind of demonstration data to make it easier to solve long-horizon robotics tasks?

This question is one focus area of hierarchical imitation learning (HIL), where solutions typically try to achieve two goals: i) learn a temporal task abstraction, and ii) discover a meaningful segmentation of the demonstrations into subtasks. These methods have not traditionally been tailored to further RL fine-tuning, making it challenging to apply them to a long-horizon setting, where pure imitation is very likely to fail. To address this need, we devise a simple and universally-applicable two-phase approach that in the first phase pre-trains hierarchical policies using demonstrations such that they can be easily fine-tuned using RL during the second phase. In contrast to HRL methods, our method takes advantage of unstructured demonstrations to bootstrap further fine-tuning, and in contrast to conventional HIL methods, it does not focus on careful subtask segmentation, making the method simple, general and very amenable to further reinforcement fine-tuning. In particular, we show that we can develop an imitation and reinforcement learning approach that while not necessarily perfect at imitation learning, is very amenable to improvement via fine-tuning with reinforcement learning and that can be scaled to challenging long-horizon manipulation tasks.

What are the advantages of using such an algorithm? First, the approach is very general, in that it can be applied to any demonstration data, including easy to provide unsegmented, unstructured and undifferentiated demonstrations of meaningful behaviors. Second, our method does not require any explicit form of skill segmentation or subgoal definition, which otherwise would need to be learned or explicitly provided. Lastly, and most importantly, since our method ensures that every low-level trajectory is goal-conditioned (allowing for a simple reward specification) and of the same, limited length, it is very amenable to reinforcement fine-tuning, which allows for continuous policy improvement. We show that relay policy learning allows us to learn general, hierarchical, goal-conditioned policies that can solve long-horizon manipulation tasks in a challenging kitchen environment in simulation, while significantly outperforming hierarchical RL algorithms and imitation learning algorithms.

Related Work

Typical solutions for solving temporally extended tasks have been proposed under the HRL framework . Solutions like the options framework , HAM , max-Q , and feudal networks present promising algorithmic frameworks for HRL. A particularly promising approach was proposed in Nachum et al. and Levy et al. , using goal conditioned policies at multiple layers of hierarchy for RL. Nevertheless, these algorithms still suffer from challenges in exploration and optimization (as also seen in our experimental comparison with Nachum et al. ), which have limited their application to general robotic problems. In this work, we tackle these problems by using additional supervision in the form of unstructured, unsegmented human demonstrations. Our work builds on goal-conditioned RL , which has been explored in the context of reward-free learning , learning with sparse rewards , large scale generalizable imitation learning , and hierarchical RL . We build on this principle to devise a general-purpose imitation and RL algorithm that uses data relabeling and bi-level goal conditioned policies to learn complex skills.

There has a been a number of hierarchical imitation learning (HIL) approaches that typically focus on extracting transition segments from the demonstrations. These methods aim to perform imitation learning by learning low-level primitives or latent conditioned policies which meaningfully segment the demonstrations. Traditionally, these approaches do not aim to and are not amenable to improving the learned primitives with subsequent RL, which is necessary as we move towards multi-task, challenging long-horizon problems where pure imitation might be insufficient. In this work, we specifically focus on utilizing both imitation and RL, and devise a method that does not explicitly try to segment out individual primitives into meaningful subtasks, but instead splits the demonstration data into fixed-length segments, amenable to fine-tuning with reinforcement learning. This allows us to leverage relabeling across different goals . We introduce a novel form of goal relabeling and demonstrate its efficiency when applied to learning robust bi-level policies. A related idea is presented in Le et al. , where the authors assume that an expert provides labelled and segmented demonstrations at both levels of the hierarchy, with an interactive expert for guiding RL. In contrast, we use a pool of unlabelled demonstrations and apply our method to learn a policy to achieve various desired goals, without needing interactive guidance or segmentation. Using imitation learning as a way to bootstrap RL has been previously leveraged by a number of deep RL algorithms , where a flat imitation learning initialization is improved using reinforcement learning with additional auxiliary objectives. In this work, we show that we can learn hierarchical policies in a way that can be fine-tuned better than their flat counterparts.

Preliminaries

Goal-conditioned imitation learning:

In typical imitation learning, instead of knowing the reward rr, the agent has access to demonstrations D\mathcal{D} containing a set of trajectories D={τi,τj,τk,...}\mathcal{D}=\{\tau^{i},\tau^{j},\tau^{k},...\} of state-action pairs τi={s0i,a0i,…,sTi,aTi}\tau^{i}=\{s_{0}^{i},a_{0}^{i},\dots,s_{T}^{i},a_{T}^{i}\}. The goal is to learn a policy π(a∣s)\pi(a|s) that imitates the demonstrations. A common approach is to maximize the likelihood of actions in the demonstration, i.e. max⁡E(s,a)∼Dlog⁡π(a∣s)\max E_{(s,a)\sim\mathcal{D}}\log\pi(a|s), referred to as behavior cloning (BC). When there are multiple demonstrated tasks, we consider a goal-conditioned imitation learning setup where the dataset of demonstrations D\mathcal{D} contains sequences that attempt to reach different goals sgi,sgj,sgk,...s_{g}^{i},s_{g}^{j},s_{g}^{k},.... The objective is to learn a goal-conditioned policy π(a∣s,sg)\pi(a|s,s_{g}) that is able to reach different goals sgs_{g} by imitating the demonstrations.

Relay Policy Learning

In this section, we describe our proposed relay policy learning (RPL) algorithm, which leverages unstructured demonstrations and reinforcement learning to solve challenging long-horizon tasks. Our approach consists of two phases: relay imitation learning (RIL), followed by relay reinforcement fine-tuning (RRF) described in Sec. 4.2 and 4.3 respectively. While RIL by itself is not able to solve the most challenging tasks that we consider, it provides a very effective initialization for fine-tuning.

We first introduce our bi-level hierarchical policy architecture (shown in Fig 3), which enables us to leverage temporal abstraction. This architecture consists of a high-level goal-setting policy and a low-level subgoal-conditioned policy, which together generate an environment action for a given state. The high-level policy πθh(sgl∣st,sgh)\pi_{\theta}^{h}(s_{g}^{l}|s_{t},s_{g}^{h}) takes the current state sts_{t} and a long-term high-level goal sghs_{g}^{h} and produces a subgoal sgl∈Ss_{g}^{l}\in\mathcal{S} which is then ingested by a low-level policy πϕl(a∣st,sgl)\pi_{\phi}^{l}(a|s_{t},s_{g}^{l}). The low-level policy takes the current state sts_{t}, and the subgoal sgls_{g}^{l} commanded by the high-level policy and outputs an action ata_{t}, which is executed in the environment.

Importantly, the goal setting policy πθh\pi_{\theta}^{h} makes a decision every HH time steps (set to 3030 in our experiments), with each of its subgoals being kept constant during that period for the low-level policy, while the low-level policy πϕl\pi_{\phi}^{l} operates at every single time-step. This provides temporal abstraction, since the high level policy operates at a coarser resolution than the low-level policy. This policy architecture, while inspired by goal-conditioned HRL algorithms , requires a novel learning algorithm to be applicable in the context of imitation learning, which we describe in Sec. 4.2. Given a high-level goal sghs_{g}^{h}, πθh\pi_{\theta}^{h} samples a subgoal sg0ls_{g_{0}}^{l}, which is passed to πθl\pi_{\theta}^{l} to generate action a0a_{0}. For the subsequent HH steps, the goal produced by πθh\pi_{\theta}^{h} is kept fixed, while πθl\pi_{\theta}^{l} generates an action ata_{t} at every time step.

2 Relay Imitation Learning

Our problem setting assumes access to a pool of unstructured, unlabeled “play" demonstrations (Lynch et al. ) D\mathcal{D}, corresponding to demonstrations of meaningful activities provided by the user, without any particular task in mind, e.g. opening cabinet doors, playing with different objects, or simply tidying up the scene. We do not assume that this demonstration data actually accomplishes any of the final task goals that we will need to solve at test-time, though we do need to assume that the test-time goals come from the same distribution of goals as those accomplished in the demonstration data. In order to take the most advantage of such data, we initialize our policy with our proposed relay imitation learning (RIL) algorithm. RIL is a simple imitation learning procedure that builds on the goal relabeling scheme described in Lynch et al. for the hierarchical setting, resulting in improved handling of multi-task generalization and compounding error. RIL assumes access to the pool of demonstrations consisting of NN trajectories D={τi,τj,τk,...}\mathcal{D}=\{\tau^{i},\tau^{j},\tau^{k},...\}, where each trajectory consists of state-action pairs τi={s0i,a0i,…,sTi,aTi}\tau^{i}=\{s_{0}^{i},a_{0}^{i},\dots,s_{T}^{i},a_{T}^{i}\}. Importantly, these demonstrations can be attempting to reach a variety of different high level goals sghs_{g}^{h}, but we do not require these goals to be specified explicitly. To learn the relay policy from these demonstrations, we construct a low-level dataset Dl\mathcal{D}_{l}, and a high-level dataset Dh\mathcal{D}_{h} from these demonstrations via “relay data relabeling", which is described below, and use them to learn πθh\pi_{\theta}^{h} and πθl\pi_{\theta}^{l} via supervised learning at multiple levels.

We construct the low-level dataset by iterating through the pool of demonstrations and relabeling them using our relay data relabelling algorithm. First, we choose a window size WlW_{l} and generate state-goal-action tuples for Dl\mathcal{D}_{l}, (s,sgl,a)(s,s_{g}^{l},a) by goal-relabeling within a sliding window along the demonstrations, as described in detail below and in Algorithms 2, 3. The key idea behind relay data relabeling is to consider all states that are actually reached along a demonstration trajectory within WlW_{l} time steps from any state sts_{t} to be goals reachable from the state sts_{t} by executing action ata_{t}. This allows us to label all states st+1,....,st+Wls_{t+1},....,s_{t+W_{l}} along a valid demonstration trajectory as potential goals that are reached from state sts_{t}, when taking action ata_{t}. We repeat this process for all states sts_{t} along all the demonstration trajectories being considered. This procedure ensures that the low-level policy is proficient at reaching a variety of goals from different states, which is crucial when the low-level policy is being commanded potentially different goals generated by the high-level policy.

We employ a similar procedure for the high level, generating the high-level state-goal-action dataset Dh\mathcal{D}_{h}. However, the actions at the high level are subgoal states that are provided to the low-level policy, so they must be chosen as states along the demonstration trajectories. We start by choosing a high-level window size WhW_{h}, which encompasses the high-level goals we would like to eventually reach. We then generate state-goal-action tuples for Dh\mathcal{D}_{h}, via relay data relabeling within the high-level window being considered, as described in Algorithm 2, 3. We also label all states st+1,....,st+Whs_{t+1},....,s_{t+W_{h}} along a valid trajectory as potential high-level goals that are reached from state sts_{t} by the high level policy, but we set the high-level action for a goal jj steps ahead st+js_{t+j}, as st+min⁡(Wl,j)s_{t+\min(W_{l},j)} choosing a sufficiently distant subgoal as the high-level action.

Given these relay-data-relabeled datasets, we train πθl\pi_{\theta}^{l} and πθh\pi_{\theta}^{h} by maximizing the likelihood of the actions taken given the corresponding states and goals:

This procedure gives us an initialization for both the low-level and the high-level policies, without the requirement for any explicit goal labeling from a human demonstrator. As we show in our experiments, this bi-level initialization is significantly more amenable to RRF than learning the high level from scratch as described in , and allows us to avoid the expensive goal labeling that is required in . Relay data relabeling not only allows us to learn hierarchical policies without explicit labels, but also provides algorithmic improvements to imitation learning: (i) it generates more data through the relay-data-relabelling augmentation, and (ii) it improves generalization since it is trained on a large variety of goals.

3 Relay Reinforcement Fine-tuning

The procedure described in Sec. 4.2 allows us to extract an effective policy initialization via relay imitation learning. However, this policy is often unable to perform well across all temporally extended tasks, due to the well-known compounding errors stemming from imitation learning . Reinforcement learning provides a solution to this challenge, by enabling continuous improvement of the learned policy directly from experience. We can use RL to improve RIL policies via fine-tuning on different tasks. We employ a goal-conditioned HRL algorithm that is a variant of natural policy gradient (NPG) with adaptive step size , where both the high-level and the low-level goal-conditioned policies πθh\pi_{\theta}^{h} and πϕl\pi_{\phi}^{l} are being trained with policy gradient in a decoupled optimization.

Given a low-level goal-reaching reward function rl(st,at,sgl)r_{l}(s_{t},a_{t},s_{g}^{l}), we can optimize the low-level policy by simply augmenting the state of the agent with the goal commanded by the high-level policy and then optimizing the policy to effectively reach the commanded goals by maximizing the sum of its rewards. For the high-level policy, given a high-level goal-reaching reward function rh(st,gt,sgh)r_{h}(s_{t},g_{t},s_{g}^{h}), we can optimize it by running a similar goal-conditioned policy gradient optimization to maximize the sum of high-level rewards obtained by commanding the current low-level policy.

To effectively incorporate demonstrations into this reinforcement learning procedure, we leverage our method via: (1) initializing both πθl\pi_{\theta}^{l} and πθh\pi_{\theta}^{h} with the policies learned via RIL, and (2) encouraging policies at both levels to stay close to the behavior shown in the demonstrations. To incorporate (2), we augment the NPG objective with a max-likelihood objective that ensures that policies at both levels take actions that are consistent with the relabeled demonstration pools DlD_{l} and DhD_{h} from Section 4.2, as described in Eqn 2 and 3:

While a similar objective has been described in , it is yet to be explored in the hierarchical, goal-conditioned scenarios, which makes a significant difference as indicated in our experiments.

In addition, since we are learning goal-conditioned policies at both the low and high level, we can leverage relay data relabeling as described in Sec. 4.2 to also enable the use of off-policy data for fine-tuning. Suppose that at a particular iteration ii, we sampled NN trajectories according to the scheme proposed in Sec. 4.1. While these trajectories did not necessarily reach the goals that were originally commanded, and therefore cannot be considered optimal for those goals, they do end up reaching the actual states visited along the trajectory. Thus, they can be considered as optimal when the goals that they were intended for are relabeled to states along the trajectory via relay data relabeling described in Algorithm 2, 3. This scheme generates a low-level dataset Dli\mathcal{D}_{l}^{i} and a high level dataset Dhi\mathcal{D}_{h}^{i} by relabeling the trajectories sampled at iteration ii. Since these are considered “optimal” for reaching goals along the trajectory, they can be added to the buffer of demonstrations DlD_{l} and DhD_{h}, thereby contributing to the objective described in Eqn 2 and Eqn 3 and allowing us to leverage off-policy data during RRF. We experiment with three variants of the fine-tuning update in our experimental evaluation: IRIL-RPL (fine-tuning with Eqn 2, 3 and iterative relay data relabeling to incorporate off-policy data as described above), DAPG-RPL (fine-tuning the policy with the update in Eqn 2, 3 without the off-policy addition) and NPG-RPL (fine-tuning the policy with the update in Eqn 2, 3, without the off-policy addition or the second maximum likelihood term). The overall method is described in Algorithm 1.

As described in Ghosh et al. , it is often difficult to learn multiple tasks together with on-policy policy gradient methods, because of high variance and conflicting gradients. To circumvent these challenges, we use RPL to fine-tune on a number of different high level goals individually, and then distill all of the learned behaviors into a single policy as described in Rusu et al. . This allows us to learn a single policy capable of achieving multiple high level goals, without dealing with the challenges of multi-task optimization.

Experimental Results

Our experiments aim to answer the following questions: (1) Does RIL improve imitation learning with unstructured and unlabelled demonstrations? (2) Is RIL more amenable to RL fine-tuning than its flat, non-hierarchical alternatives? (3) Can we use RPL to accomplish long-horizon manipulation tasks? Videos and further experimental details are available at https://relay-policy-learning.github.io/

To evaluate our algorithm, we utilize a challenging robotic manipulation environment modeled in MuJoCo, shown in Fig. 1. The environment consists of a 9 DoF position-controlled Franka robot interacting with a kitchen scene that includes an openable microwave, four turnable oven burners, an oven light switch, a freely movable kettle, two hinged cabinets, and a sliding cabinet door. We consider reaching different goals in the environment, as shown in Fig. 4, each of which may require manipulating many different components. For instance, in Fig. 4 (a), the robot must open the microwave, move the kettle, turn on the light, and slide open the cabinet. While the goals we consider are temporally extended, the setup is fully general. We collect a set of unstructured and unsegmented human demonstrations described in Sec. 4.2, using the PUPPET MuJoCo VR system . We provide the algorithm with 400 sequences containing various unstructured demonstrations that each manipulate four different elements of the scene in sequence.

Evaluation and Comparisons

Since each of our tasks consist of compound goals that involve manipulating four elements in the environment, we evaluate policies based on the number of steps that they complete out of four, which we refer to as step-completion score. A step is completed when the corresponding element in the scene is moved to within ϵ\epsilon distance of its desired position.

We compare variants of our RPL algorithm to a number of ablations and baselines, including prior algorithms for imitation learning combined with RL and methods that learn from scratch. Among algorithms which utilize imitation learning combined with RL, we compare with several methods that utilize flat behavior cloning with additional finetuning. Specifically, we compare with (1) flat goal-conditioned behavior cloning followed by finetuning (BC), (2) flat goal-conditioned behavior cloning trained with data relabeling followed by finetuning (GCBC) , and variants of these algorithms that augment the BC and GCBC fine-tuning with losses as described in Rajeswaran et al. - (3) DAPG-BC and (4) DAPG-GCBC. We also compare RPL to (5) hierarchical imitation learning + finetuning with an oracle segmentation scheme, which performs hierarchical goal conditioned imitation learning by using a hand-specified oracle to segment the demonstrations for imitation learning, followed by RRF style fine-tuning. Details of this scheme can be found in Appendix 3. For comparisons with methods that learn from scratch we compare with (6) an on-policy variant of HIRO trained from scratch with natural policy gradient instead of Q-learning and (7) a baseline (Pre-train low level) that learns low-level primitives from the demonstration data, and learns the high-level goal-setting policy from scratch with RL. The last baseline is representative of a class of HIL algorithms , which are difficult to fine-tune because it is not clear how to provide rewards for improving low-level primitives. Lastly, we compare RPL with a baseline (7) (Nearest Neighbor) which uses a nearest neighbor strategy to choose the demonstration which has the achieved goal closest to the commanded goal and subsequently executes actions open-loop.

1 Relay Imitation Learning from Unstructured Demonstrations

We start by aiming to understand whether RIL improves imitation learning over standard methods. We compare the step-wise completion scores averaged over 17 different compound goals with RIL as compared to flat BC variants. We find that, while none of the variants are able to achieve near-perfect completion scores via just imitation, the average stepwise completion score is higher for RIL as compared to both flat variants (see Table 1, first row). Additionally, we find that the flat policy with data augmentation via relabeling performs better than without relabeling. When we analyze the proportion of compound goals that are actually fully achieved (see Table 1, bottom row), RIL shows significant improvement over other methods. This indicates that, even for imitation learning, we see benefits from introducing the simple RIL scheme described in Sec. 4.2.

2 Relay Reinforcement Fine-tuning of Imitation Learning Policies

Although pure RIL does succeed at times, its performance is still relatively poor. In this section, we study the degree to which RIL-based policies are amenable to further reinforcement fine-tuning. Performing reinforcement fine-tuning individually on 17 different compound goals seen in the demonstrations, we observe a significant improvement in the average success rate and stepwise completion scores over all the baselines when using any of the variants of RPL (see Fig. 5). In our experiments, we found that it was sufficient to fine-tune the low-level policy, although we could also fine-tune both levels, at the cost of more non-stationarity. Although the large majority of the benefit is from RRF, we find a slight additional improvement from the DAPG-RPL and IRIL-RPL schemes, indicating that including the effect of the demonstrations throughout the process helps.

When compared with HRL algorithms that learn from scratch (on-policy HIRO ), we observe that RPL is able to learn much faster and reach a much higher success rate, showing the benefit of demonstrations. Additionally, we notice better fine-tuning performance when we compare RPL with flat-policy fine-tuning. This can be attributed to the fact that the credit assignment and reward specification problems are much easier for the relay policies, as compared to fine-tuning flat policies, where a sparse reward is rarely obtained. The RPL method also outperforms the pre-train-low-level baseline, which we hypothesize is because we are not able to search very effectively in the goal space without further guidance. We also see a significant benefit over using the oracle scheme described in Appendix 3, since the segments become longer making the exploration problem more challenging. The comparison with the nearest neighbor baseline also suggests that there is a significant benefit from actually learning a closed-loop policy rather than using an open-loop policy. While plots in Fig. 5 show the average over various goals when fine-tuned individually, we can also distill the fine-tuned policies into a single, multi-task policy, as described in Sec. 5, that is able to solve almost all of the compound goals that were fine-tuned. While the success rate drops slightly, this gives us a single multi-task policy that can achieve multiple temporally-extended goals (Fig 5).

3 Ablations and Analysis

To understand design choices, we consider the role of using different window sizes for RPL as well as the role of reward functions during fine-tuning. In Fig 6 (left), we observe that the window size for RPL plays a major role in algorithm performance. As window size increases, both imitation learning and fine-tuning performance decreases since the behaviors are now more temporally extended.

Next, we consider the role of the chosen reward function in fine-tuning with RRF. We evaluate the relative performance of using different types of rewards for fine-tuning - sparse reward, euclidean distance, element-wise reward (refer to Appendix A for details). When each is used as a goal conditioned reward for fine-tuning the low-level, sparse reward works much better. This indicates that when exploration is sufficient, sparse reward functions are less prone to local optima than alternatives.

Conclusion and Future Work

We proposed relay policy learning, a method for solving long-horizon, multi-stage tasks by leveraging unstructured demonstrations to bootstrap a hierarchical learning procedure. We showed that we can learn a single policy capable of achieving multiple compound goals, each requiring temporally extended reasoning. In addition, we demonstrated that RPL significantly outperforms other baselines that utilize hierarchical RL from scratch, as well as imitation learning algorithms.

In future work, we hope to tackle the problem of generalization to longer sequences and study extrapolation beyond the demonstration data. We also hope to extend our method to work with off-policy RL algorithms, so as to further improve data-efficiency and enable real world learning on a physical robot.

Acknowledgements

We would like to thank Byron David for his help in improving our kitchen environment and simulation. We would also like to thank Suraj Nair, Chelsea Finn, Ofir Nachum, Michael Ahn, Anusha Nagabandi, Dibya Ghosh for fruitful discussions. We also thank Robotics at Google for a wonderful research atmosphere.

References

Appendix A Experimental Details

We use feed-forward MLPs for all our policies, with two layer neural networks with 256 units each and ReLu nonlinearities used for both the high-level policy πθ\pi_{\theta} and the low-level policy πϕ\pi_{\phi} in all methods. Flat baselines use the same architecture as well and additional experimentation with the architecture did not yield substantially different results. We train all imitation learning algorithms with the ADAM optimizer using a batch size of 128128 and a learning rate of 0.0050.005. We choose WlW_{l} to be 3030 and WhW_{h} to be 260260 in all experiments. Our ablations suggest that the larger the window, the harder the learning problem becomes for both imitation and RL fine-tuning.

For reinforcement learning, we utilize a variant of Trust Region Policy Optimization (TRPO). We fine-tune on 17 different compound goals individually, with a path length of 280280 for every compound goal, and the low-level horizon set to 3030. We use 100100 trajectories in each iteration of on-policy fine-tuning, with a discount of 0.9950.995. When using variants of augmenting the policy gradient objective with demonstrations, we experimented with different weights λh\lambda_{h} and λl\lambda_{l}, but we found 0.00010.0001 to work well. We use a batch size of a 100100 trajectories per iteration, and fairly standard parameters for truncated natural policy gradient based on https://github.com/aravindr93/mjrl

The simulation environment has a 30-dimensional state space which consists of positions of the arm and the objects in the scene. The action space is 9 dimensional with 7 DoF for the arm and 2 DoF for the gripper. The actions are represented as the joint velocity.

Appendix B Reward Function Details

For the comparisons detailed in Section 5.3, the reward functions used for sparse, euclidean and element-wise sparse reward functions are detailed below, with ϵ\epsilon set to 0.30.3. For all our experimental results in Fig 5, we use the sparse reward variant as the reward function for fine-tuning.

In the element-wise sparse reward case, idx is selected to be the indices of state corresponding to different distinct elements of the scene such as the microwave, stove burners, light switch, sliding cabinet, hinge cabinets and so on. The robot arm is excluded from these indices.

Appendix C Oracle Baseline Details

For the oracle comparison described in Section 5, a hand-designed scheme is used to segment the demonstration into segments corresponding to semantically meaningful components, thereby generating variable sized windows rather than fixed length ones. Specifically, we split a segment any time one of [microwave, kettle, light switch, burners, slide cabinet, hinge cabinet] is moved more than ϵ=0.3\epsilon=0.3. This leads to a variable segment generation scheme, which generates splits that is shown in Fig 7.

Segments generated in this fashion can then be used for imitation learning both the low-level and high-level policies. Specifically, the actions for the high level policies are chosen to be the states at which the segments are broken, and the low level is trained via goal conditioned behavior cloning with those states set as goals.

Appendix D Visualization of Learned Behaviors

We show example visualizations of several successful learned behaviors for compound tasks, and some failed behaviors to better understand the the method. These can be best appreciated by viewing the accompanying videos on the supplementary website https://relay-policy-learning.github.io.

D.2 Failure Cases