Learning and Retrieval from Prior Data for Skill-based Imitation Learning
Abstract: Imitation learning offers a promising path for robots to learn generalpurpose behaviors, but traditionally has exhibited limited scalability due to high data supervision requirements and brittle generalization.
Inspired by recent advances in multi-task imitation learning, we investigate the use of prior data from previous tasks to facilitate learning novel tasks in a robust, data-efﬁcient manner.
To make effective use of the prior data, the robot must internalize knowledge from past experiences and contextualize this knowledge in novel tasks.
To that end, we develop a skill-based imitation learning framework that extracts temporally extended sensorimotor skills from prior data and subsequently learns a policy for the target task that invokes these learned skills.
We identify several key design choices that signiﬁcantly improve performance on novel tasks, namely representation learning objectives to enable more predictable skill representations and a retrieval-based data augmentation mechanism to increase the scope of supervision for policy training.
On a collection of simulated and real-world manipulation domains, we demonstrate that our method signiﬁcantly outperforms existing imitation learning and ofﬂine reinforcement learning approaches.
Videos and code are available at https://ut-austin-rpl.github.io/sailor
Introduction
A long-standing dream in robotics is to enable robots to perform general-purpose tasks in diverse environments.
In recent years, imitation learning has led to great progress towards this goal.
Recent work has demonstrated that well-engineered behavioral cloning methods can achieve competitive performance on diverse manipulation tasks [1, 2, 3].
Despite this promise, these imitation learning algorithms have been conﬁned to relatively small-scale domains.
For real-world problems, current imitation learning algorithms often demand a large number of task demonstrations, which can be difﬁcult and costly to obtain.
Realizing these limitations a series of recent work seeks to use prior interaction data to improve the sample efﬁciency of imitation learning on new tasks.
Such prior data come in various forms, including task-agnostic exploratory “play” data [4] or demonstrations previously collected for different tasks [5].
An open question is how best to extract knowledge from these large prior datasets and use this knowledge to facilitate learning novel tasks.
One promising approach is skill-based imitation learning [6, 7], which aims to learn a latent space of short-horizon sub-trajectories from the prior data (called skill learning ) and subsequently learn a policy to invoke the skills to solve a speciﬁc downstream task (called policy learning ).
This approach offers several appealing advantages: First, the policy beneﬁts from the temporal abstraction encapsulated by the skills, allowing the policy to focus on higher-level reasoning about what behavior to perform rather than how to execute that behavior; and second, by reasoning about the target task in relation to skills trained on prior data, the robot implicitly distills knowledge from the rich diverse interactions of the prior data into the policy.
Even so, existing skill-based imitation learning approaches bring modest improvements over simple behavioral cloning methods. This work aims at identifying the underlying limitations of existing approaches and designing new methods to address these limitations. 2. Retrieval policy learning
[Figure]
. Policy
[Figure] z Sample
We argue that a practical approach should encompass two key properties.
First, the latent skill space should serve as a predictable representation for downstream policy learning, allowing the policy to accurately infer which appropriate skill to use in new situations.
While prior methods [8, 9, 6, 7] have widely adopted variational autoencoders to transform sub-trajectories into embeddings, we empirically show that this learning objective alone is insufﬁcient for constructing a well structured latent space.
As a result, the policy tends to execute irrelevant skills.
Second, the robot should take advantage of the prior data to learn both the skills and the policy.
A major limitation of existing skill-based imitation learning approaches [6, 7] is that they primarily focus on using prior data for skill learning but not policy learning.
In these approaches, policies learned on a small number of target task demonstrations are prone to severe overﬁtting and covariate shift.
To address these limitations, we introduce S kill-
A ugmented I mitation L earning with pri o r R etrieval (SAILOR).
To improve the predictability of our skill representations, SAILOR uses an auxiliary temporal predictability objective to estimate the temporal distance between sub-trajectories in the demonstration sequences.
To use prior data for policy learning, we introduce a novel retrieval-based data augmentation procedure that selectively retrieves data relevant to the target task.
Speciﬁcally, we consider a subset of sub-trajectories from the prior dataset that have high latent skill similarity with sub-trajectories in the target task demonstrations (see Fig. 1).
We evaluate on a wide variety of manipulation tasks in simulation and the real world, and show that SAILOR signiﬁcantly outperforms state-of-the-art imitation learning and ofﬂine reinforcement learning approaches.
Through comprehensive analysis, we highlight the roles that prior data, our representation learning objective, and retrieval-based data augmentation have in data-efﬁcient learning of robust manipulation policies.
Related Work
Learning from Prior Data.
There is a large body of work on learning manipulation tasks using human demonstrations [1, 10, 11, 12, 13, 14]. While promising, most of these works learn tasks independently, without re-using knowledge from prior tasks.
As a result, they have high data requirements and exhibit brittle generalization in complex long-horizon tasks.
To address these limitations, several lines of work have investigated leveraging large ofﬂine prior datasets to facilitate learning downstream robotic tasks.
These prior datasets include task-agnostic play data [4, 15], demonstrations for related tasks [5], self-supervised agent-generated data [16], or a combination of these [17].
Alternatively to these, large video datasets are an appealing choice [18, 19, 20].
Recent work has leveraged these datasets to learn pre-trained visual representations for downstream control tasks [21, 22].
Yet despite their appeal, after experimenting with one such approach, R3M [21], we found that it can sometimes hinder downstream performance.
One hypothesis is that the prior data on which these methods are trained on exhibit signiﬁcant domain shift compared to downstream tasks, limiting transfer.
In this work, we instead consider large multi-task robotic datasets.
Multi-Task Imitation Learning.
Multi-task imitation learning methods offer a promising way to learn from diverse robotic data.
These include task-conditioned [5], language-conditioned [23, 4, 24, 25], and skill-based [6, 7] imitation learning.
While most task-conditioned and languageconditioned approaches seek to learn a single policy that performs well across a series of tasks, we focus on learning task-speciﬁc policies by utilizing a large multi-task prior dataset to learn a useful skill representation space, and supervising the task-speciﬁc policy using learned latent skills.
These multi-task imitation approaches could be complementary to our skill representation learning pipeline, ie. , we can consider incorporating language and task ID supervision into our approach. Skill-based Imitation Learning. We adopt skill-based imitation learning as the underlying framework for our method.
The objective is to learn temporally abstract representations of sensory-motor data (termed skills ) to enable more effective imitation learning. A long line of work has proposed learning skill representations by segmenting demonstrations into trajectory segments termed subtrajectories . Among these include work on segmenting demonstrations into variable-length subtrajectories, either in an unsupervised manner [26, 27, 28, 29, 30, 31, 32] or relying on additional supervision [33, 34, 35, 36]. Recent work has shown promise for encoding ﬁxed length sub-trajectories without any additional supervision [8, 9, 6, 7] using variational autoencoder-based approaches.
We adopt this setting due its relative simplicity and scalability. In contrast to these prior approaches (see Sec. 3), we enforce two key properties—a temporal predictability objective in our learned skill representation and a retrieval-based mechanism to improve task-speciﬁc policy learning.
Problem Formulation
Our goal is to leverage prior data to effectively learn novel target tasks in a data-efﬁcient manner.
Formally, we consider a target task as a Markov Decision Process M =
( S , A , r, p, p 0 , γ ) representing the state space, action space, reward function, transition probability, initial state distribution, and discount factor.
Our objective is to learn a policy that maximizes the discounted sum of rewards for the task.
To learn the policy, we assume access to a small ofﬂine dataset D target collected for the target task and a large ofﬂine dataset D prior either from previous related tasks or task-agnostic interactions.
These datasets consist of variable-length trajectories in the form { o 0 , a 0 , o 1 , · · · , o T i } with o i ∈
O denoting the observations and a i ∈
A denoting actions.
We highlight that D prior and D target may have signiﬁcant differences, i e., the two datasets may come from different environments, be collected by different human demonstrators, and demonstrate different tasks.
Skill-based Imitation Learning.
We employ a skill-based imitation learning framework [6, 7] which consists of two stages: skill learning and policy learning.
In the skill learning phase, we learn a skill embedding space Z ⊂ R d of ﬁxed-length sub-trajectories τ = { o 0 , a 0 , o 1 , · · · , o H − 1 , a H − 1 , o H } in D prior .
These skill embeddings serve as an abstract representation of the agent’s behavior and can be invoked to solve a range of downstream tasks.
A number of representation learning methods can be employed to learn the skill embeddings, spanning reconstruction-based methods and contrastive learning.
In the subsequent policy learning phase, we are given D target and our goal is to learn a policy π for the target task T .
The policy now emits skill embeddings z ∈ Z and during execution we decode z into a sequence of actions { a 0 , a 1 , · · · , a H − 1 } via a skill decoder model p ψ .
To learn the policy, prior work has proposed parametric [6] and semiparametric [7] polices that ﬁrst parse segments of D target into skills and subsequently learn to map observations in D target to these skill embeddings.
Note that we use distinct terms— skill and policy —to distinguish the role of these two components.
Skills represent short-horizon behaviors that can be re-used across many tasks, while the policy solves a speciﬁc long-horizon target task.
Skill-based Imitation Learning with Retrieval
In this section, we describe our skill-based imitation learning approach that can leverage prior multitask data to efﬁciently learn novel target tasks with a small amount of task-speciﬁc demonstrations.
As discussed in Sec. 3, this consists of two phases—a task-agnostic skill learning phase, where a latent skill space is learned using the prior data, and a task-speciﬁc policy learning phase where taskspeciﬁc data is used to learn a policy using the skills as supervision.
Compared to prior methods, our approach makes two key considerations—we (1) ensure that the learned skill space is a predictable [Table] [Table] [Table]
+d</latexit>
[Figure] [Figure] <latexit sha1_base64="RWcSzTuh7dmotVEhxCMTbgIjgQI=">AAAB8HicbVDLSgNBEOyNrxhfUY9eFhPBU9gNouIpoIccI5iHJEuYncwmQ+axzMwKYclXePGgiFc/x5t/4yTZgyYWNBRV3XR3hTGj2njet5NbW9/Y3MpvF3Z29/YPiodHLS0ThUkTSyZVJ0SaMCpI01DDSCdWBPGQkXY4vp357SeiNJXiwUxiEnA0FDSiGBkrPZZlP/Vu6tNyv1jyKt4c7irxM1KCDI1+8as3kDjhRBjMkNZd34tNkCJlKGZkWuglmsQIj9GQdC0ViBMdpPODp+6ZVQZuJJUtYdy5+nsiRVzrCQ9tJ0dmpJe9mfif101MdB2kVMSJIQIvFkUJc410Z9+7A6oINmxiCcKK2ltdPEIKYWMzKtgQ/OWXV0mrWvEvKxf31VLtLosjDydwCufgwxXUoA4NaAIGDs/wCm+Ocl6cd+dj0Zpzsplj+APn8weRKY+d [Figure]
/latexit
[Figure]
representation for downstream policy learning, and (2) improve the efﬁcacy of task-speciﬁc policy learning by retrieving task-relevant datapoints from the prior dataset.
See Fig. 2 for our model overview. 4.1 Learning a Predictable Representation of Skills We learn a skill representation by encoding sub-trajectories with a variational autoencoder (VAE) [37], and we further introduce an auxiliary objective to shape the representation.
Denoting a given sub-trajectory as τ , we employ a long short-term memory (LSTM) [38] encoder q φ that encodes τ into a Gaussian distribution over latent skills.
Our decoder is an LSTM network p ψ that, for each timestep t , decodes a latent z and the given observation o t into the reconstructed action ˆ a t .
We additionally employ a learned prior p θ to encourage sub-trajectories with similar starting and ending observations to have similar latent representations [15].
Our VAE loss objective is then
[Equation]
where β controls the effect of the KL divergence term [39] 1 .
It is important to highlight that action reconstruction is not the sole objective of our skill learning model—learning a consistent and predictable representation of behavior is critical for downstream policy learning, as shown by recent work [40, 41].
While the KL divergence term is one step towards this objective (by encouraging skills to be predictable given partial information from the subtrajectories), in this work we introduce an additional temporal predictability term that encourages the learned latent space to predict the temporal difference between two sub-trajectories.
Speciﬁcally, given two-sub-trajectories τ 1 and τ 2 from the same trajectory separated by t timesteps, we learn a model m ω to predict t given the corresponding skill mean embeddings of the trajectories:
[Equation]
where µ denotes taking the mean of the distribution.
We back-propagate L TP through the skill encoder model, allowing the this term to shape the learned skill representation.
Note that this is just one way to encourage temporal predictability, other objectives are also readily compatible with our method, such as time-contrastive networks [42].
Our overall objective is a weighted combination of the VAE and temporal predictability objectives:
[Equation]
[Figure]
Refer to Algorithm 1 for a detailed summary of our skill learning algorithm.
4.2 Retrieval-based Policy Learning
In the policy learning phase, we employ an LSTM policy that outputs the skill z to execute next.
We train the policy on a dataset D policy = { ( o ifs , z i = µ ( q φ ( τ i )) } , where τ i is an H -length subtrajectory, z i is the mean encoding of that sub-trajectory, and o ifs is the frame-stacked history of F observations preceding the sub-trajectory.
We train the policy to predict z i from o ifs using a standard behavioral cloning loss.
During execution, we roll out the LSTM skill decoder p ψ in a closed-loop manner, i.e., at each timestep we observe o t and execute the next action a t = µ ( p ψ ( z, o t )) .
After rolling out the skill for H timesteps we repeat the process by sampling a new skill from the policy.
A common approach to skill-based policy learning is training the policy with all H -length subtrajectories in D target. However, this limits the amount of supervision, especially when D target is small.
On the other hand, na ¨ ıvely training on all sub-trajectories in D target and D prior can hurt performance [43] due to divergent and conﬂicting behaviors between the prior and target datasets.
We thus introduce a retrieval-based mechanism to train on sub-trajectories in D prior that have high similarity with those in D target.
While many similarity metrics are suitable, in this work we measure similarity with respect to the skill embedding space—intuitively, sub-trajectories with similar skill embeddings demonstrate similar behaviors.
First we obtain skill embeddings of randomly sampled sub-trajectories in D prior and D target:
[Equation]
We then calculate the pairwise (cid:96) 2 distances between the prior and target dataset skill embeddings, i.e. , D[i][j] = (cid:107)
Z i prior − Z j target (cid:107) 2 .
Next, for each prior dataset skill embedding z i ∈
Z prior , we ﬁnd the closest corresponding target dataset skill, D_min[i] = min(D[i][:]) .
Finally, we retrieve the top- n sub-trajectories in D prior with the smallest distance argsort(D_min)[:n] , resulting in the retrieval dataset D ret.
We train the policy using the aggregated set of sub-trajectories in D ret and D target.
We use a behavioral cloning loss split across two terms: one for D target, and one for D ret weighted by a factor γ to control the effect of the retrieval data relative to the target dataset. At the same time as training the policy we additionally ﬁne-tune the skill model on D target. We summarize the retrieval, policy learning, and skill ﬁne-tuning steps in Algorithm 2. [Table]
Experiments 5.1 Simulated Experiment Setup
We perform empirical evaluations on two simulated robot manipulation domains (see Fig. 3):
Franka Kitchen [44]: A simulated kitchen environment involving different sub-tasks, such as opening cabinets, moving a kettle, and turning on a stove.
This environment comes with a large dataset of approximately 600 demonstrations performing various permutations of seven subtasks.
In this dataset, a subset of 18 demonstrations correspond to D target and demonstrate a speciﬁc permutation of subtasks: opening the microwave, followed by moving the kettle, ﬂipping on the light switch [Table] and opening the sliding cabinet.
We consider two prior datasets D prior : (1) using all demonstrations except the ones corresponding to the target task ( Kitchen-All ); and (2) using all demonstrations except those that involve interacting with the microwave ( Kitchen-No Microwave ).
These prior datasets have 584 and 235 demonstrations, respectively.
CALVIN [4]: A simulated tabletop playroom environment accompanies by a large dataset of taskagnostic “play” data with 2.3M transitions.
The play data encompass diverse behaviors, such as opening and closing drawers, turning on and off the lights, and picking, placing, and pushing blocks.
We use all play data as D prior to solve two target tasks.
The ﬁrst target task involves setting up the playroom environment in multiple stages ( CALVIN-Setting Up ).
Speciﬁcally, the robot must turn on the lights, and retrieve three blocks and place them on the table.
The second target task in contrast involves cleaning up the playroom environment ( CALVIN-Cleaning Up ).
Speciﬁcally, the robot must open the drawer, place all three blocks into the drawer, close the drawer, and turn off the lights.
For each task, we collect 30 demonstrations, which amounts to about half an hour of data collection.
The CALVIN domain is substantially more challenging than the Franka Kitchen domain, as the target tasks have a longer horizon and involve a greater number of objects.
Also, in contrast to Franka Kitchen, the prior and target datasets in CALVIN are collected by different human demonstrators who exhibit different styles of teleoperation.
We learn vision-based policies for both domains.
We refer readers to Appendix B for a detailed discussion of our tasks and datasets, and Appendix C.5 for environment implementation details.
5.2 Quantitative Analysis
We evaluate our method SAILOR against state-of-the-art imitation learning and ofﬂine reinforcement learning algorithms:
BC-RNN : behavioral cloning on D target without prior data We adopt
the
LSTM-based BC-RNN implementation in robomimic [ ], which has shown superior performance over other behavioral cloning approaches.
BC-RNN (FT) : BC-RNN variant that leverages prior data. We
ﬁrst pre-train BC-RNN on D prior and subsequently ﬁne-tune on D target .
This baseline aims to examine the effectiveness of supervised pre-training on interaction data for imitation learning.
BC-RNN (R3M) : behavioral cloning on D target using a frozen R3M visual representation [ pre-trained on the large-scale Ego4D video dataset [
This baseline intends to examine the effec - tiveness of using visual representations trained on natural images and videos
[Table]
IQL :
Implicit Q-Learning [45], a recent ofﬂine reinforcement learning method with state-of-the-art performance on the D4RL dataset [46]; trained on D target.
IQL (UDS) :
Implicit Q-Learning with Unlabeled Data Sharing [47, 48], which is a variant of IQL trained jointly on D prior and D target, where the transitions in D prior are labeled with the minimum reward ( 0 for our tasks). Singh et al.
[47] and Yu et al.
[48] show that this simple data augmentation procedure can effectively leverage prior data without additional rewards annotation.
FIST :
Few-shot Imitation with Skill Transition Models [7], an analogue of our method that employs a semi-parametric policy to select the skill to execute next.
This baseline uses the same underlying skill model as our method but a different policy learning scheme and does not involve retrieval.
Refer to Appendix C for additional implementation details on our model architecture, pseudocode, and evaluation protocols, and Appendix D for hyperparameter details.
We report performance of all methods in Table 1.
SAILOR greatly outperforms the baselines with an average task success rate of 89 . 0% .
Notably it outperforms the most competitive baseline BC-RNN (FT) by 12 .
BCRNN performs poorly as it fails to learn an effective policy from a small number of demonstrations.
In comparison, BC-RNN (R3M) shows signiﬁcant improvements on the Franka Kitchen tasks, but performs worse on the CALVIN tasks.
We hypothesize that this is due to the limited generalization ability of the pre-trained visual representations.
The ofﬂine reinforcement learning baselines show more promising results on the Franka kitchen tasks but struggle on the more challenging CALVIN tasks.
Finally, FIST signiﬁcantly under-performs our method. As FIST uses the same underlying skill model as our method, we attribute the limitations of FIST to its semi-parametric policy.
5.3 Ablation Study
We perform an extensive ablation study to understand the effects of various modeling choices on our method.
First, we study the effect of the temporal predictability term in Eq.
(2) on downstream task performance by removing it from Eq. (3) ( No TP ).
Next we study the role of retrieval by training the policy solely on sub-trajectories in D target ( No Retrieval ).
We also study the opposite case— training the policy on all of D prior and D target ( All Retrieval ).
Finally, we study the role of prior data on our method by training the skill and policy solely on D target ( No Prior Data ).
We also report ablations in Appendix A on the size of the prior, retrieval, and target datasets, in addition to the choice of retrieval method.
We present results in Table 2 for the more challenging CALVIN tasks.
First, we ﬁnd that both the temporal predictability objective and the retrieval mechanism have a signiﬁcant impact on the ﬁnal performance.
It is worth noting that removing these components makes our model degenerate into the skill-based imitation learning setting of OPAL [6].
In fact, these two ablations perform worse than the na¨ıve BC-RNN (FT) baseline for the CALVIN-Cleaning Up task, indicating that both an effective skill representation and the retrieval mechanism play a critical role in skill-based imitation learning.
The All Retrieval ablation also performs suboptimally—qualitatively we observe that the robot often was “distracted” and performed behaviors unrelated to the target task, likely due to the multimodal distributions in the prior data.
Finally, the No Prior Data ablation validates the role of prior data in learning effective skills and policy.
Comparing the No Retrieval ablation and the No Prior Data ablation, we see a 10 − 15% gap in performance, and this is attributed to the fact that the skills in the No Retrieval ablation are additionally trained on the prior data.
Despite the loss of performance in the No Prior Data ablation, we still see that it outperforms the BC-RNN baseline by a signiﬁcant margin.
We attribute this to the temporal abstraction afforded by our skill-based learning framework.
In sum, our ablation studies suggest that effective skill abstractions, coupled with mechanisms that effectively leverage the prior data, allow us to achieve strong results.
[Figure]
5.4 Real World Experiments
Finally, we showcase the efﬁcacy of our method in the real world with a kitchen environment involving eight food items, receptacles, a stove, and a serving area (see Fig. 4).
We ﬁrst collect a play dataset of exploratory interactions involving the food items and receptacles.
We consider three target tasks: (1) Real-Breakfast : setting up a breakfast table by placing the bread, butter, and milk in the serving area; (2) Real-Cook : cooking a meal by placing the ﬁsh, sausage, and tomato into the frying pan; (3) Real-Cook-Pan : a variant of the Real-Cook task involving placing the pan onto the stove.
We collect 30 demonstrations; refer to Appendix B.3 for detailed descriptions on the tasks datasets. We evaluate SAILOR against the most competitive baseline, BC-RNN (FT) (see Appendix C.3 for our evaluation protocol).
We ﬁnd that while on Real-Breakfast both methods achieve a success rate of 76.7%, on Real-Cook our method signiﬁcantly outperforms BC-RNN (FT) with a success rate of 73.3% vs. 23.3%, and similarly for Real-Cook-Pan (76.7% vs. 46.7%). To see the value of prior data we ran the No Prior ablation for the Real-Cook task, achieving 53.3% success rate compared to our 73.3%. Interestingly we see that the No Prior ablation largely outperforms the BC-RNN (FT) baseline on this task (23.3%) which had access to additional prior data. Overall, we observe that BC-RNN (FT) often failed to correctly grasp objects.
One hypothesis for this result is that pre-training stage biases the policy to learn the multi-modal behaviors in the prior dataset, preventing the policy from learning specialized target task behaviors during the ﬁne-tuning phase.
Limitations
While our method shows signiﬁcant promise, it leaves limitations that we hope to address in future work.
First, acquiring large amounts of multi-task prior data is difﬁcult and costly.
To amortize the high cost, large prior multi-task datasets should be useful in a diverse range of downstream tasks, rather than a handful.
In this work, we evaluate our method in a limited set of target tasks and leave it for future work to scale up the variety of tasks.
We hope (and believe) that the need for large robotic datasets will be addressed in the coming years [10, 23].
Second, our method is more computationally expensive than the BC-RNN baseline [1], due to the higher number of losses and networks used.
Third, our experiments focus on domains and datasets where the prior data and target tasks are reasonably close to each other. Notably our experiments do not evaluate generalization to unseen objects between the prior and target datasets.
It would be interesting to investigate methods that are tolerant to much larger domain shifts between prior and target task data.
Conclusion
We present SAILOR, a skill-based imitation learning framework for robot manipulation.
Our method uses prior data to construct a latent space of predictable and consistent skill representations.
It uses these latent skills as the temporal abstraction to learn policies for vision-based manipulation.
Key to its effectiveness is our newly designed representation learning objectives and retrieval-based data augmentation procedure.
We demonstrate that our method can solve long-horizon manipulation tasks in simulation and on physical hardware.
It brings forth a data-efﬁcient way of programming robots with new behaviors using a small number of target task demonstrations.
For future work, we plan to address the limitations we discussed in the previous section and investigate the effectiveness of this approach with various forms of prior data at different scales. Acknowledgments
We would like to thank Jake Grigsby, Huihan Liu, and Zhenyu Jiang for providing feedback on this manuscript.
We would also like to thank Yifeng Zhu for real robot infrastructure support.
We acknowledge the support of the National Science Foundation (1955523, 2145283), the Ofﬁce of
Naval Research (N00014-22-1-2204), and Amazon.