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: 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.   Then 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.  Experiments suggest this method outperforms SPiRL and naive BC + finetuning on unseen multi step tasks in a simulated Franka kitchen environment and on ant maze.