Transferring Hierarchical Structure with Dual Meta Imitation Learning
Abstract: Hierarchical Imitation Learning (HIL) is an effective way for robots to acquire sub-skills from long-horizon unsegmented demonstrations.
However, the learned hierarchical structure lacks the mechanism to transfer across multi-tasks or adapt 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 abilities of both the high-level network and sub-skills to reschedule new forms of sub-skills in new tasks.
In this work, we propose Dual Meta Imitation
Learning (DMIL), a hierarchical meta imitation learning method where the highlevel network and sub-skills are iteratively meta-learned with model-agnostic metalearning (MAML [1]).
DMIL uses the likelihood of state-action pairs from each subskill as the supervision for the high-level network adaptation, and uses the adapted high-level 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 Expectation-Maximization algorithm.
Empirically, we achieve state-of-the-art few-shot imitation learning performance on the Meta-world [2] benchmark and competitive results on long-horizon tasks of
Kitchen environments.
Introduction
Imitation learning (IL) has shown promising results for intelligent robots to conveniently acquire skills from expert demonstrations [3, 4].
Nevertheless, imitating long-horizon unsegmented demonstrations has been a challenge for IL algorithms, because of the well-known issue of compounding errors [5].
This is one of the crucial problems for the application of IL methods to robots since plenty of practical 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 speciﬁc 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 trained with Expectation-Maximization algorithm [6, 7, 8] or goal-conditioned IL paradigms
HIL expresses the nature of how humans solve complex tasks, and has been considered to be a valuable direction for IL algorithms [10].
However, most current HIL methods have no explicit mechanism to transfer previously learned sub-skills to new tasks with few-shot demonstrations.
This requirement comes from that the learned hierarchical structure may conﬂict with discrepant situations in new tasks.
As shown in ﬁg. 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 pulling or pushing at the same state), and each sub-skill needs to adapt to new speciﬁc forms in new tasks (for example, grasping different kinds of objects).
This drives us to develop new methods to endow HIL with the ability to simultaneously transfer both the high-level network and sub-skills with few-shot new task demonstrations.
[Figure]
[Figure]
Recently, meta imitation learning (MIL) [11, 12, 13] employs model-agnostic meta-learning (MAML) [1] into the imitation learning procedure to enable the learned policy to quickly adapt to new tasks with few shot demonstrations.
MAML ﬁrst ﬁne-tunes the policy network in the inner loop, then evaluates the ﬁne-tuned network to update its initial parameters with end-to-end gradient descent at the outer loop.
The success of MIL inspires us to integrate MAML into HIL to transfer the hierarchical structure in new tasks.
However, this is not straightforward.
HIL is a bi-level structure that is trained in an iterative and self-supervised paradigm [6], thus both the high-level network and sub-skills need to be meta-learned by MAML.
Intuitively, the question is: should sub-skills choose the ﬁne-tuned high-level network or the original high-level network for their inner loops and outer loops?
The same question applies to the high-level network.
This is a dual meta learning problem, and we need to explore appropriate methods for MAML to schedule the ﬁne-tune steps (inner loops) and meta-update steps (outer loops) of the bi-level network in HIL to ensure convergence.
In this work, we propose a novel hierarchical meta imitation learning framework called Dual Meta
Imitation Learning (DMIL) to incorporate MAML into the iterative training process of HIL, as shown in ﬁg. 1(b).
Based on the common EM-like HIL structure [6, 7] we design an elaborate bi-level MAML procedure for this hierarchical structure to make it can be fully meta-learned.
In this procedure, we ﬁrst ﬁne-tune the high-level network and sub-skills in sequence at inner loops, then meta-update them simultaneously at outer loops.
We theoretically prove the convergence of this special training procedure by leveraging previous results from [14, 15, 16] to reframe both MAML and DMIL as hierarchical Bayes inference processes and get the convergence of DMIL according to the convergence of MAML from previous results [17].
We test our method on the challenging meta-world benchmark environments [2] and the Kitchen environment of D4RL benchmarks [18].
In our experiments, we successfully acquire a set of meaningful sub-skills from a large scale of manipulation tasks, and achieve state-of-the-art few-shot imitation learning abilities in the ML suite.
In summary, the main contributions of this paper are:
• 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 in a general EM-like fashion, and theoretically prove its convergence.
• We achieve state-of-the-art few-shot imitation learning performance on meta-world benchmark environments and competitive results in the Kitchen environment.
Related Work 2.1 Hierarchical Imitation Learning
Recovering inherent sub-skills contained in expert demonstrations and reusing 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 ﬁrst one employs a set of manually designed pretraining tasks that encourage distinct skills or primitives, then learn a high-level network to master the switching of primitives to accomplish complex tasks [19, 20, 21, 22, 9].
However, for unsegmented demonstrations where no pretraining tasks are provided, which is more common in the real world, these methods can not be applied.
The second kind of methods aim to learn sub-skills with unsupervised learning methods.
Daniel et al. [6], Krishnan et al.
[7] acquire Options [23] from demonstrations with an Expectation-Maximizationlike procedure and use the Baum-Welch algorithm to estimate the parameters of different options.
Henderson et al.
[24], Jing et al.
[8] integrate generative adversarial networks into the option discovery process.
Li et al.
[25], Sharma et al.
[26], Lee and Seo [27] incorporate generative-adversarial imitation learning [28] framework and an information-theoretic metric [29] 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 also 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 [30, 31, 32].
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 generated from the latent space to reuse sub-skills [33, 34, 35, 36, 37].
In this work, we use the EM-like HIL methods used in [6, 7] as they can acquire semantic and separated sub-skill networks which have shown better results in recent works [38], and compare other kinds of HIL methods as baselines to show the superiority of our method in experiments.
None of above HIL methods do not take the fast adaptation ability into consideration: they assume testing tasks are in the same distribution of training tasks.
Although some work ﬁne-tune the whole structure in new tasks [36], the performance of ﬁne-tuning all depends on the generalization of deep networks, which may vary among different tasks and network designs.
Thus we aim to enhance HIL methods with meta learning abilities, as described below.
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.
[39], Cachet et al.
[40] employ self-attention modules to process the whole demonstration and the current observation to predict the current action.
Yu et al. [13], Finn et al.
[12], Yu et al.
[11] use model-agnostic meta-learning (MAML)
[1] to achieve one-shot imitation learning ability for various manipulation tasks with robot or human visual demonstrations.
Xu et al.
[41], Yu et al.
[42] 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 conﬂicts with our goal of few-shot adaptation.
Most 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 [43, 44, 45, 46], 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 [13, 45], 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.
Method 3.1 Meta Imitation Learning
We denote a discrete-time ﬁnite-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 s 0 .
The goal of meta imitation learning is to extract some common knowledge from a set of robot manipulation tasks {T i } 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) [1], we formalize the common knowledge as the initial parameter θ of the policy network π θ that can be efﬁciently adapted with new task gradients.
For each task T i ∼ p ( T ) , a set of demonstrations
D i is provided, where D i consists of N demonstration trajectories: D i = { τ ij } Nj =1 , and τ ij consists of a sequence of state-action pairs: τ ij = { ( s t , a t ) } T ij t =1 , where T ij is the length of τ ij .
Each D
i is randomly split into support set D tri and query set D vali for meta-training and meta-testing respectively.
During the training phase, we sample m tasks from p ( T ) , and in each task T i , we use D tri to ﬁne-tune π θ to get the adapted task-speciﬁc parameters
λ i with gradient descent, and then evaluate it with D vali to get the meta-gradient of T i , and we optimize the initial parameters θ with the average of meta-gradients from all m tasks. The policy [Figure] ො𝑎
[Figure]
ො𝑎 [Figure] 𝑎 [Figure]
Multiply
[Figure] selecting ℒ [Figure] π θ : S → A is trained to maximize the likelihood such that θ
[Equation] , where N is the number of provided state-action pairs.
We denote the loss function of this optimization problem as L BC ( θ, D ) , and the general objective of meta imitation learning problem is:
[Equation]
(1) where λ i = θ − α ∇ θ L BC ( θ, D tri ) , and α is a hyper-parameter which represents the inner-update learning rate.
3.2 Dual Meta Imitation Learning (DMIL)
In this work we assume at each time step t , the robot may switch to different sub-skills to accomplish the task.
We deﬁne the sub-skill category at each time step t as z t = 1 , · · · , K , where K is the maximum number of sub-skills.
We assume a successful trajectory τ ij of a task T i is generated from several (at least one) sub-skill policies, i.e., τ ij = (cid:80) T ij t =1 { ( s t , π E ( s t | z t )) } , where π E represents the expert policy.
Our goal is to learn a hierarchical structure from multi-task demonstrations {D 1 , · · · , D m } in an unsupervised fashion.
The high-level network π θ h that parameterized by θ h determines the sub-skill category ˆ z t at each time step t , and the z -th sub-skill among K different sub-skills π θ l 1 , · · · , π θ lK will be called to predict the corresponding action ˆ a t of state s t , where the hat symbol denotes the predicted result.
We use λ h and λ l 1 , · · · , λ lk to represent the adapted parameters of θ h and θ l 1 , · · · , θ lK respectively.
We condition the high-level network only on states, i.e., ˆ z t =
π θ h ( s t ) , to ﬁt the actual situation at the testing phase.
DMIL aims to ﬁrst ﬁne-tune both π θ h and π θ l 0 , · · · , π θ lK and then meta-update them.
In a new task, π θ h may not provide correct sub-skill categories as stated in the introduction.
However, sub-skills still retain the ability to give out supervision for the high-level network with knowledge learned from previously learned tasks and few-shot demonstrations.
This is because most robot manipulation tasks are made up of a set of shared basis skills like reach, push and pick-place .
As shown on the left side of ﬁg. 2, the sub-skill that gives out the closet ˆ a t to a t can be seen as supervision for the high-level network to classify s t into this sub-skill.
On the other hand, the adapted high-level network can classify each data point in provided demonstrations to different sub-skills for them to perform ﬁnetuning, as shown on the right side of ﬁg. 2.
In summary, DMIL contains four steps for one training iteration.
We call them High-Inner-Update (HI) , Low-Inner-Update (LI) , High-Outer-Update (HO) , and Low-Outer-Update (LO) , which represents the ﬁne-tuning and meta-updating process of the bi-level networks respectively.
The key problem is how to arrange these optimization orders to ensure convergence.
We ﬁrst introduce these steps formally here, then discuss why they can achieve convergence in the next section.
The whole procedure is summarized in algorithm 1 in appendix.
HI:
For each sampled task T i , we sample the ﬁrst batch of trajectories { τ i 1 } from D tri .
The principle of this step is to use sub-skill that can predict the closest action to the expert action to provide self-supervised category ground truths for the training of high-level network, which is a classiﬁer in form.
We make every state-action pair passed directly to each sub-skill and compute
L BC ( θ lk , τ i 1 ) , k = 1 , · · · , K , and choose the ground truth at each time step as the sub-skill category k that minimizes L BC ( θ lk , ( s t , a t )) :
[Equation]
Then we get predicted sub-skill categories from the high-level network:
ˆ z i 1 t = π θ h ( s t ) , and use a cross-entropy loss to train the high-level network:
[Equation]
Finally we perform gradient descent on the high-level network and get λ
h
= θ h − α ∇ θ h L h
( θ h , τ
i 1 ) .
Note θ l 1 , · · · , θ lk are freezed here.
LI: We sample the second batch of trajectories { τ i 2 } from D tri .
The adapted high level network π λ h will process each state in τ i 2 to get sub-skill category ˆ z i 2 t = π λ h ( s t ) at each time step, thus we get
K separate data sets for different sub-skills: D 2 k
= { ( s i 2 t , a i 2 t ) | ˆ z
i 2 t
= k } , k = 1 , · · · , K .
Then we compute the adaptation loss for each sub-skill with the corresponding dataset.
In case of continuous action space, we assume that actions belong to Gaussian distributions, so we have:
[Equation]
where N k is the number of state-action pairs in D 2 k .
Finally we perform gradient descent on sub-skills and get λ lk
= θ lk − α ∇ θ lk L BC ( θ lk , D 2 k ) .
Note π λ h is frozen in this process.
HO:
We sample the third batch of trajectories { τ i 3 } from D tr i and get L ( λ h , τ i 3 ) as in the HI process.
Then we use it to compute the meta-gradient ∇ θ h L ( λ h , τ i 3 ) which equals to:
[Equation]
LO:
We sample τ i 4 and get L ( λ lk , D 4 k ) , k = 1 , · · · , K as in the LI process, then we use it to compute the meta-gradient ∇ θ lk L ( λ lk , D 4 k ) which equals to:
[Equation]
Note after the training of m tasks, we average all meta-gradients from m tasks and perform gradient descents on the initial parameters together to update high-level parameters θ (cid:48)
h
= θ h − β (cid:80) mi =1 ∇ θ h L ( λ h , τ i 3 ) and sub-skill policies parameters θ (cid:48)
lk = θ lk − β (cid:80) mi =1 ∇ θ lk L ( λ lk , τ i 4 ) , k = 1 , · · · , K , i.e., we do not update them at step 5 and 6.
This is crucial to ensure convergence.
For testing, although our method needs two batches of trajectories for one round of adaptation, in practice we ﬁnd only using one trajectory to perform HI and LI also works well in new tasks, thus
DMIL can satisfy the one-shot imitation learning requirement.
Besides the above process, we also add an auxiliary loss to better drive out meaningful sub-skills to avoid the excessively frequently switching between different sub-skills along with time.
Detailed information can be found in C.
Theoretical Analysis
In this section, we show the above algorithm can converge by rewriting both MAML and DMIL as hierarchical variational Bayes problems to establish the equivalence between them since the convergence of MAML can be proved in Fallah et al.
Proofs of all theorems are in Appendix D.
4.1 Hierarchical Variational Bayes Formulation of MAML
According to [14], MAML is a hierarchical variational Bayes inference process.
The general metalearning objective (1), which can be rewritten as L g =
L ( θ, λ 1 , · · · , λ m ) =
log (cid:81) m i =1 p
( D i | θ ) , can be formulated as follows:
[Equation]
where φ i , i = 1 , · · · , m represent the local latent variables for task T i , and λ 1 , · · · , λ M are the variational parameters of the approximate posteriors over φ 1 , · · · , φ M .
We denote λ
i as λ i ( D i , θ ) and p ( φ i |D i , θ ) as p ( φ i |D tri , θ ) to mean that λ
i and φ i are determined with prior parameters θ and support data D tri .
First we need to minimize KL( q ( φ i ; λ i ) (cid:107) p
( φ i |D tri , θ )) w.r.t.
λ i .
According to D.2, we have:
[Equation]
We can establish the connection between 8 and the inner loop in MAML by the following Lemma:
Lemma
In case q ( φ i ; λ i ) is a Dirac-delta function and choosing Gaussian prior for p ( φ i | θ ) , equation 8 equals to the inner-update step of MAML, that is, maximizing log p ( D tri ) w.r.t.
λ
i by early-stopping gradient-ascent with choosing µ θ as initial point:
[Equation]
µ θ . Then we need to optimize L ( θ, λ 1 , · · · , λ M ) w.r.t.
θ .
Since we evaluate p ( D i | λ i ( D tri , θ )) with only
D vali , we assume p ( D i |
λ i ( D tri , θ ))
=
p ( D vali | λ i ( D tri , θ )) .
We give out the following theorem to establish the connection between the meta-update process and the optimization of L g :
Theorem
In case that Σ θ → 0 + , i.e., the uncertainty in the global latent variables θ is small, the following equation holds:
[Equation]
A general EM algorithm will ﬁrst compute the distribution of latent variables (E-step), then optimize the joint distribution of latent variables and trainable parameters (M-step), and the likelihood of data can be proved to be monotone increasing to guarantee the convergence since the evidence lower bound of likelihood is monotone increasing.
Here φ i , i = 1 , · · · , M are the latent variables, and θ corresponds to the trainable parameters.
Lemma 1 and Theorem 1 correspond to the E-step and
M-step respectively.
In the following part we establish the equivalence between 9 with 3 and 4, and between 10 with 5 and 6 to prove the equivalence between DMIL and MAML.
4.2 Modeling DMIL with Hierarchical Variational Bayes Framework
For simplicity, here we only derive in one speciﬁc task T
i , since derivatives of parameters from multi-task can directly add up.
We ﬁrst establish the connection between the maximization of log p ( D tri | θ h , θ l 1 , · · · , θ lK ) with the particular loss functions in DMIL:
Theorem
In case of p ( a t | s t , θ lk ) ∼ N ( µ θ lk ( s t ) , σ 2 ) , we have:
[Equation]
D tri
[Equation] [Equation]
and
[Equation]
where k = 1 , · · · , K. Note in 12, D 2 k corresponds to data sets determined by the adapted high level network λ h , and this connects with 3 and 4 in DMIL.
According to 8, ﬁnding λ
i equals to maximize log p ( D tri | θ ) in speciﬁc conditions, and here in Theorem 2, we prove that maximize log p ( D tri | θ h , θ l 1 , · · · , θ lK ) corresponds to 3 and 4 in DMIL.
Thus theorem 2 corresponds to the Estep of DMIL, where we take τ i 1 and τ i 2 as D tri , and optimize arg max λ i E q ( φ i ; λ i ) [log p ( D tri | φ i )] −
KL( q ( φ i ; λ i ) (cid:107) p
( φ i | θ )) with coordinate descent method, which can be proved to be equal to 9 in D.5.
For the M-step, we take τ i
and τ i 4 as D vali .
According to Theorem 1, we can take the derivative of λ ih , λ il 1 , · · · , λ ilK to maximize the joint distribution of latent variables and trainable parameters to maximize the likelihood of dataset, so we have:
[Equation]
(13) where θ il =
[ θ i 1 , · · · , θ iK ] T and λ il =
[ λ i 1 , · · · , λ iK ] T .
This is exactly the gradients computed in
HO and LO steps.
Note this computation process can be automatically accomplished with standard deep learning libraries such as PyTorch [51].
To this end, we establish the equivalence between DMIL and MAML, and the convergence of DMIL can be proved.
For a clearer comparison, MAML is an iterative process of θ → λ → θ (cid:48) , and DMIL is an iterative process of θ h , θ l → λ h , θ l → λ h , λ l → θ (cid:48)
h , θ (cid:48) l , where the posterior estimation stages θ h , θ l → λ h , θ l → λ h , λ l has no effect on parameters θ h , θ l , thus can be divided to two steps as in DMIL.
This decoupled ﬁne-tuning fashion is exactly what we need to ﬁrst adapt the high-level network and then adapt sub-skills.
If we end-to-end ﬁne-tune parameters like θ h , θ l → λ h , λ l , sub-skills will receive supervisions from an unadapted high-level network, which may provide incorrect classiﬁcations.
Different to this, the meta-updating process λ h , λ l → θ (cid:48)
h , θ (cid:48)
l must be done at the same time, since if we update θ h and θ l successively, the later one will receive different derivative (for example, ∇ θ l log p ( D vali | θ (cid:48) ih , λ il ) ) from derivatives in MAML ( ∇ θ l log p
( D vali | λ ih , λ il ) ), and the equivalence would not be proved.
Experiments
In experiments we aim to answer the following questions: (a) Can DMIL successfully transfer the learned hierarchical structure to new tasks with few-shot new task demonstrations?
(b) Can DMIL achieve higher performance compared to other few-shot imitation learning methods?
(c) What are the effects of different parts in DMIL, such as the skill number K , the bi-level meta-learning procedure, and the continuity constraint?
Codes and video results are provided in supplementary materials.
[Table] [Table]
[Table]
Environments and Baselines
We choose to evaluate DMIL on two representative robot manipulation environments.
The ﬁrst one is
Meta-world benchmark environments [2], which contains 50 diverse robot manipulation tasks, as shown in ﬁg. 6 and ﬁg. 7.
We use both the ML10 suite (10 meta-training tasks and 5 meta-testing tasks) and ML45 suite (45 meta-training tasks and 5 meta-testing tasks) to evaluate our method, and collect 2K demonstrations for each task.
We choose Meta-world since we think a large scale of diverse manipulation tasks can drive semantic skills.
We use the following approaches for comparison in this environment: Option-GAIL : a hierarchical generative adversarial imitation learning method to discover options from unsegmented demonstrations [8].
We use Option-GAIL to evaluate the effect of meta-learning in DMIL.
MIL : a transformer-based meta imitation learning method [40].
We use
MIL to evaluate the effect of hierarchical structures.
MLSH : the meta-learning shared hierarchies method [45] that relearns the high-level network in every new task.
We use MLSH to evaluate the effect of ﬁne-tuning (rather than relearning) the high-level network in new tasks.
PEMIRL :
a contextual meta inverse RL method which transfers the reward function in the new tasks [42].
We use
PEMIRL to show DMIL can transfer to new tasks that have signiﬁcantly different reward functions.
The second one is the Kitchen environment of the D4RL benchmark [18], which contains ﬁve different sub-tasks in the same kitchen environment.
The accomplishment of each episode requires sequentially completions of four speciﬁc sub-tasks, as shown in ﬁg. 9.
We use an open demonstration dataset [48] to train our method.
During training, we exclude interactions with selected objects and at test time we use demonstrations that involve manipulating the excluded object to make them unseen tasks.
We choose this environment to show DMIL can be used in long-horizon tasks.
We use two approaches for comparison in this experiment: SPiRL : an extension of the skill extraction methods to imitation learning over skill space [37]; FIST : an algorithm that extracts skills from ofﬂine data with an inverse dynamics model and a distance function [36].
We use fully-connected neuron networks for both the high-level network and sub-skills.
More details of experiments can be found in appendix F.
5.2 Results
Table 1 shows success rates of different methods in ML10 and ML45 suites with sub-skill number
K = 3 .
We perform 1-shot and 3-shot experiments respectively to show the progressive few-shot performance of different methods.
DMIL achieves the best results in ML10 testing suite and ML training and testing suites.
OptionGAIL achieves high success rates in both ML10 and ML45 training suites.
These results show the adequate capacity of hierarchical structures to ﬁt potential multi-modal behaviors in multi-task demonstrations.
MIL achieves comparable results for all meta-testing tasks but is worse than DMIL.
This shows the necessity of meta-learning processes.
Compared to them,
PEMIRL and MLSH are mediocre among all suites.
This comes from that the reward functions across different tasks are difﬁcult to transfer with few shot demonstrations, and the relearned high-level network of MLSH damages previously learned knowledge.
We also illustrate t-sne results of these methods in ﬁg. 4(a) to further analyze them in appendix E.2.
Table 2 shows the rewards of different methods on four unseen tasks in the Kitchen environment.
FIST-no-FT refers to a variant of FIST that does not use future state conditioning, which makes the comparison fairer.
DMIL achieves higher rewards on two out of four tasks and comparable results on the other two tasks, which exhibits the effectiveness of the bi-level meta-training procedure.
The poor performance of DMIL on the ﬁrst task may come from the choice of skill number K or from low-quality demonstrations.
We perform ablation studies on K in the next section.
Fig.
3 shows curves of sub-skill probabilities along time of two tasks window-close and peg-insertside of Meta-world and the microwave-kettle-top burner-light task in Kitchen environment.
We can see the activation of sub-skills shows a strong correlation to different stages of tasks.
In ﬁrst two tasks, π θ l 0 activates when the robot is closing to something, π θ l 1 activates when the robot is picking up something, and π θ l 2 activates when the robot is manipulating something.
In the third task, π θ lactivates when the robot is manipulating the microwave, π θ l 0 activates when the robot is manipulating the kettle or the light switch, and π θ l 3 activates when the robot is manipulating the burner switch.
This shows that DMIL learns semantic sub-skills from unsegmented multi-task demonstrations.
5.3 Ablation Studies
[Table]
In this part, we perform ablation studies on different skill numbers K .
Due to the limited space, we put ablations on different ﬁne-tuning steps, meta-training processes, continuity constraints, and hard/soft EM choices in appendix E.
Effect of different skill number K : Tableshows the effect of different sub-skill number K in Meta-world experiments.
We can see that a larger K can lead to higher success rates on meta-training tasks, but a smaller K can lead to better results on meta-testing tasks.
This tells us that an excessive number of sub-skills may result in over-ﬁtting on training data, and a smaller K can play the role of regularization.
In Kitchen experiments, we can see similar phenomenons in Table 5.
It is worth noting in both environments, we did not encounter collapse problems, i.e., every sub-skill gets well-trained even when K = 8 in kitchen environment or K = 10 in Meta-world environments.
This is because more sub-skills can help the whole structure get lower loss in the meta-training stage.
However, in our supplementary videos, we can see that sub-skills trained with a large K (for instance, K = 10 in Meta-world environments) are not as semantic as sub-skills trained by a small K (for instance, K = 3 in Meta-world environments) during the execution of a task.
Limitations
The limitations of DMIL come from several aspects, and future works can seek meaningful extensions in these perspectives.
Firstly, DMIL models all tasks as bi-level structures.
However, in real-world situations, tasks may be multi-level structures.
One can extend DMIL to multi-level hierarchical structures like done in recent works [49].
Secondly, DMIL does not capture temporal information in demonstrations.
Future state conditioning in Hakhamaneshi et al.
[36] seems an effective tool to improve few-shot imitation learning performance in long-horizon tasks such as in the Kitchen environments.
Future works can employ temporal modules such as transformer [50] as the high-level network of DMIL to improve its performance. Acknowledgments
We would like to thank Tianren Zhang and Haichuan Gao for their insightful comments of the whole work, and Qualcomm China WRD UR Program.
This work was supported by the National Natural
Science Foundation of China 62176133.
A Code
Code and demonstrations are in https://github.com/HeegerGao/DMIL.
B Algorithm
Require: task distribution p ( T ) , multi-task demonstrations {D i } , i = 1 , · · · , m , initial parameters of high-level network θ h and sub-skill policies θ l 1 , · · · , θ lK , inner and outer learning rate α, β . while not done do
Sample batch of tasks T
i ∼ p ( T ) for all T i do
Sample { τ i 1 } , { τ i 2 } , { τ i 3 } , { τ i 4 } from {D i }
Evaluate ∇ θ h L h ( θ h , τ i 1 ) according to 3 and τ i
Compute adapted parameters of high-level network: λ h
= θ h − α ∇ θ h L h
( θ h , τ
i 1 )
Evaluate ∇ θ lk L BC ( θ lk , D 2 k ) according to 4 and τ i 2 , k = 1 , · · · , K
Compute adapted parameters of sub-skills: λ lk = θ lk − α ∇ θ lk L BC ( θ lk , D 2 k ) , k
= 1 , · · · , K
Evaluate ∇ θ h L T i ( λ h , τ i 3 ) and ∇ θ lk L T i ( λ lk , D 4 k ) , k = 1 , · · · , K end for
Update θ h ← θ h − β ∇ θ h (cid:80) T i ∼ p ( T ) L T i ( λ h , τ i 3 )
Update θ lk ← θ lk − β ∇ θ lk (cid:80) T i ∼ p ( T ) L T i ( λ lk , D 4 k ) , k = 1 , · · · , K end while
C Auxiliary Loss
We adopt an auxiliary loss for DMIL to better drive out meaningful sub-skills by punishing excessive switching of sub-skills along the trajectory.
This comes from an intuitive idea: each sub-skill should be a temporal-extended macro-action, and the high-level policy only needs to switch to different skills few times along a task, as the same idea of macro-action in MLSH [45].
We denote:
[Equation]
and the auxiliary loss is:
[Equation]
Although this operation seems discrete, in practice we can use the operations in modern deep learning framework such as PyTorch [51] to make it differentiable.
We add this loss function to the L h ( θ h , τ i 1 ) and ∇ θ h L ( λ h , τ i 3 ) with a coefﬁcient λ = 1 .
We also perform ablation studies of L aux and results are in table 7.
D Proofs D.1 Proof of Lemma
{+Lemma 
+}
In case q ( φ i ; λ i ) is a Dirac-delta function and choosing Gaussian prior for p ( φ i | θ ) , equation 8 equals to the inner-update step of MAML, that is, maximizing log p ( D tri ) w.r.t.
λ
i by early-stopping gradient-ascent with choosing µ θ as initial point:
[Equation]
Proof: in case of the conditions of Lemma 1, we have:
[Equation]
As stated in [52], ﬁrstly in the case of linear models, early stopping of an iterative gradient descent process of λ equals to the maximum posterior estimation (MAP) [53].
In our case the posterior distribution refers to q ( φ i | λ i ) , and MAML is a Bayes process to ﬁnd the MAP estimate as the point estimate of λ
i ( D tri ; θ ) .
In the nonlinear case, this point estimate is not necessarily the global mode of the posterior, and we can refer to [54] for another implicit posterior distribution over φ i and making the early stopping procedure of MAML acting as priors to get the similar result.
D.2 Proof of Equation
Equation 8 can be written as:
[Equation]
where in MAML we assume p ( D tri | φ i )
=
p ( D tri | φ i , θ ) , and use the joint distribution p ( D tri , φ i | θ ) to replace p ( φ i |D tri , θ ) since we assume that p ( D tri ) subjects to uniform distribution.
Thus 8 can be proved.
D.3 Proof of Theorem
{+Theorem 
+}
In case that Σ θ → 0 + , i.e., the uncertainty in the global latent variables θ is small, the following equation holds:
[Equation]
Proof:
[Equation]
where the ﬁrst approximate equal holds because the VI approximation error is small enough, and the second approximate equal holds because that in case Σ θ → 0 + and assuming λ
i be a neuron network, log p ( λ i ( D tri , θ ) | θ ) ≈ 0 holds almost everywhere, so ∇ θ log p ( λ i ( D tri , θ ) | θ ) ≈ 0 .
Note the condition of theorem 1 is usually satisﬁed since we are using MAML, and the initial parameters θ are assumed to be deterministic.
From another perspective, the right side of above equation is the widely used meta-gradient in MAML, and it is equal to 1 m (cid:80)
[Equation]
D
[Equation]
, which is proved to be converged by [17].
D.4 Proof of Theorem
{+Theorem 
+}
In case of p ( a t | s t , θ lk ) ∼ N ( µ θ lk ( s t ) , σ 2 ) , we have:
[Equation]
and
[Equation]
Proof: since p ( D tri | θ h , θ l 1 , · · · , θ lK ) =
[Equation]
Nt =1 p
[Equation]
a t | s t , θ h , θ l 1 , · · · , θ lK ) p
[Equation]
s t | θ h , θ l 1 , · · · , θ lK ) and the second term is independent of θ , we consider the ﬁrst conditional probability:
[Equation]
In the HI step, p ( a t | s t , θ lk ) is ﬁxed, thus 23 becomes a convex optimization problem:
[Equation]
The solution of this problem is λ ∗ h which satisﬁes p ( z k | s t , a t , λ ∗ h )
= 1 , k = arg max k p (
[Equation]
t | s t , θ lk ) .
This means that π θ h needs to predict the sub-skill category at time step t as k , in which case π θ lk can maximize p ( a t | s t , θ lk ) .
In case we choose π θ h to be a classiﬁer that employs a Softmax layer at the end, minimizing the cross entropy loss 3 equals to maximize 24, thus 21 can be proved.
In the LI step, p ( z k | s t , λ h ) is ﬁxed, and the data sets for optimizing θ l 1 , · · · , θ lK are also ﬁxed as D 2 k
= { ( s ijt , a ijt ) | ˆ
z
i 2 t = k } N k t =1 .
Thus we need to maximize each p ( a t | s t , θ lk ) with D 2 k .
In case of p ( a t | s t , θ lk ) ∼ N ( µ θ lk ( s t ) , σ 2 ) ∝ exp[ − ( a t − π θlk ( s t )) 2 2 σ] , we have max θ lk p ( a t | s t , θ lk ) ⇔ min θ lk ( a t − π θ lk ( s t )) 2 , which leads to the loss function 4, thus 22 can be proved, which ﬁnishes the prove of Theorem 2.
D.5 Proof of the E-step of DMIL
According to Theorem 1, we aim to maximize 17 w.r.t λ i from the initial point θ i with coordinate gradient ascent.
We here need to prove that in DMIL, we could also achieve the global maximum point of λ
i as in MAML.
We ﬁrst give out the following Lemma:
Lemma 2 Let x be the solution found by coordinate gradient descent of f ( x ) .
Let x i , i = 1 , · · · , n be the n coordinate directions used in the optimization process.
If f ( x ) can be decomposed as:
[Equation]
where g ( x ) is a differentiable convex function, and each h i ( x ) is a convex function of the coordinate direction x
i , then x is the global minimum of f ( x ) .
Proof: Let y be another arbitrary point, we have:
[Equation]
Now let’s consider our problem.
Consider log p ( D tri | θ ih , θ il )
=
log
[Equation]
In our case, two coordinate directions are θ ih and θ il .
Let’s consider the terms inside the brackets.
According to Lemma 2, we can think log p ( s t ) as g ( x ) (here it equals to constant), (cid:80)
Kk =1 log p ( z k | s t , a t , θ ih ) as h 1 ( x ) and (cid:80)
Kk =1 log p ( a t | s t , θ il ) as h 2 ( x ) .
Thus the optimum can be proved.
E Additional Ablation Studies E.1 Effects of the Bi-level Meta-learning Process
We use two variants of DMIL to see the effectiveness of bi-level meta-earning process.
DMIL-High : a variant that only meta-learns the high-level network, and DMIL-Low : a variant that only metalearns sub-skills.
We use Option-GAIL as a comparison that does not meta-learn any level of the hierarchical structure.
table 4 shows the results of this ablation study.
DMIL-High achieves close results with OptionGAIL in all meta-training suites and better results in all meta-testing suites, but worse than DMIL in all cases.
This shows that meta-learning the high-level network can help the hierarchical structure adapt to new tasks, but only transferring the high-level network is not enough for accomplish all kinds of new tasks.
DMIL-Low achieves poor results in all suites except in ML10 meta-training suites.
This shows that transferring the high-level network is necessary when training on a large scale of tasks or testing in new tasks.
[Table]
For better understand the effect of bi-level meta-learning process, we perform ablation study for three DMIL variants in an manually-designed new task push-around-wall (ﬁg. 8).
In this task, the robot needs to grasp a cube and circle it around the wall.
This is a brand new skill that is not in the meta-world suite.
The accomplishment of this new task requires quickly adapting abilities of both the high-level network and sub-skills.
We sample two demonstrations and use the ﬁrst one as few-shot data, and illustrate the sub-skill categories of the second demonstration given by the high-level network at ﬁg. 4(b).
Before adaptation, all variants give out approximately random results.
However, after one-shot adaptation, DMIL classiﬁes almost every state into sub-skill 0 and sub-skill 2, which indicates these two sub-skills in DMIL have been adapted to the new task, and the high-level
[Figure]
network has also been adapted with the supervision from adapted sub-skills.
Compared to DMIL,
DMIL-Low still can not give out reasonable results after adaptation, since its high-level network lacks the ability to quickly transfer to new tasks.
Instead, DMIL-High gives out plausible results after adaptation.
This shows the high-level network has adapted to the new task according to the supervision from adapted sub-skills, but no sub-skill can dominate for a long time period since all sub-skills lack the quickly adaptation ability.
E.2 T-sne Results of Different Methods
For comparison of different methods, we illustrate the t-sne results of states of each sub-skill in an
ML45 meta-testing task hand-insert in ﬁg. 4(a).
We use 3 demonstrations for adaptation, and draw the t-sne results on another 16 demonstrations.
This task is a meta-testing task, so no method has ever been trained on this task before.
MLSH shows almost random clustering results no matter before and after adaptation, since its high-level network is relearned in new tasks.
OptionGAIL clusters to three sub-skills after adaptation.
Compared to them, DMIL clusters the data to only two sub-skills after adaptation.
We believe fewer categories reﬂect more meaningful sub-skills are developed in DMIL.
E.3 Effects of Sub-skill Number K in The Kitchen Environment
We perform ablation studies of sub-skill number K on the Kitchen environments and choose K = 2 , 4 , 8 respectively.
table 5 shows the results.
We can see that a smaller number of sub-skills can achieve better results on such four unseen results that a large number of sub-skills.
This may indicate that the sub-skill number K can work as a ’bottleneck’ like the middle layer in an auto-encoder.
[Table]
E.4 Effects of Fine-tuning Steps
As all few-shot learning problems, the ﬁne-tuning steps in new tasks to some extent determine the performance of the trained model.
It controls the balance between under-ﬁtting and over-ﬁtting.
We perform ablation studies of ﬁne-tune steps in Meta-world benchmarks with K = 5 and lr=1e-2 in table 6.
Results are as follows:
[Table]
E.5 Effects of Continuity Regularization
We perform ablation studies of the effect of continuity regularization as following table 7.
DMIL nc means no continuity regularization.
We can see that the continuity constraint would damage metatraining performance slightly, but increase the meta-testing performance greatly
[Table]
[Table]
E.6 Effects of Hard/Soft EM Choices
In DMIL
[Table]
we use hard EM algorithm to train the high-level network.
One may think about to use soft cross entropy loss to train the high-level network to get better results.
We perform this ablation study in the following table 8.
We can see that a soft cross entropy training won’t help increase the whole success rates.
This may comes from that, usually we use soft cross entropy (such as label smoothing) to prevent over-ﬁtting.
However, in our situation, this may cause under-ﬁtting, since training on such
[Table]
[Table]
a large scale of diverse manipulation tasks is already very difﬁcult.
Future works can seek more comparisons about this choice.
Additionally
[Table]
we found an interesting phenomena that the training loss of the high-level network with a softmax shows a trend of rising ﬁrst and then falling, as shown in ﬁg. 5.
In our experiments, a softmax loss may regularize the optimization process to make the high-level network be under-ﬁtting on training data.
This may come from that, the experiment environment (Meta-world) contains a large scale of manipulation tasks, in which the training of the high-level network can be difﬁcult and unstable.
Thus a soft-max cross entropy loss cannot help that much here like how it works as a regularizer to prevent over-ﬁtting in the label smoothing [55].
F Experiment Details F.1 Environments
See ﬁg.
7,ﬁg.
6,ﬁg.
9 and ﬁg. 8.
F.2 Model Setup
DMIL:
The high-level network and each sub-skill is modeled with a 4-layer fully-connected neuron network, with 512 ReLU units in each layer.
We use Adam as the meta-optimizer.
DMIL-High and
DMIL-Low use the same architecture with DMIL.
Hyper-parameter settings are available in table 9.
MIL: We use a transformer [50] as the policy to perform MAML.
The input of the encoder is the whole one-shot demonstration or 3-shot demonstrations with concatenated state and action.
The input of the decoder is current state.
The output is the predicted action.
Hyper-parameter settings are available in table 10.
[Table]
MLSH: We use the same settings of network with DMIL here.
The macro step of high-level network is 3.
Since our problem is not a reinforcement learning process, we use the behavior cloning variant of MLSH.
The pseudo reward is deﬁned by the negative mean square loss of the predicted action and the ground truth, and we perform pseudo reinforcement learning process with off-policy demonstration data.
We use PPO as our reinforcement learning algorithm.
Note in this way we ignore the importance sampling weights that required by replacing the sampling process in the environments with the demonstrations in the replay buffer, which has been shown to be effective in practice in [56, 57].
Hyper-parameter settings are available in table 11.
PEMIRL:
We use the same setting of the high-level network as the policy π ω and the inference model q ψ in PEMIRL.
We use PPO as our reinforcement learning algorithm.
We use a 3-layer fully-connected neuron network as the context-dependent disentangled reward estimator r θ ( s, m ) and the context-dependent potential function h φ ( s, m ) .
Here we also use the behavior cloning variant of PEMIRL to only train models on the off-policy data.
Hyper-parameter settings are available in table 12.
F.3 Training Details
For ﬁne-tuning, OptionGAIL, DMIL-Low and DMIL-High have no meta-learning mechanism for (some parts of) the trained model.
In the few-shot adaptation process, we have different ﬁne-tune method for these baselines:
OptionGAIL: We train OptionGAIL models on the provided few-shot demonstrations for a few epochs.
DMIL-Low: We ﬁx the high-level network and only ﬁne-tune sub-skills with few-shot demonstrations.
DMIL-High : We ﬁx sub-skills and only ﬁne-tune the high-level network with few-shot demonstrations.
[Table]
[Table]