Transferring Hierarchical Structure with Dual Meta Imitation Learning

Hierarchical Imitation learning ( HIL ) is an effective way for robots to learn subskills from long-horizon unsegmented demonstrations . However , the learned hierarchical structure lacks the mechanism to transfer across multi-tasks or to new tasks , which makes them have to learn from scratch when facing a new situation . Transferring and reorganizing modular sub-skills require fast adaptation ability of the whole hierarchical structure . In this work , we propose Dual Meta Imitation Learning ( DMIL ) , a hierarchical meta imitation learning method where the high-level network and sub-skills are iteratively meta-learned with model-agnostic meta-learning . DMIL uses the likelihood of state-action pairs from each sub-skill as the supervision for the high-level network adaptation , and use the adapted highlevel network to determine different data set for each sub-skill adaptation . We theoretically prove the convergence of the iterative training process of DMIL and establish the connection between DMIL and the Expectation-Maximization algorithm . Empirically , we achieve state-of-the-art few-shot imitation learning performance on the meta-world ( Yu et al. , 2019b ) benchmark and comparable results on the Kitchen environment . 1 INTRODUCTION . Imitation learning ( IL ) has shown promising results for intelligent robots to conveniently acquire skills from expert demonstrations ( Zhu et al. , 2018 ; Peng et al. , 2018 ) . Nevertheless , acquiring skills from long-horizon unsegmented demonstrations has been a challenge for IL algorithms , because of the well-known issue of compounding errors ( Ross et al. , 2011 ) . This is one of the crucial problems for the application of IL methods to robots since plenty of manipulation tasks are long-horizon . Hierarchical Imitation Learning ( HIL ) aims to tackle this problem by decomposing long-horizon tasks with a hierarchical model , in which a set of sub-skills are employed to accomplish specific parts of the long-horizon task , and a high-level network is responsible for determining the switching of sub-skills along with the task . Such a hierarchical structure is usually modeled with Options ( Daniel et al. , 2016 ; Krishnan et al. , 2017 ; Jing et al. , 2021 ) or goal-conditioned structure ( Le et al. , 2018 ) . HIL expresses the nature of how humans solve complex tasks , and has been considered to be a valuable direction for IL algorithms ( Osa et al. , 2018 ) . However , most current HIL methods have no explicit mechanism to transfer previously learned skills to new tasks with few-shot demonstrations . This requirement comes from that the learned hierarchical structure may conflict with discrepant situations in new tasks . As shown in figure 1 ( a ) , both the high-level network and sub-skills need to be transferred to new forms to satisfy new requirements : the high-level network needs new manners to schedule sub-skills in new tasks ( for example , calling different sub-skills at the same state ) , and each sub-skill needs to adapt to new specific forms in new tasks ( for example , grasping different kinds of objects ) . Quickly adapting to new tasks can significantly increase the generalization ability of HIL methods and make them be applied to a wider range of scenarios , thus an appropriate approach needs to be exploited to endow HIL with such kind of ability with few-shot new task demonstrations . Recently , meta imitation learning ( MIL ) ( Yu et al. , 2018b ; Finn et al. , 2017b ; Yu et al. , 2018a ) employ model-agnostic meta-learning ( MAML ) ( Finn et al. , 2017a ) to imitation learning procedure to enable the learned policy to quickly adapt to new tasks with few shot demonstrations . MAML first fine-tunes the policy network in the inner loop , then evaluates the fine-tuned network to update its initial parameters with end-to-end gradient descent at the outer loop . This inspires us to integrate MAML into HIL to transfer the hierarchical structure . However , most MIL methods meta-learn a monolithic policy in an end-to-end fashion , which does not conform to the iterative training of the bi-level structure in HIL , where the training supervision of each level comes from the other level . Integrating MAML into such a self-supervised bi-level optimization procedure is a novel problem . The main challenge comes from how to schedule the inner loops and outer loops of the meta-learning process of bi-level modules in HIL to ensure convergence . In this work , we propose a novel hierarchical meta imitation learning framework called dual meta imitation learning ( DMIL ) to successfully incorporate MAML into the iterative training process of HIL , as shown in figure 1 ( b ) . We make the high-level network and sub-skills as mutual supervision during the bi-level MAML process : the likelihood of each state-action pair in few-shot demonstrations of each sub-skill can provide the supervisions for the meta-training of high-level network , and the fine-tuned high-level network in turn determines the data sets for the meta-training of each subskill , like a general EM process . We further design an elaborate training procedure for DMIL that first fine-tunes the high-level network and sub-skills in sequence at inner loops , then meta-updates them simultaneously at outer loops . We theoretically prove the convergence of this special training procedure by leveraging previous results from Ravi & Beatson ( 2019 ) ; Finn et al . ( 2018 ) ; Zou & Lu ( 2020 ) to reframe both MAML and DMIL as hierarchical Bayes inference processes and get the convergence of DMIL according to the convergence of MAML Fallah et al . ( 2020 ) . We test our method on the challenging meta-world benchmark environments ( Yu et al. , 2019b ) and the Kitchen environment of D4RL benchmarks ( Fu et al. , 2020 ) . In our experiments , we successfully acquire a set of meaningful sub-skills from a large scale of manipulation tasks , and achieve state-ofthe-art few-shot imitation learning abilities in the ML45 suite . In summary , the main contributions of this paper are as follows : • We propose DMIL , a novel hierarchical meta imitation learning framework that meta-learns both the high-level network and sub-skills from unsegmented multi-task demonstrations . • We propose a special EM-like training algorithm to schedule the meta-learning and hierarchical imitation learning processes in DMIL and theoretically guarantee its convergence . • We achieve state-of-the-art few-shot imitation learning performance on meta-world benchmark environments and comparable results in the Kitchen environment . 2 RELATED WORK . 2.1 HIERARCHICAL IMITATION LEARNING . Recovering inherent sub-skills contained in expert demonstrations and then reuse them with hierarchical structures has long been an interesting topic in the hierarchical imitation learning ( HIL ) domain . According to whether there are pretraining tasks , we can divide HIL methods into two categories . The first kind aims to manually design a set of simple pretraining tasks that could encourage distinct skills or primitives , and then they learn a high-level network to master the switching of primitives to accomplish complex tasks ( Florensa et al. , 2017 ; Peng et al. , 2019 ; Liu & Hodgins , 2017 ; Merel et al. , 2019 ; Le et al. , 2018 ) . However , for unsegmented demonstrations where no pretraining tasks are provided , which is the situation in our paper , these methods can not be applied . The second kind of methods aim to learn sub-skills with unsupervised learning methods . Daniel et al . ( 2016 ) ; Krishnan et al . ( 2017 ) acquire Options ( Sutton et al. , 1999 ) from demonstrations with an Expectation-Maximization-like procedure and use the Baum-Welch algorithm to estimate the parameters of different options . Henderson et al . ( 2018 ) ; Jing et al . ( 2021 ) integrate generative adversarial networks into option discovery process . Li et al . ( 2017 ) ; Sharma et al . ( 2019 ) ; Lee & Seo ( 2020 ) incorporate generative-adversarial imitation learning ( Ho & Ermon , 2016 ) framework and an information-theoretic metric ( Chen et al. , 2016 ) to simultaneously imitate the expert and maximize the mutual-information between latent sub-skill categories and corresponding trajectories to acquire decoupled sub-skills . There are some methods called mixture-of-expert ( MoE ) that compute the weighted sum of all primitives to get the action rather than only using one of them at each time step ( Hausknecht & Stone , 2016 ; Neumann et al. , 2009 ; Jacobs et al. , 1991 ) . Other methods aim to seek an appropriate latent space that can map sub-skills into it and then condition a policy on the latent variable to reuse sub-skills ( Lynch et al. , 2019 ; Haarnoja et al. , 2018 ; Hausman et al. , 2018 ; Hakhamaneshi et al. , 2021 ; Price & Boutilier , 2003 ) . For transferring learned sub-skills , some work fine-tune the whole structure in new tasks ( Hakhamaneshi et al. , 2021 ) . However , the performance of fine-tuning all depends on the generalization of deep networks , which may vary among different tasks and network designs . 2.2 META IMITATION LEARNING . Meta imitation learning , or one-shot imitation learning , leverages various meta-learning methods and multi-task demonstrations to meta-learn a policy that can be quickly adapted to a new task with few-shot new task demonstrations . Duan et al . ( 2017 ) ; Cachet et al . ( 2020 ) employ self-attention modules to process the whole demonstration and the current observation to predict current action . Yu et al . ( 2018a ) ; Finn et al . ( 2017b ) ; Yu et al . ( 2018b ) use model-agnostic meta-learning ( MAML ) ( Finn et al. , 2017a ) to achieve one-shot imitation learning ability for various manipulation tasks with robot or human visual demonstrations . Xu et al . ( 2019 ) ; Yu et al . ( 2019a ) propose to meta-learn a robust reward function that can be quickly adapted to new tasks and then use it to perform IRL in new tasks . However , they need downstream inverse reinforcement learning after the adaptation of reward functions , thus conflicts with our goal of few-shot adaptation . Most of above methods only learn one monolithic policy , lacking the ability to model multiple sub-skills in long-horizon tasks . Some works aim to tackle the multi-modal data problem in meta-learning by avoiding single parameters initialization across all tasks ( Vuorio et al. , 2019 ; Alet et al. , 2018 ; Frans et al. , 2018 ; Yao et al. , 2019 ) , but they lack the mechanism to schedule the switching of different sub-skills over time . There are some works that also meta-learn a set of sub-skills in a hierarchical structure ( Yu et al. , 2018a ; Frans et al. , 2018 ) , but they either use manually designed pretraining tasks or relearn the high-level network in new tasks , which is not appropriate in few-shot imitation learning settings . 3 METHOD . 3.1 PRELIMINARIES . We denote a discrete-time finite-horizon Markov decision process ( MDP ) as a tuple ( S , A , T , P , r , ρ0 ) , where S is the state space , A is the action space , T is the time horizon , P : S × A × S → [ 0 , 1 ] is the transition probability distribution , r : S × A → R is the reward function , and ρ0 is the distribution of the initial state s0 . 3.2 FORMULATION OF META IMITATION LEARNING PROBLEM . We firstly introduce the general setting of the meta imitation learning problem . The goal of meta imitation learning is to extract some common knowledge from a set of robot manipulation tasks { Ti } that come from the same task distribution p ( T ) , and adapt it to new tasks quickly with few shot new task demonstrations . As in model-agnostic meta-learning algorithm ( MAML ) ( Finn et al. , 2017a ) , we formalize the common knowledge as the initial parameter θ of the policy network πθ that can be efficiently adapted with new task gradients . For each task Ti ∼ p ( T ) , a set of demonstrations Di is provided , where Di consists of N demonstration trajectories : Di = { τij } Nj=1 , and τij consists of a sequence of state-action pairs : τij = { ( st , at ) } Tij t=1 , where Tij is the length of τij . Each Di is randomly split into support set Dtri and query set Dvali for meta-training and meta-testing respectively . During the training phase , we sample m tasks from p ( T ) , and in each task Ti , we use Dtri to fine-tune πθ to get the adapted taskspecific parameters λi with gradient descent , and then evaluate it with Dvali to get the meta-gradient of Ti , and we optimize the initial parameters θ with the average of meta-gradients from all m tasks . As in Finn et al . ( 2017b ) ; Duan et al . ( 2017 ) , we use behavior cloning ( Atkeson & Schaal , 1997 ) loss as our metrics for meta-training and meta-testing . It aims to train a policy πθ : S → A that maximizes the likelihood such that θ∗ = arg maxθ ∑N i=1 log πθ ( ai|si ) , where N is the number of provided state-action pairs . We denote the loss function of this optimization problem as LBC ( θ , D ) , and the general objective of meta imitation learning problem is : min θ m∑ i=1 LBC ( λi , Dvali ) , ( 1 ) where λi = θ − α∇θLBC ( θ , Dtri ) , and α is a hyper-parameter which represents the inner-update learning rate .

Summary: This paper proposes a meta imitation learning framework aimed at learning long-horizon robot control tasks with fast adaptation capabilities. Specifically, the proposed approach adapts the model-agnostic meta learning framework for learning a hierarchical policy, where both levels of the hierarchy are meta-trained and fine-tuned at test time to learn new tasks. On the metaworld benchmark, the proposed method outperforms prior work, and qualitative analysis shows that the method discovers meaningful skills.