Hierarchical Few-Shot Imitation with Skill Transition Models
Authors: Kourosh Hakhamaneshi
Ruihan Zhao
Albert Zhan
Pieter Abbeel
Michael Laskin
Venue: ICLR 2022 Poster (2022)
Keywords: behavioral priors
skill extraction
imitation learning
few-shot learning

Abstract: A desirable property of autonomous agents is the ability to both solve long-horizon problems and generalize to unseen tasks. Recent advances in data-driven skill learning have shown that extracting behavioral priors from offline data can enable agents to solve challenging long-horizon tasks with reinforcement learning. However, generalization to tasks unseen during behavioral prior training remains an outstanding challenge. To this end, we present Few-shot Imitation with Skill Transition Models (FIST), an algorithm that extracts skills from offline data and utilizes them to generalize to unseen tasks given a few downstream demonstrations. FIST learns an inverse skill dynamics model, a distance function, and utilizes a semi-parametric approach for imitation. We show that FIST is capable of generalizing to new tasks and substantially outperforms prior baselines in navigation experiments requiring traversing unseen parts of a large maze and 7-DoF robotic arm experiments requiring manipulating previously unseen objects in a kitchen.

Under review conference paper at ICLR 2022
HIERARCHICAL FEW-SHOT IMITATION WITH SKILL TRANSITION MODELS
Anonymous authors Paper under double-blind review
ABSTRACT
A desirable property of autonomous agents the ability to both solve long-horizon prob- lems and generalize to unseen tasks. Recent advances in data-driven skill learning have shown that extracting behavioral priors from offline data can enable agents to solve chal- lenging long-horizon tasks with reinforcement learning. However, generalization tasks unseen during behavioral prior training remains an outstanding challenge. To this end, present Few-shot Imitation with Skill Transition Models (FIST), algorithm that extracts skills from offline data and utilizes them to generalize to unseen tasks given a few downstream demonstrations. FIST learns inverse skill dynamics model, a distance function, and utilizes a semi-paramerre approach for imitation. We show that FIST is capable of generalizing to new tasks and substantially outperforms prior baselines in navi- gation experiments requiring traversing unseen parts of a large maze and 7-DoF robotic arm experiments requiring manipulating previously unseen objects in kitchen.
INTRODUCTION
We are interested in developing control algorithms that enable robots to solve complex and practical tasks such as operating kitchens or assisting humans with everyday chores at home. There are two general characteristics of real-worl world tasks long-horizon planning and generalizability. Practical tasks are often long-horizon in the sense that they require robot to complete a sequence of subtasks. For example, to cook a meal a robot might need to prepare ingredients, place them in pot, and operate the stove before the full meal is ready. Additionally, in the real world many tasks we wish our robot to solve may differ from tasks the robot has completed in the past but require similar skill set. For example, if robot learned to open the top cabinet drawer it should be able to quickly adapt that skill to open the bottom cabinet drawer. These considerations motivate our research question: how can we learn skills that enable robots to generalize to new long-horizon downstream tasks?
Recently, learning data-driven behavioral priors has become a promising approach solving long- nogghorizonn tasks. Given a large unlabeled offline dataset of robotic demonstrations solving diverse set of tasks this family of approaches Singh et al. (2020); Pertsch al. (2020); Ajay et al. (2021) extract behavioral priors by fitting maximum likelihood expectation latent variable models to the offline dataset. The behavioral priors are then used to guide a Reinforcement Learning (RL) algorithm to solve downstream tasks. By selecting skills from the behavioral prior, the RL algorithm is able to explore in structured manner and can solve long horizon navigation and manipulation tasks. However, the generalization capabilities of RL with behavioral priors are limited since different RL agent needs to be trained for each downstream task and training each RI. agent often requires millions of environment interactions.
On the other hand, few-shot imitation learning has been a promising paradigm for generalization. In the few- shot imitation learning setting, an imitation learning policy is trained on an offline dataset of demonstrations and is then adapted in sho- to downstream task Duan al. (2017). Few-shot imitation learning has the

---

Under review conference paper at ICLR 2022
Offline dataset multi-task demonstrations
Extraction of from offline
adaptation skills
data
new tasks
Assembly
Pickup
Reach
Figure I: In this work we are interested in enabling autonomous robots to solve complex Long-horizon tasks that were unseen during training. To do so, assume access to large multi-task dataset of demonstrations, extract skills from the offline dataset, and adapt those skills new tasks that unseen during training.
added advantage over RL that is often easier for a human to provide handful of demonstrations than it is to engineer new reward function for downstream task. However, unlike RL with behavioral priors, few-shot imitation learning is most often limited to short-horizon problems. The reason is that imitation learning policies quickly drift away from the demonstrations due to error accumulation Ross et al. (201 Ib), and especially so in the few-shot setting when only handful of demonstrations are provided.
While it is tempting to simply combine data-driven behavioral priors with few-shot imitation learning, it is not obvious how to do so since the two approaches are somewhat orthogonal. Behavioral priors are trained on highly multi-modal datasets such that given state can correspond to multiple skills. Given sufficiently large dataset demonstrations for the downstream task the imitation learning algorithm will learn to select the correct mode. However, in the few-shot setting how do we ensure that during training on downstream data we choose the right skill? Additionally, due to the small sample size and long task horizon it is highly likely that naive imitation learning policy will drift from the few-shot demonstrations. How do we prevent the imitation learning policy from drifting away from downstream demonstrations?
The focus of our work is the setup illustrated in Figure 1; we introduce Few-Shot Imitation Learning with Skill Transition Models (FIST), new algorithm for few-shot imitation learming with skills that enables gener- alization unseen but semantically similar long-horizon tasks to those seen during training. Our approach addresses the issues with skill selection and drifting in the few-shot setting with two main components. First, we introduce an inverse skill dynamics model that conditions the behavioral prior not only on the current state but also future state, which helps FIST learn uni-modal future conditioned skill distribution that can then be utilized few-shot. The inverse skill model is then used as a policy to select skills that will take the agent the desired future state. Second, train distance function find the state for conditioning the inverse skill model during evaluation. By finding states along the downstream demonstrations that are closest to the current state, FIST prevents the imitation learning policy from drifting. We show that our method results in policies that are able to generalize to new long-horizon downstream tasks in navigation environments and multi-step step robotic manipulation tasks in kitchen environment. To summarize, we list our three main contributions:
1. We introduce FIST an imitation learning algorithm that learns an inverse skill dynamics model and distance function that used for semi-parametrie few-shot imitation.
2. We show that FIST can solve long-horizon tasks in both navigation and robotic manipulation settings that were unseen during training and outperforms previous behavioral prior and imitation learning baselines.
3. We provide insight into how different parts of the FIST algorithm contribute final performance by ablating different components of our method such as future conditioning and tue-tunnng on downstream data.
2

---

Under review as conference paper at ICLR 2022
2 RELATED WORK
Our approach combines ingredients from imitation learning and skill extraction to produce policies that can solve Long-horizon tasks and generalize to tasks that are out of distribution but semantically similar to those encountered in the training set. We cover the most closely related work in imitation learning, skill extraction, and few-shot generalization.
Imitation Learning: Imitation learning is a supervised learning problem where an agent extracts a policy from a dataset of demonstrations Billard et al. (2008); Osa et al. (2018). The two most common approaches to imitation are Behavior Cloning Pomerleau (1988); Ross et al. (201 Ia) and Inverse Reinforcement Learning (IRL) Ng & Russell (2000). BC approaches learn policies Ty(a(s) that most closely match the state- conditioned action distribution of the demonstration data. IRL approaches learn reward function from the demonstration data assuming that the demonstrations are near-optimal for a desired task and utilize Reinforcement Learning to produce policies that maximize the reward. For simplicity and to avoid learning a reward function, in this work we aim to learn generalizable skills and using the BC approach. However, two drawbacks of BC are that the imitation policies require a large number of demonstrations and prone to drifting away from the demonstration distribution during evaluation due to error accumulation Ross et al. (201 1b). For this reason, BC policies work best when the time-horizon of the task is short.
Skill Extraction with Behavioral Priors: Hard- Hard-coding prior knowledge into a policy or dynamics model has been considered as solution to more sample efficient learning, especially in the context of imitation learning Chatzilygeroudis et al. (2019); Bahl et al. (2020); Stulp & Sigaud (2013). For example Stulp & Sigaud (2013) utilizes dynamic movement primitives instead of the raw action space to simplify the learning problem and improve the performance of evolutionary policy search methods. While imposing the structural prior on policy or dynamic models can indeed improve few-shot learning, our method is complementary to these works and proposes to learn behavioral priors from a play dataset. Methods that leverage behavioral priors utilize offline datasets of demonstrations to bias a policy towards the most likely skills in the datasets. While related closely to imitation learning, behavioral priors have been mostly applied to improve Reinforcement Learning. Behavioral priors learned through maximum likelihood latent variable models have been used for structured exploration in RL Singh et al. (2020), to solve complex long-horizon tasks from sparse rewards Pertsch et al. (2020), and regularize offline RL policies Wu et al. (2019); Peng et al. (2019); Nair et al. (2020). While impressive, RL with data-driven behavioral priors does not generalize to tasks efficiently, often requiring millions of environment interactions to converge to an optimal policy for a new task.
Few-Shot Learning: Few-shot learning Wang et al. (2020) has been studied in the context of image recog- nition Vinyals et al. (2016); Koch et al. (2015), reinforcement learning Duan et al. (2016), and imitation learning Duan et al. (2017). In the context of reinforcement and imitation learning, few-shot learning is often cast as a meta-learning problem Finn et al. (2017); Duan et al. (2016; 2017), with often specialized datasets of demonstrations that labeled by tasks. However, there are other means of attaining few-shot generalization and adaptation that do not require such expensive sources of data. For instance. Cully et al. (2015) and Pautrat al. (2018) use Bayesian Optimization for transfer learning to a possibly damaged robot in real-world via learned priors based on big repertoires of controllers in different settings from simulation. However, our problem of interest is skill adaptation that requires no further interaction with the downstream environment during few-shot imitation.
Recently, advances in unsupervised representation learning in natural language processing Radford et al. (2019); Brown et al. (2020) and vision He et al. (2020); Chen et al. (2020) have shown how network pre-trained with a self-supervised objective can be finetuned or adjusted with a linear probe to generalize in few-shot or even zero-shot Radford et al. (2021) to downstream task. Our approach to few-shot imitation learning is loosely inspired by the generalization capabilities of networks pre-trained with unsupervised objectives. Our approach first fits behavioral prior unlabeled offline dataset of demonstrations extract skills and then fits an imitation learning policy over the previously acquired skills to generalize in few-shot to new tasks.
3

---

Under review conference paper at ICLR 2022
(A) Skil Extraction
(b) Stil Adaptation
(d Few Shot imitation
Decoder
Few shot trajactories
Un- Un-labeledd
Few-shot demos
offline dataset
H steps
Skill Encoder
mijn diy
X
LSTM
LSTM
MLP
io-e(e (ns ((r
Inverse Shill Dynamics g(a (5,5))
Figure 2: Our algorithm Few-Shot Imitation Learning with Skill Transition Models (FIST) is composed of three parts: (a) Skill Extraction: we fit a skill encoder, decoder, inverse skill dynamics model, and distance function to the offline dataset; Skill Adaptation: For downstream task, given few demonstrations and adapt the skills learned in (a), by fine-tuning encoder, decoder, and the inverse model. (c) Few-Shot Imitation: finally, to imitate the downstream demonstrations, we utilize the distance function to perform a look ahead along the demonstration to condition the inverse model and decode action.
3 APPROACH
3.1 PROBLEM FORMULATION
Few-shot Imitation Learning: We denote a demonstration as sequence of states and actions:
few-shot setting we assume access small dataset of M such expert demonstrations D {Tx JEM that fulfill specific long horizon task in the environment. For instance sequence of sub-tasks in a kitchen environment such as moving the kettle, turning on the burner and opening a cabinet door. The goal is imitate this behavior using only few available example trajectories.
{s1, 01, $2, 02 ..2,,............... ................................
-..,.-...77 ST, ).).
Skill Extraction: In this work we assume access to an unlabeled offline dataset of prior agent interactions with the environment in the form of N reward-free trajectories {T; = {(*,, at) Ji=N JEY. We further assume that these trajectories include semantically meaningful skills that are composable to execute long horizon tasks in the environment. This data can be collected from past tasks that have been attempted, or be provided by ouman-experts through teleoperation Zhang et al. (2018).
Skill extraction refers to an unsupervised learning approach that utilizes this reward-free and tusk-agnostic dataset leam skill policy in form of oa(als, (), where is action, is the current state, and is the skill. Our hypothesis is that by combining these skill primitives we can solve semantically similar long-horizon tasks that have directly been seen during the training. In this work propose new architecture for skill extraction based on continuous latent variable models that enables a semi-parametric evaluation procedure for few-shot imitation learning.
3.2 HIERARCHICAL TEW-SHOT IMITATION WITH SKILL TRANSITION MODELS
Our method, shown in Fig. 2, has three components: (i) Skill extraction, (ii) Skill adaptation via fine-tuning on few-shot data, and (iii) Evaluating the skills using semi- paramerreet approach to enable few-shot imitation.
4

---

Under review conference paper at ICLR 2022
(i) Skill Extraction from Offline Data: We define continuous skill 2, Z as an embedding for sequence of state-action pairs {st, Qtyy ...., -1+ with fixed length H. This temporal abstraction of skills has proven to be useful in prior work Pertsch et al. (2020); Ajay al. (2021), by allowing hierarchical decomposition of skills to achieve long horizon downstream tasks. To learn the latent space 2 we propose training a continuous latent variable model with the encoder as qo dt....., 8t+H-11 (++ and the decoder as (ala, 2). The encoder outputs a distribution over the latent variable z that best explains the variation in the state-action pairs in the sub-trajectory.
The encoder is an LSTM that takes in the sub-trajectory of length H and outputs the parameters of Gaussian distribution as the variational approximation over the true posterior p(z[st, decoder is policy that maximizes the log- seglikelihhod of actions of the sub-trajectoyy conditioned on the current state and the skill. We implement the decoder as a feed-forward network which takes the current state 8, and the latent vector and regresses the action vector directly. This architecture resembles prior works on skill extraction Pertsch et al. (2020).
(++H-1). The
To learn parameters Φ and 0, randomly sample batches of H-step continuous sub-trajectories from the training data D and maximize the evidence lower bound (ELBO):
logp(a)(s)) 2 E log To (uelss, z) +B (logp(=) - log do (=|T)
where the posterior (=(r) is regularized by its Kullback-Leible (KL) divergence from a unit Gaussian prior p(2) = (0, I) and ,3 is parameter that tunes the regularization term Higgins et al. (2016).
To enable quick sho-shot adaptation over skills we learn an inverse skill dynamics model go (2( St-H-1) that infers which skills should be used given the current state and future state that is H steps away. To train the inverse skill dynamics model we minimize the KI. divergence between the approximated skill posterior (=([r) and the output of the state conditioned skill prior. This will result in minimizing the following loss with respect to the parameters v:
Cprio (v) E,-D (9kz(90(1))) S1HH-
(2)
We use reverse KL divergence to ensure that inverse dynamics model has broader distribution than the approximate posterior to ensure mode coverage Bishop (2006). In our implementation we use a feed-forward network that takes in the concatenation of the current and future state and outputs the parameters of Gaussian distribution over z. Conditioning on the future enables us to make a more informative decision on what skills execute which is key enabler to few-shot imitation. We jointly optimize the skill extraction and inverse model with the following loss:
((0,0,0) Croe 0) + (0) Cptior())
(ii) Skill Adaption via Fine-tuning on Downstream Data downstream demonstrations and the prior over skills, we use the demonstrations to fine-tune the parameters of the architecture by taking gradient steps over the loss in Equation 3. In the experiments we ablate the performance of FIST with and without fine-tuning to highlight the differences.
To improve the consistency between the unseen
(iii) Sema-etram Evaluation for Few-shot Imitation Learning: To the agent, we need to first sample a skill ~ (2]s,, ii+r) based on the current state and the future state that it seeks to reach. Then, we can use the level decoder π (aas $t) to convert that sampled skill 2 and the current state 8; to the corresponding action a, During evaluation use the demonstrations Deemo to decide which state to use
5

---

Under review conference paper at ICLR 2022
the future state condition on. For this purpose we use learned distance function d(s, to measure the distance between the current state and every other state in the demonstrated trajectories. Then, from the few-shot data we find the closest state sy to the current state according to the distance metric:
$5 = min d(s1, 865)
where 8ij the 7th state in the 7th trajectory in Ddemo We then condition the inverse dynamics model on the current state and the state +H+ H steps ahead of sj within the trajectory that & belongs to. If by adding H steps we reach the end of the trajectory, we use the end state within the trajectory as the target future state. The reason for this look-ahead adjustment is to ensure that the sampled skill always makes progress towards the future states the demonstration. After the execution action at according to the low-level decoder, the process is repeated until the fulfillment of the task. The procedure is summarized in Algorithm 1.
Algorithm 1 FIST: Evaluation Algorithm
1: Inputs: Fine-tuned inverse skill dynamics model qu (2|ss, St+H-1), fine-tuned skill policy (e(als 2), learned distance function d(s, s'), downstream demonstration Ddemo
2: Initialize the environment to 80
3: for each t [1... T do
4: Pick sy" Look AoedAhhammminn..m ge d(8t,
5: Sample skill (2
6: Sample action a ~ (a (a)a81,
7: 8t < env.step(a)
3.3 LEARNING THE DISTANCE FUNCTION
To enable few-shot imitation learning we use learned distance function to search for goal state that is "close" the current state of the agent. Learning contrastive distance metrics has been successfully applied in prior work both in navigation and manipulation experiments. Both Liu et al. (2020) and Emmons et al. (2020) utilize contrastive distance functions to build semi- coni-parametric topological graphs, and Savinov et (2018) and Shah et al. (2020) use similar approach for visual navigation. Inspired by the same idea, we also use a contrastive loss such that states that are H steps the future are close to the current state while all other states are far. We refer the reader to the Appendix A.1 for further details.
4 EXPERIMENTS
In the experiments we interested in answering the following questions: (i) Can method successfully imitate unseen long-horizon downstream demonstrations? (ii) What is the importance of semi-parametric approach vs. future conditioning? (iii) Is pre-training and fine-tuning the skill embedding model necessary for achieving high success rate?
4.1 ENVIRONMENTS
We evaluate the performance FIST on two simulated navigation environments and robotic manipulation task from the D4RL benchmark shown in Figure 3. To ensure generalizability to out-of-distribution tasks we remove some category of trajectories from the offline data and at test-time, we see if the agent can generalize to those unseen trajectories. Specifically for Pointmass and Ant environments block some
6

---

Under review conference paper at ICLR 2022
Maze Environments
Kitchen Environment
Bottom
Right
Left
Turn Top
Burner
Microwave Light Door
Switch
Skill Data
Kettle
Downstream Demo
E
Figure 3: Top: In cach environment, we block some part of the environment and collect task-agnostic and reward-free trajectories for extracting skills. In the kitchen environment, red markers indicate the objects that excluded. Bottom: For downstream demonstrations, unseen objects.
expert trajectories that involve parts of the maze manipulation of
regions of the maze and test-time provide demonstration that move into the excluded region; and for Kitchen we exclude interactions with few selected objects and at test-time we provide demonstrations that involve manipulating the excluded object. Further details on the environment setup and data collection procedure is provided in appendix B.
4.2 RESULTS
We use the following approaches for comparison: BC+FT: Trains behavioral cloning agent (i.e. To (als)) on the offline dataset 'D and fine-tunes the downstream dataset D SPiRL: This is an extension of the existing skill extraction methods to imitation learning over skill space Pertsch et al. (2020). The skill extraction method SPiRL Pertsch et (2020) is very similar to FIST, instead conditioning the skill prior the future state it only uses the current state. To adapt SPiRL imitation learning, after pre-training the skills module, we fine-tune tune it on the downstream demonstrations Delemo (instead of finetuning with RL as proposed in the original paper). After fine-tuning execute the skill prior for execution. FIST (ours): This runs our semi-paramen parammarmetricc approach after learning the future conditioned skill prior. After extracting skills from D we fine-tune the parameters the downstream demonstrations Delmmo and perform the proposed semi-parametric approach for evaluation.
Figure 4 compares the normalized average presents either an excluded region (for Maze) or an excluded object (for Kitchen). The un-normalized scores are included in tables and 7 in Appendix C. Here, we provide summary of our key findings !
each unseen task from each domain. Each tick x-axis
(i) In the PointMaze environment, FIST consistently succeeds in navigating the point mass into all three goal locations. The skills learned by SPiRL fail generalize when the point mass falls outside of the training distribution, causing it to get stuck in corners. While BC+FT also solves the task frequently in the Left and Bottom goal location, the motion of the point mass sub-optimal, resulting in longer episode lengths.
'Tee video of of experiments included in the supplementary material.
7

---

Under review conference paper at ICLR 2022
Jears 3PL
05
ArtMnn
Kichen
Figure 4: Normalized Reward of our environments, and their excluded regions / ohjects. For maze, the episode length subtracted from the maximum episode length, which then divided by maximum episode length. For kitchen, the reward is the number of sub-tasks completed in order, and normalized by the maximum of 4.0
(ii) In the AntMaze environment, FIST achieves the best performance compared to the baselines. SPiRL and BC+FT make no progress in navigating the agent towards the goal while FIST is able to frequently reach the goals in the demonstrated trajectories. We believe that the low success rate numbers in this experiment is due to the low quality of trajectories that exist the offline skill dataset D. In the dataset, we see many episodes with ant falling over, and FIST's failure cases also demonstrate the same behavior, hence resulting in low success rate. We hypothesize that with a better dataset FIST can achieve a higher success rate number.
(iii) In the kitchen environment, FIST significantly outperforms SPiRL and BC+FT. FIST successfully complete out of 4 long-horizon object manipulation tasks in the same order required by the task. In one of these long-horizon tasks all algorithms perform similarly poor. We believe that such behavior is due the fact that the majority of the trajectories in the pre-training dataset start with microwave task and removing this sub-task can substantially decrease the diversity and number of trajectories seen during pre-training.
4.3 ABLATION STUDIES
In this section we study different components of the FIST algorithm to provide insight on the contribution of each part. We provide further ablations Appendix C.2.
Effect of the Semi-parameteriy evaluation and future conditioning of the skill prior. We recognize two major differences between FIST and SPiRL: (1) future conditioning of the skill prior during skill extraction (2) the semi-paramteric approach for picking the future state to condition on. In order to separately investigate the contribution of each part we performed an ablation with the following modifications of FIST and SPiRL the kitchen environment: FIST (Euc.) is similar distance on the raw state space, rather than our learned contrastive distance for lookup. Comparison to this baseline measure the importance of contrastive distance method, compared simpler alternative; SPiRL (closest) uses the contrastive distance function to look-up the closest state the demonstrations to sample the skill via Sclosest:: ) and SPiRL (H-step) uses the same H-step-ahead look-up approach as FIST and samples skill based on p(z Selosest++
FIST experiment, except that
the Euclidean
Table 1: Ablation on the impact of semi parameerrttt approach and the future conditioning of skill prior
Task (Unseen)
(ours)
(Boc.)
(closest)
SpiRL (H-steps)
Microwave, Kettle, Top Burner, Light Switch
3.6 1: 0.16
0.15
21 0.48
0.0 0.0
0.2 0.13
Microwave, Bottom Bumner, Light Switch, Slide Cabinet
2.3 0.5
2.8 0.42
2.3 0.49
2.8 a 0.44
0.0 0.0
Microwave, Kettle, Slide Cabúnet, Cabinet
3.0=0032
-92
0.0 0.0
0.0
Micnoaae, Kettle, Slide Cabinet, Hinge Cabinet
0.14
33:0338
1,5
8

---

Under review conference paper at ICLR 2022
Comparison of SPiRL results suggests that the performance would degrade if the semi- cemi-parameteric lookup method is used to pick the conditioning state in SPIRL. This is not surprising since the model in SPiRL is trained condition on the current state and is unable to pick a good skill based on p(2) (ctosss)) in SPiRL (closest) or p(æ Sclosest++ H1 in SPiRL (H-steps). Therefore, it is crucial to have the prior p(z]st, 8c+ so that we can condition on both the current and future goal state.
We also note that FIST (Eue.) results are slightly worse in 3 out of 4 cases. This suggests that (1) contrastive distance is better than Euclidean to some degree (2) the environment's 's state space is simple enough that Euclidean distance still works certain extent. It is worth noting that the Euclidean distance, even with comparable results to the contrastive distance in the kitchen environment, is not a general distance metric to and would quickly lose advantage when state representation gets complex (e.g. pixel representation).
The effect of skill pre-training and fine-tuning on FIST: To adjust the pre-trained skill-set to OOD tasks (e.g. moving the kettle while it is excluded from the skill dataset) FIST requires fine-tuming demonstrations. We hypothesize that without fine-tuning, the agent should be able to perfectly imitate the demonstrated sub-trajectories that it has seen during training, but should start drifting away when encountered with an OOD skill. We also hypothesise that pre-training on a large dataset, even if it does not include the downstream demonstration sub-trajectories, is crucial for better generalization. Intuitively, pre-training provides a behavioral prior that is easier to adapt to unseen tasks than a random initialization.
the downstream
To examine the impact of fine-tuning, we compare FIST with FIST-Ro-FT which directly evaluates the cemi-parameteric approach with the model parameters trained on the skill dataset without fine-tuning on the downstream trajectories. To understand the effect of pre-training. we compare FIST with FIST-no-pretren ---pttrain which is not pre-trained on the skill dataset. Instead, we directly train the latent variable and inverse skill dynamics model on the downstream data and perform the semi-parametric evaluation of the FIST algorithm.
From the results in Table 2, we observe that fine-tuning is critical component for adapting to OOD tasks. The scores on FIST-no-FT indicate that the agent is capable fulfilling the sub-tasks seen during skill pre-training but cannot progress onto unseen tasks without fine-tuning, Based on the scores on FIST-mo-pretrain, we also find that the pre-training on a rich dataset, even when the downstream task is directly excluded, provides sufficient prior knowledge about the dynamics of the environment and can immensely help with generalization to unseen tasks via fine-tuning.
Table 2: We ablate the use of pre-training on offline data, as well as fine-tuning on downstream demonstrations. FIST-no- FT removes the fine-tuning tine-tuning on downstream demonstration step in FIST, while FIST- FIST-no-pretrain trains the skills purely from the given downstream data. Without seeing the subtask, FIST-no-FT is unable to solve the downstream subtask. Trained on only downstream data, TIST--o--pretie
unable to properly manipulate the robot.
Task (Unseen)
FIST (ours)
FIST-no-FT
FIST-no-pretrain
Mierowave, Kettle, Top Burner, Light Switch
3.6 0.16
2.0 0.0
0.5 + 0.16
Microwave, Bottom Burner. Light Switch, Slide Cabinet
0.5
0,0 0.0
0.7 : 0.15
Mierowave, Kettle, Slide Cabinet, Hinge Cabinet Microwave, Kettle, Slide Cabinet, Hinge Cabinet
0.3
1.0
0.0 0.0
4.0
2. 0.0
0.8 1 0.13
5 CONCLUSION
We present FIST, semi parametric algorithm for few-shot imitation learning for long- mori-horizon tasks that are unseen during training. We use previously collected trajectories of the agent interacting with the environment learn set of skills along with inverse dynamics model that is then combined with a non-parametric approach to keep the agent from drifting away from the downstream demonstrations. Our approach is able to solve long-horizon challenging tasks in the few-shot setting where other methods fail.
9

---

Under review conference paper at ICLR 2022
REFERENCES
Anurag Ajay, Aviral Kumar, Pulkit Agrawal, Sergey Levine, and Ofir Nachum. (OPAL): Offline primitive
discovery for accelerating offline reinforcement learning. In International Conference on Learning Representations, 2021. URL https: /ttp://////p/iiie/eeeeeeeee net fert/forumiii ftte: d=V69T.GW001 TN.
Shikhar Bahl, Mustafa Mukadam, Abhinav Gupta, and Deepak Pathak. Neural dynamic policies for end-to-end sensorimotor learning. arXiv preprint arXiv:2012.02788, 2020.
Aude Billard, Sylvain Calinon, Ruediger Dillmann, and Stefan Schaal. Survey: Robot programming by demonstration. Technical report, Springrer, 2008.
Christopher M Bishop. Pattern recognition and machine learning. springer. 2006.
Tom B Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. arXiv preprint arXiv:2005. 14165, 2020.
Konstantinos Chatzilygeroudis, Vassilis Vassiliades, Freek Stulp. Sylvain Calinon, and Jean- Saan-Baptiste Mouret. survey on policy search algorithms for learning robot controllers in handful of trials. IEEE Trunsactions on Robotics, 36(2):328-344 347, 2019.
Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. In International conference machine learning, 2020.
Antoine Cully, Jeff Clune, Danesh Tarapore, and Jean-Baptiste Mouret. Robots that can adapt like animals. Nature, 521(7553):503-500) 2015.
Yan Duan, John Schulman, Xi Chen, Peter L. Bartlett, Πya Sutskever, and Picter Abbeel. RI2: Fast reinforcement learning via slow reinforcement learning. arXiv:1611.02279, 2016.
Yan Duan, Marcin Andrychowicz, Bradly C. Stadic, Jonathan Ho, Jonas Schneider, Ilya Sutskever, Pieter Abbeel, and Wojeiech 7aremba. One-shot imitation learning. In Isabelle Guyon, UI- rike von Luxburg, Samy Bengio, Hanna M. Wallach, Rob Fergus, S. V. N. Vishwanathan, and Roman Garnett (eds.), Advances in Neural Information Processing Systems 30: Annual Confer- ence on Neural Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA, pp. 1087-1098, 2017. URL httpt:t/ttppeeeeeeeeeeeeeee http///eoooceeeeeeeeeeeeeepce.ee neurips cc wwwpppape//00// 2017 / wwt.:995555eee0eeaaaaaaaaaaaaaaa -bsttrract html.
Scott Emmons, Ajay Jain, Michael Laskin, Thanard Kurutach, Pieter Abbeel, and Deepak Pathak. Sparse graphical memory for robust planning. In Hugo Larochelle, Mare' Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan-Tien Lin (eds.), Advances Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual, 2020. URL httpa: / pttpe/edoccceeeeeeee/eeeeeee www.nnuuuuccmmmmm http:pp.cc/ppppppp/p 2020 haahh/ 385822e359ate seeaaaaaaaaa22bb22222222222 -bt. html.
Chelsea Finn, Tianhe Yu, Tianhao Zhang, Pieter Abbeel, and Sergey Levine. One-shot visual imitation learning via meta-learning In 1st Annual Conference Robot Learning, CoRL 2017, Mountain View, California, USA, November 13-15, 2017, Proceedings, volume 78 of Proceedings of Machine Learning Research, pp- 357-368. PMLR, 2017. URL httt:tt///cceeeeiiiee/
mlr wwwps/esssssee....p v finni 7a html.
Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4rl: Datasets for deep data-driven reinforcement learning. arXiv preprint arXiv:2004.07219, 2020.
10

---

Under review conference paper at ICLR 2022
Abhishek Gupta, Vikash Kumar, Corey Lynch, Sergey Levine, and Karol Hausman. Relay policy learning: Solving long-horizon tasks via imitation and reinforcement learning, arXiv preprint arXiv: 1910. 1956, 2019.
Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross Girshick. Momentum contrast for unsupervised visual representation learning. In Conference on Computer Vision and Pattern Recognition, 2020.
Irina Higgins, Loic Matthey, Arka Pal, Christopher Burgess, Xavier Glorot, Matthew Botvinick, Shakir Mohamed, and Alexander Lerchner. beta-vae: Learning basic visual concepts with constrained variational framework. 2016.
Gregory Koch, Richard Zemel, and Ruslan Salakhutdinov. Siamese neural networks for one-shot image recognition. 2015.
Kara Liu, Thanard Kurutach, Christine Tung, Pieter Abbeel, and Aviv Tamar. Hallucinative topological mem- ory for zero-shot visual planning. In Proceedings the the 37th International Conference on Machine Learning, ICML 2020, 13-18 July 2020, Virtual Event, volume 119 of Proceedings of Machine Learning Research, pp. 6259-6270, PMLR, 2020. URL http: /ttpc//eeoceieooeeoooooooooooc ngs ml wtt.:.aww.........mmm ppw.s//evvlllv 9/1iu20h, html.
Ashvin Nair, Abhishek Gupta, Murtaza Dalal, and Sergey Levine. Awac: Accelerating online reinforcement learning with offline datasets, 2020.
Andrew Y. Ng and Stuart J. Russell. Algorithms for inverse reinforcement learning. In Pat Langley (ed.), Proceedings of the Seventeenth International Conference on Machine Learning (ICML 2000), Stanford University, Stanford, CA, USA, June 29 July 2, 2000, pp. 663-670. Morgan Kaufmann, 2000.
Takayuki Osa, Joni Pajarinen, Gerhard Neumann, J Andrew Bagnell, Pieter Abbeel, and Jan Peters. An algorithmic perspective on imitation learning. arXiv preprint arXiv:1811.0611, 2018.
Rémi Pautrat, Konstantinos Chatzilygerouddss, and Jean-Bapticn -aptistt Mouret. Bayesian optimization with auto- matic prior selection for data-efficient direct policy search. In 2018 IEEE International Conference on Robotics and Automation (ICRA), pp.
IEEE, 2018
Xue Bin Peng, Aviral Kumar, Grace Zhang, and Sergey Levine. Advantage- regression: Simple and scalable off- off-policy reinforcement learning. CoRR, abs/1910.00117, 2019. URL http: http://axxvvoooooo/ abs 1910.91100 00177.
Karl Pertsch, Youngwoon Lee, and Joseph J. Lim. Accelerating reinforcement learning with learned skill priors. In Conference Robot Learning (CoRL), 2020.
Dean Pomerleau. ALVINN: autonomous land vehicle in neural network. In David S. Touretzky (ed.), Advances in Neural Information Processing Systems 1, (NIPS Conference, Denver, Colorado, USA, 1988], pp. 305-313. Morgan Kaufmann, 1988. URL http: / hapers.nipa cc/paper/ 95-alvinn an- autooooooos--a-rd--eiiiaaaa -narrre--nn----errraatttttttttto -newworkk...
Alee Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models are unsupervised multitask learners. 2019.
Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, Gretchen Krueger, and Ilya Sutskever. Learning transferable visual models from natural language supervision, 2021.
11

---

Under review conference paper at ICLR 2022
Stéphane Ross, Geoffrey J. Gordon, and Drew Bagnell. A reduction of imitation learning and structured prediction to no-regret online learning. In Geoffrey J. Gordon, David B. Dunson, and Miroslav Dudík (eds.). Proceedings of the Fourteenth International Conference on Artificial Intelligence and Statistics, AISTATS 2011, Fort Lauderdale, USA, April 11-13, 2011, volume 15 of JMLR Proceedings, pp. 627-635. JMg.. 2011a. URL http: / ppwpeedinggggggooooooooo mlr, www......ee...... / vv 5 rosssla/ rasiaa/rooslllll pd..
Stéphane Ross, Geoffrey J. Gordon, and Drew Bagnell. A reduction of imitation learning and structured prediction to no-regret online learning. In Geoffrey J. Gordon, David B. Dunson, and Miroslav Dudík (eds.), Proceedings of the Fourteenth International Conference on Artificial Intelligence and Statistics, AISTATS 2011, Fort Lauderdale, USA, April 11-13, 2011, volume 15 of JMLR Proceedings, pp. 627-635. JMLR.org, 2011b. URL http:tttppppoceeeeeeeeee http///ooceeen...e//eeeepoe.. ngg. mlr. arl.. prwss/eessveeee vv 5 rros 1a
ross laossslll.. aal.la.pplll pdd..
Nikolay Savinov, Alexey Dosovitskiy, and Vladlen Koltun. Semi-parametric topological memory for navigation. In 6th International Conference on Learning Representatiinss ICLR 2018, Vancouver, BC, Canada, April 30 May 3, 2018, Conference Track Proceedings. OpenReview.net, 2018. URL https:
//openreview. net/ forumm fo SygwwGbRR..
Dhruv Shah, Benjamin Eysenbach, Gregory Kahn, Nicholas Rhinehart, and Sergey Levine. Ving: Learning open- navigation with visual goals. arXiv preprint arXiv:2012.09812. 2020.
Avi Singh, Huihan Liu, Gaoyue Zhou, Albert Yu, Nicholas Rhinehart, and Sergey Levine. Parrot: Data-driven behavioral priors for reinforcement learning. In International Conference on Learning Representations, 2020.
Freek Stulp and Olivier Sigaud. Robot skill learning: From reinforcement learning to evolution strategies. Paladyn, Journal of Behavioral Robotics, 4(1):49-61, 2013.
Aaron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predictive coding. arXiv preprint arXiv: 1807.03748, 2018.
Oriol Vinyals, Charles Blundell, Tim Lillicrap, Koray Kavukcuoglu, and Daan Wierstra. Matching net- works for one shot learning. In Daniel D. Lee. Masashi Sugiyama, Ulrike von Luxburg, Isabelle Guyon, and Roman Garnett (eds.). Advances in Neural Information Processing Systems 29: An- nual Conference on Neural Information Processing Systems 2016, December 5-10, 2016, Barcelona, Spain, pp. 3630-3633600 3638, 2016. URL https: / http://eooccengss//ee/ee/ inw.rrreedinnnnnnnrirrr........e www.ins..uu..ee.. 30e13578336567
wttplpssscppppppppe.epe 201 6/hash/ 212661222600 www.eeeeeeeeeee.0 88 Awwoce-bbssttttttae html.
Yaqing Wang, Quamming Yao, James Kwok, and Lionel M. Ni. Generalizing from a few examples: survey few-shot learning. ACM Comput. Surv., 53(3):63:1-633:34 2020. doi: 10.1145/338625.. URL
https: http///o..og/..////////o http///1.111111111555515
Yifan Wu, George Tucker, and Ofir Nachum. Behavior regularized offline reinforcement learning. CoRR, abs/1911.11361, 2019. URL http: //tp/v/axvvvv/v/ org/ abs /11//.91.... 11361.
Tianhao Zhang, Zoe McCarthy, Owen Jow, Dennis Lee, Xi Chen, Ken Goldberg, and Pieter Abbeel. Deep imitation learning for complex manipulation tasks from virtual reality teleoperation. In 2018 IEEE International Conference on Robotics and Automation, ICRA 2018, Brishane, Australia, May 21-25, 2018, 8.6614944.4 URL https: /tts:://dooooooo 1109/ 109//CCA. ICRA.
pp. 1-8. IEEE, 2018. doi: 2018. 8461249.
12

---

Under review conference paper at ICLR 2022
A IMPLEMENTATION DETAILS
A.1 DISTANCE FUNCTION
As mentioned in Section 3, we wish to learn an encoding such that our distance metric is the euclidean distance between the encoded states.
d(s,s) [1(ss) - h(s))11
(5)
To learn the encoder h, we optimize a contrastive loss on encodings the current and future states along the same trajectory. At each update step have batch of trajectories 71,--? TB where т; Si.15.-.,.8 8i For $111 the positive sample is s and the negative samples are 8j where j o i. We use the InfoNCE Loss van den Oord al. (2018),
expla Wk) E exp(g"
Lg log exp
(6)
with query q = 1((s)) as the encoded starting state, and the keys along the K trajectories in the dataset D. By doing so, we are attracting the states that are reachable after steps and pushing those not reachable far out. This approach has been used successfully prior work Liu et al. (2020); Emmons et al. (2020) and we use this method for generalizable solution for finding the state to condition the future on. This particular choice of distance function a means to an end for obtaining the future state condition and further research is required for investigating similar choices for a robust performance across a diverse set of tasks. We leave this investigation for future work.
((st+)) as the encoded future states
A.2 TRAINING
The training for both skill extraction and fine-tuning were done on single NVIDIA 2080Ti GPU. Skill extraction takes approximately 3-4 hours, and fine-tuning requires less than 10 minutes. Our codebase builds upon the SPiRL released code and is located at https: / anonymous, 4open science/r / fist-C5DFF SDtt-C5DF/README README, md. Hyperparameters used for training are listed in Table 3.
A.3 FINE-TUNING
Fine-tuning for both FIST and SPiRL follows the same implementation details. done only on Ddomo which includes 10 trajectories of the agent fulfilling the long horizon task. These trajectories are segmented into uubttrajectories of length H = 10 first (similar to pre-training). and then we update all the network parameters by minimizing the loss equation 3, for 50 epochs of the small dataset. The original training was done on 200 epochs of the large and diverse dataset. Everything else, including batch size, learning rate, and the optimizer remain the same between pre- pre-training and fine-tuning. The hyper-parameter for fine-tuning are listed in Table 4.
"During our initial experimental analysis, we tried a version of FIST that only finetuned the inverse skill dynamics model and that did not show any meaningful results. We hypothesise that without fine-tuning the VAE components, the skill-set of the agent will not include the of distribution parts of the task.
13

---

Under review conference paper at ICLR 2022
Table 3: Training Hyperparameters
Hyperparameter
Value
Contrastive Distance Metric
Encoder output dim Encoder Hidden Layers
32
128
Encoder # Hidden Layers
2
Optimizer
Adam (31 0.9, B2 0.999, LR=1e-3)
Skill extraction
Epochs
200
Batch size
128
Optimizer
Adam(3: = 0.9, A8 = 0.999, LR=1e-3)
(sub-trajectory length)
10
3
5c-4 (Kitchen), 1e-2 (Maze)
Skill Encoder
dim-2 in VAE
128
hidden dim # LSTM Layers
128
Skill Decoder hidden dim
128
# hidden layers
Inverse Skill Dynamic Model hidden dim
128
# hidden layers
Fine-tuning
Epochs
50
Batch size
128
Optimizer
Adam = 0.9, B2 = 0.999, LR=1e-3)
Table 4: Fine-tuning hyperparameters
Hyperparameter
Value
Epochs
50
Epoch cycle train
10
B ENVIRONMENT AND DATASET DETAILS
In this section we explain the tasks and the procedure for data collection for both pre-training the skills and the downstream fine-tuning in each environment separately. The instruction for downloading the dataset as well as its generation is included in code repository.
PointMaze. In this environment, the task is to navigate point mass through a maze, from a start to goal location. The outline of the maze is shown in Figure 3. We train the skills three different datasets, each
14

---

Under review conference paper at ICLR 2022
blocking one side of the maze. To test the method's ability to generalize to unseen long-horizon tasks, we use 10 expert demonstrations that start from random places in maze, but end at goal within the blocked region. This ensures that our demonstrated trajectories are out of distribution compared training data. We evaluate the performance by measuring the episode length and the success rate in reaching the demonstrated goals. The data for PointMaze collected using the same scripts provided the D4RL dataset repository Fu et al. (2020). To generate the pre-training dataset we modified the maze outline to have parts of it blocked. Then we run the same oracle way-point controller on the new maze map and collect continuous trajectories of the agent navigating to random goals, for total of 4 million transitions. For downstream demonstrations we unblock the blocked region, pick a fixed goal within that region, and have the agent navigate from random reset points to the picked goal location. The observation consists of the (x, y) location and velocities.
AntMaze. The task is to control a quadruped ant to run to different parts of the maze. The layout of the maze is the same PointMaze, and the same sides blocked off. Similar to PointMaze we measure the episode length and success rate as our evaluation metric. The data for AntMaze is solely based on the cunt-lareediierse--.
dataset in D4RL. To construct the pre- pre-training data we filter out sub-trajectories that contain states within the blocked region. By doing so, we effectively exclude the region of interest from the pre-training dataset. The result are datasets with 47165, 58329, and 50237 number of transitions, for the Bottom, Right, and Left blocked regions respectively. For the downstream demonstrations we randomly select 10 from the excluded trajectories that start outside the blocked region and end in the blocked region, shown in Figure 3.
Kitchen. The task is to use 7-DoF robotic arm to manipulate different parts of a kitchen environment in a specific order (e.g. open microwave door move the kettle)'. During skill extraction we pre-process the offline data to exclude interactions with certain objects in the environment (c.g. we exclude interactions with the kettle). However, for the demonstrations we pick four sub-tasks one of which includes the objects that were excluded from the skill dataset (e.g. if the kettle was excluded, we pick the task to be to open the microwave, move the kettle, turn the top burner, and slide the cabinet door). In evaluation, for completion of each sub-task in the order consistent with the downstream demonstrations, the agent is awarded with a reward of 1.0 for a total max reward of 4.0 per episode.
The pool of demonstrations are downloaded from the repository of Gupta et al. (2019) located at ht httpp: thub.
total of 24 multi- task long horizon sets of trajectories that the data is collected from. Each trajectory set is sampled at least 10 times via VR tele-operation procedure and the filenames indicate what the agent is trying to achieve (e.g. microwave-kettke kww.kel--sstchh-hiiieeeeeeee on the filenames that do not include the keyword for the task of interest. This will essentially remove all multi-task trajectories that include the sub-task and not just the part that we do not want. It is also worth noting that the excluded object (e.g. the kettle) is still part of the environment and all its state vectors are still visible to the agent, despite the exclusion of the "interaction" with that object. The dataset size for each case is shown in Table 5.
/ //github. comm /gogg] e- researchh re 1 ay-poll Ecy-learme ng. There
For creating the pre- tre-training data, we filter the trajectory sets, solely based
"The environment dynamics are still stochastic and therefore, simple replication of actions would not fulfill the tasks robustly
15

---

Under review conference paper at ICLR 2022
Table 5: Kitchen pre-trainin pre-training dataset sizes
Tasks
# of trajectories
Microwave, Kettle, Top Burner, Light Switch
285
Microwave, Bottom Burner, Light Switch, Slide Cabinet
236
Microwave, Kettle, Slide Cabinet, Hinge Cabinet
230
Microwave, Kettle, Slide Cabinet, Hinge Cahinet
198
C EXPERIMENTAL RESULTS
C.I TABLE OF RESULTS
Table 6: Comparison our approach other baselines on the Maze environments. For each experiment we report the average episode length from 10 fixed starting positions with the standard error across 10 evaluation runs (lower is better). We also report success rate and its standard deviation. The maximum episode length for PointMaze and AntMaze are 2000 and 1000, respectively.
FIST (Ours)
SPiRL
BC+FT
Blocked Region
Environment
Episode Length
Success Rate
Episode Length
Success Rate
Episode Length
Success Rate
Left
PointMaze
363.87
18.73
0.99 + 0.03
1966.7 32.54
0.02 0.04
1089.76 + 173.74
0.74 0.11
Right Bottom
PointMaze
571.21
38.82
0.91 + 0.07
2000 = 0
0.0 0.0
1918.99 + 43.65
0.07 0.06
PointMuze
359.82 +: 3.62
1.0 0.0
1127.47 148.24
0.87 0.10
AntMaze
0.32
Right
AntMaze
913.98 12.01
0.22 + 0.12
1000
0.0 ::
1000
0.0 0.0
Bottom
AntMuze
923.22 + 6.36
0.21 + 0.07
957.85 +: 8.62
0.12 1 0.07
1000 = 0
0.0 0.0
Table 7: Comparison of average episode reward for our approach against other baselines the KitchenRobot environment. The average episode reward (with a max. of 4) along with its standard error measured across evaluation runs (higher is better). Each bolded keyword indicates the task that was excluded during skill data collection.
Task (Unseen)
Environment
FIST (Ours)
SPiRI
BC+FT
Microwave, Kettle, Top Burner, Light Switch Microwave, Bottom Burner, Light Switch, Slide Cabinet
KitchenRabot
3.6 + 0.16
2.1 + 0.48
0.0 + 0.0
KitchenRobul
2.3 1 0.5
2.3 2.3
2.2 : 0.28
Microwave, Kettle, Slide Cabinet, Hinge Cabinet Microwave, Kettle, Slide Cabinet, Hinge Cabinet
KitchenRobot
3.5 3.5 0.3
1.9 0.09
1.3 + 0.47
KitchenRobot
4.0
3.3
1.0 = 0.32
C.2 EXTRA ABLATIONS
Imitation Learning over skills vs. atomic actions. FIST is comprised of two coupled pieces that are both critical for robust performance: the inverse dynamics model over skills and the non-parametric evaluation algorithm. In this experiment measure the influence of inverse skill dynamics model ( (z|ss,, --+H-1).
An alternative baseline learning skill dynamics model is to learn an inverse dynamics model on atomic actions (a st, 8t-H-1) and perform goal-conditioned behavioral cloning (Goal-BC). This model outputs the first action a, required for transitioning from 3, to $ over H steps. We can combine this model with FIST's non-parametric module determine the 8t+ H -1 condition during the evaluation of the policy. As shown in Table 8, temporal abstraction obtained learning an inverse skill dynamics model a critical factor in the performance of FIST.
16

---

Under review conference paper at ICLR 2022
Table 8:
We ablate the use of our inverse skill dynamics model by replacing it with an inverse dynamics model on atomic actions. The baseline ablations only sucoeed on one out of the four tasks. BC learns an inverse dynamics model that takes in state input and outputs distribution over atomic actions. Goal-BC uses both state and the goal subb-task) as input.
Task (Unseen)
FIST (ours)
Goal-BC
Microwave, Kettle, Top Burner, Light Switch Microwave, Bottom Burner, Light Switch, Slide Cabinet Microwave, Kettle, Slide Cabinct, Hinge Cabinet Microwave, Kettle, Slide Cabinet, Hinge Cabinet
3.6 + 0.16
0.0 + 0.0
2.3 - 0.5
1.2 0.3
3.5 - 0.3
1,8 + 0,44
4.0 0.0
0.9 0.1
Is our contrastive distance function an optimal approach for picking the future state to condition on? For environments such as PointMaze, where we have access to the same waypoint controller that generates the demonstrations, we can use the ground truth environment dynamics calculate the oracle state that the waypoint controller would be at, H steps in the future, and use that as the foal. This waypoint controller is not available for the kitchen environment, since the demonstrations were collected using VR tele-operation.
In Table 9. FIST (Oracle) uses the waypoint controller oracle look-up. The results show that with better future conditioning, an oracle approach can solve the pointmaze task even better than FIST with the contrastive distance. Improving the current panimparametric approach for obtaining the future conditioning state could be a very interesting direction for future work.
Table 9:
We ablate FIST against an oracle version pointmaze which uses the ground truth way point planner that has access to the exact state that the agent will end up at H-steps in the future (if it commits to the optimal path).
Section
FIST
FIST (oracle)
Episode Length
Success Rate
Episode Length
Success Rate
Left
363.87 + 18.73
0.99 + 0.03
236.00 + 1.02
1.0 1.0+ 0.00
Right
571.21 = 38.82
0.91 + 0.07
280.93 = 5.61
1.0 0.00
Bottom
359.82 + 3.62
1.0 0.00
269.89 + 3.75
1.0 0.00
One-shot Imitation Learning: The FIST algorithm can be directly evaluated on one-shot in-distribution downstream tasks without any fine-tuning. In this experiment, we want to see if the agent can pick up the right mode within its skill-set with only one demonstration for fulfilling long-horizon task in the kitchen environment. The difference between this experiment and our main result is that the down- task is within the distribution of its pre-trained skill-set. This is still a challenging task since the agent needs to correctly identify the desired mode of skills.
Our hypothesis that in SPiRL, the skill prior is only conditioned on the current state and therefore is, by definition, multi-modal distribution and would require more data to adapt to a specific long-horizon trajectory. For instance, in the kitchen environment, after opening the microwave door, the interaction with any other objects in the environment is possible choice of skills that can be invoked. However, in FIST, by conditioning the skill prior on the future states, we fit a uni-modal distribution over skills. In principle, there should be no need for fine-tuning for invoking those skills within the distribution of the pre- pre-trained skill set.
We compare our approach to SPiRL (Section 4.2) as bascline. In addition, we can provide supervision which skills to invoke to fulfill the long-horizon task by fine-tuning SPiRL (hence SPiRL-F7) for few epochs on the downstream demonstration. As summarized in Table 10, FIST, without any fine-tuning, can fulfill all the long-horizen tasks listed with almost drift from the expert demonstration. We also see that
17

---

Under review conference paper at ICLR 2022
it is tricky to fine-tune SPiRL in one-shot setting, fine-tuning only on one demonstration may cause over-fitting and degradation of performance.
Table 10:
With all subtasks seen in the skill dataset, FIST is able imitate long mongzooizzon task in the kitchen environment. We compare a baseline method, SPiRL, which fails follow the single demo.
Order of tasks (seen in the skill dataset)
FIST (ours)
SPiRL-FT
SPiRL-no-F FT
Kettle, Bottom Burner, Slide Cabinet, Hinge Cabinet Kettle, Top Burner, Light Switch, Slide Cabinet
4.0 = 0.0
0.8 0.19
2.4 + 0.35
3.8 - 0.19
0.16
1.1 # 0.22
Microwave, Kettle, Slide Cabinet, Hinge Cahinet
4.0 + 0.0
1.22
1.0 + 0.37
Top Burner, Bottom Burner, Slide Cabinet, Hinge Cabinet
4.0 0.0
0.1
0.25
18

---

Under review conference paper at ICLR 2022
D BROADER IMPACTS AND LIMITATIONS
Limitations As with all imitation learning methods, the perfor- mance of FIST is related to the quality of the provided demonstra- poor,
Environment
Episode Length
Success Rate
tions. Concretely, when the skill training demonstrations we expect the extracted skills to also sub-optimal, thus, hurting downstream imitation performance. To better understand this lim- itation, we analyze an extremely noisy versions the PointMaze dataset and it for skill extraction. As shown in Table 11, despite achieving a high success rate, the episode length is substantially worse than FIST trained expert data.
Pointaae
621.02 69.87
1.0 0.0
Table 11: We evaluate FIST on the maze env environment with goal at the bottom when the inverse skill model is trained tremely noisy dataset. this case, FIST achieves sub-optimal performance, or unable to imitate the test time demonstra- tion.
Learning structured skills from noisy offline data is an exciting direction for future research.
Broader Impacts The ability to extract skills from offline data and adapt them to solve new challenging tasks in few-shot could impactful domains where large offline datasets are available but control is challenging and cannot be manually scripted. Examples of such domains include autonomous vehicle navigation, warchouse robotics, digital assistants and perhaps in the future, home robots. However, there are also negative potential consequences. First, since in real-world settings offline data will be collected from users at scale there will likely be privacy concerns, especially for video data collected from users' cars homes. Additionally, since FIST extracts skills, without labeled data, quality for large datasets becomes increasingly opaque and if there are harmful skills or behavior present in the dataset FIST may extract those and use them during deployment which could have unintended consequences. A promising direction for future work is to include a human in the loop for skill verification.
19


Reviews:
[{"review_id": "zXXSWqKUfE", "reviewer": "Reviewer_enSE", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths\n-------------\n\n- I like the inverse skill dynamics model idea\n- Nice implementation of skill learning + fine-tuning\n- Thorough analysis, ablations and results\n\n\nWeaknesses\n------------------\n\n- The learned distance function feels like it is \"forced\" because of the inverse skills model and not thought out in detail. Although the experiments/analysis in the supplementary reduced this feeling a lot. This is not a big issue, but if the authors can shed more light to this, it'll make the paper better.\n- The algorithm uses very big offline datasets (on unstructured data) and at least 10 full demonstrations of successfully completing the task. Here I have 2 concerns:\n    - The success rate even after these big datasets is relatively low (<80% on average). Couldn't we do better?\n    - What is practical usage of the proposed method? If I have access to big datasets, I think I can do many other more effective things. For example, insert priors in the form of structures in the policies (see {1,2,3}). Of course I understand that here we do not use any prior at all, but I'd like to see at least a discussion on this.\n- The related work section and motivation is quite narrow. Overall the motivation for developing such a method is rather weak. In other words, I see limited impact. A few examples:\n    - Imitation Learning: the field is vast and has been tackled from many different perspectives. Perhaps the biggest literature for imitation learning comes from the robotics community (see {3,4,5}). The related work section/motivation needs quite some effort to improve.\n   - Skill Extraction with Behavioral Priors and Few-Shot Learning: here again the evolutionary-based perspective is not even mentioned. There are many approaches that compute behavioral repertoires to enable fast (or few-shot) adaptation (see {3,6,7}). Here I'd also like to see some discussion.\n- The supplementary material needs to better describe the datasets (samples of offline data, a few visual examples, ...)\n- The supplementary video does not show ANY video of the learned behaviors!!\n\n\nReferences\n----------------\n\n{1}: Bahl, S., Mukadam, M., Gupta, A. and Pathak, D., 2020. Neural dynamic policies for end-to-end sensorimotor learning. NeurIPS.\n\n{2}: Stulp, F. and Sigaud, O., 2013. Robot skill learning: From reinforcement learning to evolution strategies. Paladyn, Journal of Behavioral Robotics, 4(1), pp.49-61.\n\n{3}: Chatzilygeroudis, K., Vassiliades, V., Stulp, F., Calinon, S. and Mouret, J.B., 2019. A survey on policy search algorithms for learning robot controllers in a handful of trials. IEEE Transactions on Robotics, 36(2), pp.328-347.\n\n{4}: Billard, A., Calinon, S., Dillmann, R. and Schaal, S., 2008. Survey: Robot programming by demonstration (No. BOOK_CHAP, pp. 1371-1394). Springrer.\n\n{5}: Osa, T., Pajarinen, J., Neumann, G., Bagnell, J.A., Abbeel, P. and Peters, J., 2018. An algorithmic perspective on imitation learning. Foundations and Trends® in Robotics.\n\n{6}: Cully, A., Clune, J., Tarapore, D. and Mouret, J.B., 2015. Robots that can adapt like animals. Nature, 521(7553), pp.503-507.\n\n{7}: Pautrat, R., Chatzilygeroudis, K. and Mouret, J.B., 2018, May. Bayesian optimization with automatic prior selection for data-efficient direct policy search. In 2018 IEEE International Conference on Robotics and Automation (ICRA) (pp. 7571-7578). IEEE.\n", "summary_of_the_paper": "The paper presents an approach for skill learning and fine-tuning. The main idea of the manuscript is to combine offline skill learning with few online demonstrations to provide efficient learning experience in out of distribution (OOD) tasks. There are two main novelties in the paper:\n\n- Learning of an inverse skill dynamics model that infers which skills should be used given the current state and a future state\n- Learning a distance function between states that can be used to select states from the demos as \"target/future\" states for better generalization\n\nApart from the technical novelties, the paper presents ablation studies with intuitions on which of the parts of the new approach (FIST) are important and why.", "summary_of_the_review": "The paper was an interesting and nice read. Skill learning and fast adaptation are key features that can enable widespread adoption of robots. I have however three main concerns: (a) the algorithm uses very big offline datasets + 10 full demonstrations (states and actions) and success rate of completing the tasks is relatively low (<80% on average); we can certainly do better than this with minor prior information, (b) the related work/motivation part needs more work (it is not convincing me at the moment) and (c) the supplementary video does not show the performance of the algorithm.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "ZLPjzcZQeUj", "reviewer": "Reviewer_JEuK", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "*Strengths*\n\nThe method proposes some simple extensions to the method in SPiRL that seem to be effective. It makes sense that having the skill grounding take not just a current state but a goal state should make mapping to the correct skill easier. Adding in finetuning on the few demos is also sensible idea. Lastly, using a learned temporal distance to query the closest state in the provided demo makes sense, and the results indicate it performs better than using Euclidean distance.\n\nThe experimental results show that the method seems to perform SPiRL significantly across a range of environments.\n\n*Weaknesses*\n\nI have some questions about the justification for the performance gap between this method and SPiRL. First, to confirm - SPiRL is also finetuned on the few shot demos correct? So when the paper states that SPiRL \"The skills learned by SPiRL fail to generalize when the point mass falls outside of the training distribution, causing it to get stuck in corners\", why is this an issue for SPiRL but not FIST when both are finetuned?\n\nIn general I think it makes sense that skill selection should be conditioned on not just the current state but a goal state H steps ahead. However from what I can tell this is the primary difference between FIST and SPiRL, is it the case that just this change is the difference between the 0% -> 80% success between SPiRL and FIST? It would be great if the authors could expand more on why exactly SPiRL fails, and why the change to having the skill prior take s_{t+H} impacts performance so much. It seems like a minor change but the results indicate that its important. \n\nThere are a few typos and grammatical errors throughout the paper, and in general the presentation/figures could be improved. Fig 4 should have some measure of uncertainty/standard error. I also can't find how many trials/seeds are used in the experiments.", "summary_of_the_paper": "The paper presents an approach for skill extraction from an offline dataset of demos/interaction data, for use in few shot imitation learning of previously unseen long horizon tasks. The method trains a skill encoder mapping trajectories to a latent skill space, as well as a decoder which outputs actions conditioned on a state and skill (similar to SPiRL, Pertsch et al). The method also learns a skill prior which maps a current state and goal state (H steps ahead of current state) to the latent skill space. The main different to SPiRLis that the skill prior takes as input both a current and goal state, while in SPiRL it took in only takes in the current state. The skill prior and skill encoder/decoder are jointly trained on the demos. \n\nThen during evaluation of a new long horizon task given a few demos, the agent takes in as input the current state, queries a nearby state in the provided demos, recover the state H steps ahead, and apply the skill prior given the s, and s_{t+H}. The querying nearby states is using a learned distance function trained to capture temporal distance. All components of the method are finetuned on the few demos provided for the new task.\n\nExperiments suggest this method outperforms SPiRL and naive BC + finetuning on unseen multi step tasks in a simulated Franka kitchen environment and on ant maze.", "summary_of_the_review": "Overall the paper presents some technical extensions over SPiRL, and seems to have significant performance improvements. I have some questions about the experiments, but if addressed I think the papers results are solid.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "RZEaiRVTOCu", "reviewer": "Reviewer_LjxN", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "### Strengths\n\n- This paper tackles an important problem of solving a new task without a huge number of interactions utilizing skills and a few downstream demonstrations.\n- The proposed method is intuitive and shows potential for a versatile agent which can quickly adapt to a new task. \n- The ablation study dissects components of the proposed method and shows the importance of each of the components, supporting the paper's claims.\n\n### Weaknesses\n\n- In the kitchen environment, all tasks start from Microwave. In some of the tasks, the baseline approaches (Goal-BC, FIST-no-pretrain, SPiRL (closest), SPiRL (H-steps)) fail to accomplish the Microwave sub-task for all random seeds while in some other cases, they succeed. How can we interpret these results?\n\n- In Algorithm 1, a skill $z$ is sampled every environment step. Is this better than executing $H$ actions from the sampled skill $z$? To adjust to the given few-shot demonstrations, inferring a lookahead state and skill every time step makes sense. But, in terms of temporally consistent action execution, executing a skill as a whole may work better. \n\n\n### Minor comments\n\n- In Section 3.2 (iii), $z \\sim q_\\psi(z|s_t,s^*_{t})$ should be $z \\sim q_\\psi(z|s_t,s^*_{t+H})$.\n", "summary_of_the_paper": "This paper proposes a novel semi-parametric way of few-shot imitation learning. It assumes access to a large task-agnostic dataset to extract fixed-length skills $\\pi(a|s,z)$ and an inverse skill dynamics model $p(z|s_t, s_{t+H})$, which infers required skills given the current and future states. With the inverse skill dynamics model, an agent can sample a skill to reach a given goal state. When tackling a new downstream task with a few demonstrations, this method seeks for the closest state $s^*_{t}$ in these demonstrations and sets $s^*_{t+H}$ as a new goal. Since we assume $s_t$ and $s^*_{t}$ are similar, reaching $s^*_{t+H}$ should be plausible with the learned skills. The experiments show that this approach successfully solves unseen tasks with a few demonstrations and outperforms the behavioral cloning and skill prior RL baselines. ", "summary_of_the_review": "Overall, this paper proposes an intuitive and novel approach for skill-based few-shot imitation learning. The experimental results show the advantage of learning the inverse skill dynamics and the proposed semi-parametric imitation learning approach.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "JDCcjAXfNEs", "reviewer": "Reviewer_s3dr", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Overall, this paper is quite interesting to read and the presented problem is worth being investigated. I will write its strength and weakness respectively.\n\n[Major concern:]\n- As I summarized above, the first contribution on long-horizon task learning by extracting behaviour skills and by inverse skill model to selecting skills seems trivial since the basic components were proposed in previous work (Pertsch et al. 2020, Ajay et al. 2021). \n- The second contribution I think is more interesting but still lacks in-depth investigation. \n  - For example, the skill level is defined on state-action sequences, but the learned distance function for finding similar states is based on the state level. Will there be any issues? Any distributional perspective to explain why your distance function can make the learned skills generalize or adapts?\n\n[Minor comments:]\n- It seems to me this paper involves lots of work there, however, the main draft is underdeveloped that it even lacks a conclusion section and a Reference section name. The format needs to be double-checked.\n- Could the authors explain the details in the process \"LookAhead\" in Algorithm 1?", "summary_of_the_paper": "This paper tackles the problem of learning generalizable long-horizon tasks from offline human demonstrations with the below two insights:\n\n- (1) The long-horizon task learning is enabled by extracting behaviour priors Z as basic skills from offline data, which has been seen in previous publications  (Pertsch et al. 2020).\n- (2) The generalization is enabled by matching the new states from adapted envs with the states from demonstrations using a pre-trained distance function, which is common in few-shot imitation learning.\n\nExperiments on a simulated navigation task (Pointmass and Ant) and a robotic manipulation task (Kitchen) were performed to prove the effectiveness of the proposed approach. \n\n======================================================\n\nClarifications:\n\nSince task generalization benchmark is not well established in our research community currently, I feel it is worthy to state clearly what is the claimed generalization here. Specifically, the demonstrated generalization performance is under the few-shot learning setting. For example, \n\n- (1) In Pointmass and Ant env, some regions are blocked in offline human demonstrations. \n    - When testing generalization, the authors provide a few more human demonstrations covering the blocked regions. The proposed method will thus generalize.\n- (2) In the Kitchen task, some objects are removed from offline human demonstrations.\n    - When testing generalization, the authors provide a few more human demonstrations covering the removed objects.\n", "summary_of_the_review": "Overall, it is a good paper to read and the authors have done a lot of works. However, the major contribution is not sufficiently novel to be clearly accepted in ICLR. As a result, I recommend a weak acceptance. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Meta-review: The reviewers agree that addressing long-horizon tasks with off-line learning and fine tuning afterwards from demonstrations is an interesting and relevant topic. The technical ideas about learning a relevance metric to select relevant off-line data, and to learn an inverse skill dynamics models. The experimental results are convincing, even if success rates are sometimes lower than expected. All reviewers recommend acceptance of the paper.
Decision: Accept (Poster)