Hierarchical Few-Shot Imitation with Skill Transition Models

1 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-world tasks – long-horizon planning and generalizability . Practical tasks are often long-horizon in the sense that they require a robot to complete a sequence of subtasks . For example , to cook a meal a robot might need to prepare ingredients , place them in a 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 a similar skill set . For example , if a 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 to solving long-horizon tasks . Given a large unlabeled offline dataset of robotic demonstrations solving a diverse set of tasks this family of approaches Singh et al . ( 2020 ) ; Pertsch et 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 a structured manner and can solve long-horizon navigation and manipulation tasks . However , the generalization capabilities of RL with behavioral priors are limited since a different RL agent needs to be trained for each downstream task and training each RL agent often requires millions of environment interactions . On the other hand , few-shot imitation learning has been a promising paradigm for generalization . In the fewshot imitation learning setting , an imitation learning policy is trained on an offline dataset of demonstrations and is then adapted in few-shot to a downstream task Duan et al . ( 2017 ) . Few-shot imitation learning has the added advantage over RL in that it is often easier for a human to provide a handful of demonstrations than it is to engineer a new reward function for a 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 . ( 2011b ) , and especially so in the few-shot setting when only a 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 a given state can correspond to multiple skills . Given a sufficiently large dataset of 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 a 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 ) , a new algorithm for few-shot imitation learning with skills that enables generalization to 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 on a future state , which helps FIST learn uni-modal future conditioned skill distribution that can then be utilized in few-shot . The inverse skill model is then used as a policy to select skills that will take the agent to the desired future state . Second , we train a distance function to 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 robotic manipulation tasks in a 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 a distance function that is used for semi-parametric 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 to final performance by ablating different components of our method such as future conditioning and fine-tuning on downstream data . 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 . ( 2011a ) and Inverse Reinforcement Learning ( IRL ) Ng & Russell ( 2000 ) . BC approaches learn policies πθ ( a|s ) that most closely match the stateconditioned action distribution of the demonstration data . IRL approaches learn a 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 are prone to drifting away from the demonstration distribution during evaluation due to error accumulation Ross et al . ( 2011b ) . For this reason , BC policies work best when the time-horizon of the task is short . Skill Extraction with Behavioral Priors : Hard-coding prior knowledge into a policy or dynamics model has been considered as a 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 new 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 recognition 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 et 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 a 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 a 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 a behavioral prior to an unlabeled offline dataset of demonstrations to extract skills and then fits an imitation learning policy over the previously acquired skills to generalize in few-shot to new tasks . 3 APPROACH . 3.1 PROBLEM FORMULATION . Few-shot Imitation Learning : We denote a demonstration as a sequence of states and actions : τ = { s1 , a1 , s2 , a2 , . . . , sT , aT } . In a few-shot setting we assume access to a small dataset of M such expert demonstrations Ddemo = { τi } i=Mi=1 that fulfill a specific long horizon task in the environment . For instance a 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 to imitate this behavior using only a few available example trajectories . 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 { τi = { ( st , at ) } t=Tit=1 } i=Ni=1 . 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 human-experts through teleoperation Zhang et al . ( 2018 ) . Skill extraction refers to an unsupervised learning approach that utilizes this reward-free and task-agnostic dataset to learn a skill policy in form of πθ ( a|s , z ) where a is action , s is the current state , and z is the skill . Our hypothesis is that by combining these skill primitives we can solve semantically similar long-horizon tasks that have not directly been seen during the training . In this work we propose a new architecture for skill extraction based on continuous latent variable models that enables a semi-parametric evaluation procedure for few-shot imitation learning .

Summary: 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.