# One-shot Imitation in a Non-Stationary Environment via Multi-Modal Skill

Sangwoo Shin  Daehee Lee  Minjong Yoo  Woo Kyung Kim  Honguk Woo

###### Abstract

One-shot imitation is to learn a new task from a single demonstration, yet it
is a challenging problem to adopt it for complex tasks with the high domain
diversity inherent in a non-stationary environment. To tackle the problem, we
explore the compositionality of complex tasks, and present a novel skill-based
imitation learning framework enabling one-shot imitation and zero-shot
adaptation; from a single demonstration for a complex unseen task, a semantic
skill sequence is inferred and then each skill in the sequence is converted
into an action sequence optimized for environmental hidden dynamics that can
vary over time. Specifically, we leverage a vision-language model to learn a
semantic skill set from offline video datasets, where each skill is
represented on the vision-language embedding space, and adapt meta-learning
with dynamics inference to enable zero-shot skill adaptation. We evaluate our
framework with various one-shot imitation scenarios for extended multi-stage
Meta-world tasks, showing its superiority in learning complex tasks,
generalizing to dynamics changes, and extending to different demonstration
conditions and modalities, compared to other baselines.

Machine Learning, ICML

##  1 Introduction

In the context of learning control tasks for autonomous agents, one-shot
imitation aims at learning such a task on a single demonstration presented
(Duan et al., 2017; Finn et al., 2017; Li et al., 2022). One-shot imitation is
particularly desirable, as it allows efficient use of expert demonstrations,
which are often costly to collect for diverse tasks. However, it is considered
challenging to achieve reliable one-shot imitation for complex tasks across
various domains and environments (Pertsch et al., 2022; Jang et al., 2022;
Brohan et al., 2022).

In this paper, we investigate skill-based imitation learning techniques to
conduct one-shot imitation for long-horizon tasks in a non-stationary
environment, where underlying dynamics can be changed over time. To do so, we
leverage the compositionality in control tasks and explore the semantic
representation capability of a large scale vision-language pretrained model.
Specifically, we develop a novel one-shot imitation learning framework using
semantic skills, OnIS to enable one-shot imitation for new long-horizon tasks
across different dynamics in a non-stationary environment. The framework
employs a vision-language pretrained model and contrastive learning on expert
data to represent semantic skills and environment dynamics separately, and
explores meta-learning to adapt semantic skills to unseen dynamics. This
establishes the capability of dynamics-aware skill transfer.

In deployment, given a single video demonstration for a new task, the
framework translates it to a sequence of semantic skills through a sequence
encoder built on the vision-language model. It then combines the skills with
inferred current dynamics to generate actions optimized for not only the
demonstration but also the current dynamics.

As such, through the decomposition of skills and dynamics, each skill can be
well represented in the vision-language aligned embedding space without much
deformation by time-varying environment dynamics. This decomposition structure
with multi-modality renders OnIS unique to explore the compositionality of
control tasks with skills, different from prior skill-based reinforcement
learning (RL) approaches, e.g., (Pertsch et al., 2021; Nam et al., 2022).

The main contributions of our work are as follows.

  * •

We present the OnIS framework to enable one-shot imitation and zero-shot
adaptation.

  * •

We develop a dynamics-aware skill transfer scheme to adapt skills in a non-
stationary environment.

  * •

We create an expert dataset for long-horizon, multi-stage Meta-world tasks
with diverse robotic manipulation skills, and make it publicly available for
other research works.

  * •

We test the OnIS framework with several imitation learning use cases, and
evaluate it in terms of one-shot imitation performance, generalization ability
to different environment conditions, and extensibility to demonstrations in
different modalities.

##  2 Problem Formulation

In vision-based tasks, one-shot imitation aims at learning a task upon a
single expert demonstration presented in a video clip. Among several related
problems in one-shot imitation learning approaches, our work considers a
challenging yet practical problem setting, inherent in tackling complex tasks
via imitation in a non-stationary environment.

Unlike recent one-shot imitation approaches, our problem setting requires the
followings to be addressed in a unified framework. (1) A policy should not
only be learned from a single expert demonstration, but also adapt to unseen
dynamics conditions inherent in a non-stationary environment, (2) an expert
demonstration specifies a long-horizon multi-stage task, and (3) an expert
demonstration can be presented either in video or language. Here, we formally
describe our problem, “one-shot imitation and zero-shot adaptation” in a non-
stationary environment.

Figure 1: Multi-stage Meta-world task

Given a demonstration $d^{\mathcal{T}}$ for task
$\mathcal{T}\sim\mathbb{T}_{\text{train}}$, we consider a model $\phi$ which
maps $d^{\mathcal{T}}$ to a policy $\pi_{\mathcal{T}}=\phi(d^{\mathcal{T}})$.
Then, the objective of one-shot imitation is to find the optimal model
$\phi^{*}$ such that

$\underset{\phi}{\text{argmax}}\left[\underset{\mathcal{T}^{\prime}\sim\mathbb{T}_{\text{eval}}}{\mathbb{E}}\left[\text{E}(\mathcal{T}^{\prime},\phi(d^{\mathcal{T}^{\prime}}))\right]\right]$
(1)

where E is an evaluation function for a policy
$\phi(d^{\mathcal{T}^{\prime}})$ against a new task $\mathcal{T}^{\prime}$.

To deal with one-shot imitation upon various dynamics, we consider a problem
setting of a Hidden Parameter Markov decision process without reward (HiP-
MDP$\setminus\text{R}$) with a tuple
$(\mathcal{S},\mathcal{A},\mathcal{P},\gamma,\mathcal{H},P_{\mathcal{H}})$.
Here, $\mathcal{S}$ is a state space, $\mathcal{A}$ is an action space,
$\mathcal{P}$ is a transition probability, $\gamma$ is a discount factor, and
$\mathcal{H}$ is a hidden parameter space, where $h\in\mathcal{H}$ determines
the dynamics $s^{\prime}\sim\mathcal{P}(s,a;h)$ for states
$s,s^{\prime}\in\mathcal{S}$ and action $a\in\mathcal{A}$. $P_{\mathcal{H}}$
is a prior distribution of $\mathcal{H}$. Then, the objective of one-shot
imitation and zero-shot adaptation across various dynamics is to find the
optimal model $\phi^{*}$ such that

$\underset{\phi}{\text{argmax}}\left[\underset{h\sim
P_{\mathcal{H}}}{\mathbb{E}}\left[\underset{\mathcal{T}^{\prime}\sim\mathbb{T}_{\text{eval}}}{\mathbb{E}}\left[\text{E}_{h}({\mathcal{T}^{\prime}},\phi(d^{\mathcal{T}^{\prime}}))\right]\right]\right]$
(2)

where $\text{E}_{h}$ is an evaluation function for a policy
$\phi(d^{\mathcal{T}^{\prime}})$ against task $\mathcal{T}^{\prime}$ with
fixed $h$ in a HiP-MDP$\setminus$R.

Regarding task definitions, we consider long-horizon, multi-stage tasks, in
the form of $K$-staged tasks $\mathcal{T}_{0:K-1}$ (e.g., in Figure 1, a
4-stage task with sequential subtasks, slide puck and others). Note that
$\mathcal{T}_{0:K-1}$ should be done in a correct order, e.g., slide puck,
push button, close drawer, and then open door, specifying that the evaluation
function $\text{E}_{h}$ yields some positive value for performing subtask
$\mathcal{T}_{j}$ only when subtasks $\mathcal{T}_{0:j-1}$ have been
completed.

##  3 Our Approach

###  3.1 Overall Framework

To enable the one-shot imitation and zero-shot adaptation in a non-stationary
environment, we develop the OnIS framework, leveraging the semantic
representation capability of a large scale vision-language pretrained model
and compositionality of complex tasks.

As illustrated in Figure 2, the framework has two phases: (a) training phase
and (b) deployment phase. In the training phase, we leverage task
compositionality and pretrained vision-language models to extract semantic
skills from video demonstrations. Specifically, we explore skill decomposition
(task compositionality) for dealing with complex tasks where each task
consists of sequential skills (achievable goals). Each task is decomposed and
translated into a sequence of skills that are task-agnostically learned on
offline datasets. From demonstrations for various tasks and diverse dynamics,
the framework employs contrastive learning with video and text retrieval tasks
to disentangle dynamics-invariant task-relevant features from the
demonstrations, thus representing them in a vision-language semantic space,
i.e., in the CLIP embedding space (Radford et al., 2021).

At the same time, for dealing with adaptation across different dynamics, we
develop a novel skill representation, dynamics-aware skills. The framework
uses contrastive learning with dynamics reconstruction tasks to disentangle
dynamics information from the trajectories in state and action pairs. These
two individual contrastive learning procedures tend to decompose unstructured
features contained in the demonstrations and trajectories into semantic skill
sequences and environment dynamics, respectively. Then, to render a skill
sequence adaptive to different dynamics, the decomposed skill sequence and
dynamics are combined in training the skill transfer capability through a
meta-learning procedure.

Figure 2: One-shot imitation and zero-shot adaptation in OnIS: (a) In the
training phase, a semantic skill set is established from video demonstrations
in offline datasets by leveraging a pretrained vision-language model. For each
demonstration in the dataset, semantic skill sequence and environment dynamics
are first disentangled, and then they are combined to enable dynamics-aware
skill transfer to different environments. (b) In the deployment phase, given a
single demonstration for an unseen task, its semantic skill sequence can be
immediately inferred (one-shot imitation), where each in the sequence adapts
temporally to the hidden, time-varying dynamics in the environment (zero-shot
adaptation).

(b) In the deployment phase, given a single video demonstration for an unseen
task, the framework derives such a policy not only imitating the demonstrated
expert behavior to complete the unseen task but also being capable of adapting
to dynamics changes in the environment. In this phase, the framework uses only
a single demonstration (one-shot imitation), and it enables rapid adaptation
to dynamics changes (zero-shot adaptation) in a non-stationary environment, by
making use of the knowledge constructed in the training phase.

For one-shot imitation, in the training phase, we exploit offline datasets of
video demonstrations and their matched trajectories in state and action pairs,
presuming the availability of such datasets from the environment. The
contrastive learning procedures on the datasets and the meta-learning
procedure together enable to construct the knowledge required for one-shot
imitation in a non-stationary environment. The contrastive learning procedures
are implemented in the framework, as a semantic skill sequence encoder and a
dynamics encoder, and the meta-learning procedure is implemented as a skill
transfer module. We explain these modules in Figure 3, and Sections 3.2 and
3.3.

Figure 3: OnIS framework: In (a), the semantic skill sequence encoder
$\Phi_{enc}$ and the semantic skill decoder $\Phi_{dec}$ are trained offline
using the CLIP vision-language pretrained model, where $\Phi_{enc}$ translates
video demonstrations to semantic skill sequences and is contrastively learned,
and $\Phi_{dec}$ learns to infer an optimal skill (from a sequence) upon a
state. The language prompt $\theta_{p}$ is only used for the subtask-level
instruction (S-OnIS) case, and the additional encoders $\theta_{v},\theta_{l}$
are only used for episode-level instruction case (U-OnIS). In (b), the skill
transfer $\pi_{tr}$ and the dynamics encoder $\psi_{enc}$ are trained offline,
where $\pi_{tr}$ learns to infer an action sequence optimized for the
deployment setting from a given semantic skill sequence and inferred dynamics,
and $\psi_{enc}$ learns to infer dynamics from sub-trajectories. These modules
establish dynamics-aware skill transfer. In (c), for a given demonstration,
$\Phi_{enc}$ first infers a sequence of semantic skills, $\Phi_{dec}$ infers a
current semantic skill, and $\psi_{enc}$ infers current dynamics in the non-
stationary deployment environment. Then, $\pi_{tr}$ yields actions optimized
through the current semantic skill and dynamics.

###  3.2 Learning Semantic Skill Sequences

Semantic skill sequence encoder. A semantic skill sequence encoder
$\Phi_{enc}$ is responsible for mapping an expert demonstration $d$ to a
sequence of semantic skills, in that $d$ is segmented into a sequence of
dynamics-invariant behavior patterns (action sequences). Each pattern
corresponds to a short sequence of expert actions and it can be described in a
language instruction in the environment. Such an expert behavior pattern
associated with some language instruction is referred to as a semantic skill,
as for an expert behavior pattern, its associated language instruction is used
to represent expert behaviors on the semantic embedding space of a vision-
language aligned model.

To implement the encoder $\Phi_{enc}$, we use the CLIP vision-language
pretrained model and sample-efficient prompt learning techniques (Zhou et al.,
2022). Specifically, we assume that a dataset of expert trajectories
$\mathcal{D}=\\{\tau_{1},\tau_{2},\cdots,\tau_{N}\\}$ contains language
instruction data, in that each $T$-length trajectory
$\tau=\\{(s_{1},v_{1},l_{1},a_{1}),\cdots,(s_{T},v_{T},l_{T},a_{T})\\}$
consists of state $s$, visual observation $v$, language instruction $l$, and
action $a$, where each language instruction is an element of the instruction
set $\mathbf{L}$, similar to (Zheng et al., 2022).

Furthermore, we consider two distinct cases of language instructions annotated
in the trajectories, one is the subtask-level instruction, (e.g.,“Push lever”
is annotated on the transitions for its relevant single subtask, as in
(Pertsch et al., 2022)), and the other is the episode-level instruction,
(e.g., “Push lever, open door and close the box” is annotated on all the
transitions in an episode, as in (Garg et al., 2022; Liu et al., 2022)).

For the subtask-level instruction case, we assume that language instructions
exist only for a small subset $\mathcal{D}_{0}$ of $\mathcal{D}$, and consider
the CLIP model with a video encoder $\Phi_{V}$ and a language encoder
$\Phi_{L}$ such that

$\begin{split}\Phi_{V}:v_{t:t+H}\mapsto z_{v},\ \ \ \Phi_{L}:l_{t}\mapsto
z_{l}.\end{split}$ (3)

These $\Phi_{V}$ and $\Phi_{L}$ encoders map a video demonstration $v_{t:t+H}$
and a language instruction $l_{t}$ into the same embedding space for timesteps
$t$. To construct $\Phi_{enc}$ based on $\Phi_{V}$ and $\Phi_{L}$ using a
small number of samples, we use the language prompt (Zhou et al., 2022). Given
$v_{t:t+H}$ or $l_{t}$ as input $d$, the semantic skill sequence encoder
$\Phi_{enc}$ learns to translate each input into a sequence of semantic skills
by

$\begin{split}\Phi_{enc}(d)=\begin{cases}\underset{z\in
z_{\mathbf{L}}}{\mathrm{argmax}}\\{\mathrm{sim}(z,\Phi_{V}(v_{t:t+H}))\\},&d=d_{v}\\\
\\{\Phi_{L}(l_{t};\theta_{p})\\},&d=d_{l}\end{cases}\end{split}$ (4)

where $d_{v}=(v_{t:t+H}:0\leq t\leq T^{\prime})$, $d_{l}=(l_{t}:0\leq t\leq
N)$, $z_{\mathbf{L}}=\Phi_{L}(\mathbf{L};\theta_{p})$, and $\mathbf{L}$
denotes a set of language instructions. Specifically, a language prompt
$\theta_{p}$ in learnable parameters are trained via contrastive learning on
positive pairs $(v_{t:t+H},l_{t})$ in $\mathcal{D}_{0}$, where the contrastive
loss is defined as $\mathcal{L}_{CON}(\theta_{p})=$

$\begin{split}-\log\left(\frac{\mathrm{sim}(\Phi_{V}(v_{t:t+H}),\Phi_{L}(l_{t};\theta_{p}))}{\sum_{l_{t}\neq
l\in\mathbf{L}}\mathrm{sim}(\Phi_{V}(v_{t:t+H}),\Phi_{L}(l;\theta_{p}))}\right).\end{split}$
(5)

The similarity of latent vectors $z$ and $z^{\prime}$ is calculated by

$\mathrm{sim}(z,z^{\prime})=\frac{1}{\alpha}\exp\left(\frac{\langle
z,z^{\prime}\rangle}{\|z\|\|z^{\prime}\|}\right)$ (6)

where $\alpha$ is a temperature coefficient. As such, for the subtask-level
instruction case, the semantic skill sequence encoder $\Phi_{enc}$ can be
established from the pretrained visual and language encoders $\Phi_{V}$,
$\Phi_{L}$ through prompt-based contrastive learning on a small set of video
and text retrieval samples.

In the case of episode-level instruction, all the transitions in an episode in
$\mathcal{D}$ are associated only with a single instruction without detailed
subtask-level instructions. Thus, we adapt unsupervised skill learning
techniques in (Garg et al., 2022) with contrastive learning on video features.
Accordingly, $\Phi_{enc}(d)$ learns to translate a video demonstration
$d=v_{0:T^{\prime}}$ or language instruction $d=l$ to a sequence of semantic
skills by

$\begin{split}\begin{cases}\\{\underset{z\in
z_{\mathbf{L}}}{\mathrm{argmax}}\\{\mathrm{sim}(z,\theta_{v}(\Phi_{V}(v_{0:T^{\prime}}),s_{t-H:t})\\}\\},&d=v_{0:T^{\prime}}\\\
\\{\theta_{l}(\Phi_{L}(l),s_{t-H:t})\\},&d=l\end{cases}\end{split}$ (7)

where $z_{\textbf{L}}=\theta_{l}(\Phi_{L}(\mathbf{L}),s_{t-H:t})$. Similar to
the skill predictor used in (Garg et al., 2022), two additional encoders
$\theta_{v}$ and $\theta_{l}$ such as

$\begin{split}&\theta_{v}:(\Phi_{V}(v_{0:T^{\prime}}),s_{t-H:t})\mapsto
z_{v}\\\ &\theta_{l}:(\Phi_{L}(l),s_{t-H:t})\mapsto z_{l}\end{split}$ (8)

are trained. Specifically, the encoder $\theta_{v}$ is contrastively learned
on positive pairs $(v_{0:T^{\prime}},l)$ in $\mathcal{D}$, where the
contrastive loss in (5) is rewritten as $\mathcal{L}_{CON}(\theta_{v})$ =

$\begin{split}-\log\left(\frac{\mathrm{sim}(\theta_{v}(\Phi_{V}(v_{0:T^{\prime}}),s_{t-H:t}),z_{l})}{\sum_{z_{l}\neq
z\in
z_{\mathbf{L}}}\mathrm{sim}(\theta_{v}(\Phi_{V}(v_{0:T^{\prime}}),s_{t-H:t}),z)}\right)\end{split}$
(9)

for $z_{l}=\theta_{l}(\Phi_{L}(l),s_{t-H:t})$ and
$z_{\textbf{L}}=\theta_{l}(\Phi_{L}(\mathbf{L}),s_{t-H:t})$. The encoder
$\theta_{l}$ is trained via behavior cloning that maximizes the mutual
information between the actions and semantic skills, where the loss is defined
as

$\mathcal{L}_{BC}(\theta_{l},f)=\mathbb{E}[\|a_{t}-f(s_{t},z_{l})\|]$ (10)

for the action reconstruction model $f:(s_{t},z_{l})\mapsto a_{t}$.

Semantic skill decoder. Given a sequence of semantic skills, the semantic
skill decoder $\Phi_{dec}$ is responsible for obtaining a semantic skill for a
current state. $\Phi_{dec}$ is implemented as a binary model to determine
whether the current semantic skill is terminated or not, similar to the task
specification interpreter in (Xu et al., 2018), i.e.,

$\Phi_{dec}:(s_{t_{0}},s_{t},z_{t})\mapsto\\{0,1\\}$ (11)

for a semantic skill $z_{t}$ being currently executed, the initial state
$s_{t_{0}}$ for $z_{t}$, and a current state $s_{t}$. $\Phi_{dec}$ is learned
on $\mathcal{D}$ with binary cross entropy (BCE),

$\mathcal{L}_{BCE}(\Phi_{dec})=\text{BCE}(\mathbbm{1}_{z_{t}\neq
z_{t+1}},\Phi_{dec}(s_{t_{0}},s_{t},z_{t}))$ (12)

where $z_{t}=\Phi_{v}(v_{t:t+H})$.

As illustrated in Figure 3(a), the semantic skill sequence encoder
$\Phi_{enc}$ is trained through contrastive learning to extract the inherent
semantic skills in the video. The semantic skill decoder $\Phi_{dec}$, which
is basically a binary classifier, is conditioned on a skill sequence and is
trained to predict the appropriate semantic skill at the current timestep.

###  3.3 Learning Dynamics-aware Skill Transfer

Given a semantic skill sequence encoder $\Phi_{enc}$, the skill transfer
module $\pi_{tr}$ is responsible for transferring a semantic skill
$z_{t}=\Phi_{enc}(v_{t:t+H})$ to an action sequence adapted for environment
dynamics. That is,

$\pi_{tr}:(s_{t},z_{t},h^{q}_{t})\mapsto a_{t}$ (13)

for a current state $s_{t}$, a semantic skill $z_{t}$, and the dynamics
embedding $h^{q}_{t}=\psi_{enc}(\tau_{t})$. For $H_{0}$-length sub-
trajectories $\tau_{t}=(s_{t-H_{0}},a_{t-H_{0}},...,s_{t-1},a_{t-1})$, the
dynamics encoder $\psi_{enc}$ takes it as input, and maps to a quantized
vector $h^{q}_{t}$, i.e.,

$\psi_{enc}=\mathbf{q}\circ\psi^{c}_{enc}:\tau_{t}\mapsto h^{q}_{t}$ (14)

where $\psi^{c}_{enc}$ maps $\tau_{t}$ to a continuous latent vector. We use
the vector quantization operator $\mathbf{q}$ to avoid the posterior collapse
(van den Oord et al., 2017). The output of $\mathbf{q}$ is the nearest vector
among the learnable parameters, called codebook,
$h^{q}\in\\{h^{q_{1}},\cdots,h^{q_{M}}\\}$. Then, the quantized dynamics
embedding $h^{q}_{t}$ is obtained by

$h^{q}_{t}=\mathbf{q}(\psi^{c}_{enc}(\tau_{t}))=\underset{j\in\\{1,\cdots,M\\}}{\textnormal{argmin}}\\{\|\psi^{c}_{enc}(\tau_{t})-h^{q_{j}}\|\\}.$
(15)

For a semantic skill $z_{t}=\Phi_{enc}(v_{t:t+H})$ and dynamics embedding
$h^{q}_{t}=\psi_{enc}(\tau_{t})$, the skill transfer module and dynamics
encoder are jointly trained by minimizing the behavior cloning (BC) loss
$\mathcal{L}_{BC}(\pi_{tr},\psi_{enc})=$

$\begin{split}\mathbb{E}_{s_{T},a_{T}\sim\mathcal{D}}\left[\|a_{t}-\pi_{tr}(s_{t},z_{t},\psi_{enc}(\tau_{t}))\|^{2}\right].\end{split}$
(16)

Furthermore, to disentangle task-irrelevant dynamics from sub-trajectories, we
also use contrastive learning on sub-trajectories in various dynamics.
Specifically, suppose that a batch of $H_{0}$-length sub-trajectories
$\\{\tau_{t_{i}}\\}_{1\leq i\leq N}$ contains one positive sample
$\\{(\tau_{t_{j}},\tau_{t_{k}})\\}$ which comes from the same trajectory
starting at different timesteps. Then, $\psi_{enc}$ is learned on the positive
and negative pair, where the contrastive loss is defined as
$\mathcal{L}_{CON}(\psi_{enc})=$

$\begin{split}&-\log\left(\frac{\mathrm{sim}(\psi_{enc}(\tau_{t_{j}}),\psi_{enc}(\tau_{t_{k}}))}{\sum_{i\neq
i^{\prime}}\mathrm{sim}(\psi_{enc}(\tau_{t_{i}}),\psi_{enc}(\tau_{t_{i^{\prime}}}))}\right).\end{split}$
(17)

To maximize the mutual information between the inferred embedding $h^{q}_{t}$
and underlying dynamics, we adopt reconstruction-based feature extraction by
using the inverse dynamics decoder
$\psi_{dec}:(s_{t},s_{t+1},h_{t}^{q})\mapsto a_{t}$, where the action
reconstruction loss is defined as

$\begin{split}&\mathcal{L}_{REC}(\psi_{enc},\psi_{dec})\\\
&=\mathbb{E}_{t-H_{0}\leq
i<t}\left[\|a_{i}-\psi_{dec}(s_{i},s_{i+1},\psi_{enc}(\tau_{t})\|^{2}\right].\end{split}$
(18)

Algorithm 1 Learning to transfer skills

1: Semantic skill sequence encoder $\Phi_{enc}$

2: Dynamics encoder $\psi_{enc}$, Inverse dynamics decoder $\psi_{dec}$

3: Skill transfer module $\pi_{tr}$, Dataset $\mathcal{D}$

4: repeat

5: Sample a batch $\\{(\tau_{t_{i}},v_{t_{i}:t_{i}+H})\\}_{i}\sim\mathcal{D}$

6: $\\{z_{t_{i}}\\}_{i}=\Phi_{enc}(\\{v_{t_{i}:t_{i}+H}\\}_{i})$

7: /* Calculate loss with $\\{(\tau_{t_{i}},z_{t_{i}})\\}$ */

8: $\text{loss}_{bc}\leftarrow\mathcal{L}_{BC}(\pi_{tr},\psi_{enc})$ using
(16)

9: $\text{loss}_{con}\leftarrow\mathcal{L}_{CON}(\psi_{enc})$ using (17)

10: $\text{loss}_{rec}\leftarrow\mathcal{L}_{REC}(\psi_{enc},\psi_{dec})$
using (18)

11:
$\psi_{enc}\leftarrow\psi_{enc}-\nabla_{\psi_{enc}}(\text{loss}_{bc}+\text{loss}_{con}+\text{loss}_{rec})$

12: $\psi_{dec}\leftarrow\psi_{dec}-\nabla_{\psi_{dec}}\text{loss}_{rec}$

13: $\pi_{tr}\leftarrow\pi_{tr}-\nabla_{\pi_{tr}}\text{loss}_{bc}$

14: until converge

Overall, the skill transfer module $\pi_{tr}$, the dynamics encoder
$\psi_{enc}$, and the inverse dynamics decoder $\psi_{dec}$ are jointly
trained to minimize the losses $\mathcal{L}_{BC},\mathcal{L}_{CON}$ and
$\mathcal{L}_{REC}$, where these modules are presented in Figure 3(b). During
the deployment phase, the skill transfer module acts as a policy network that
determines actions ($\hat{a}_{t}$) upon the skill and dynamics embeddings
($z_{t}$ and $h^{q}_{t}$) along with the current state ($s_{t}$), as
illustrated in Figure 3(c).

Algorithm 1 lists the learning procedures for skill transfer in Section 3.3.

##  4 Evaluations

In this section, we evaluate the performance of our OnIS framework under
various configurations of non-stationary environments.

Environment settings. For evaluation, we devise multi-stage robotic
manipulation tasks, namely multi-stage Meta-world, using the Meta-world
simulated benchmark (Yu et al., 2020), where each multi-stage task is composed
of a sequence of existing Meta-world tasks (subtasks). An agent manipulates a
robot arm using actions such as slide puck, close drawer, etc., to achieve
certain objectives. For emulating non-stationary environments, we exploit
kinematic parameters affecting the dynamics of robotic manipulation, which are
not explicitly revealed through states or videos. This non-stationarity is
consistent with the prior works in robotics (Kumar et al., 2021; Qi et al.,
2022).

Offline datasets. To generate the dataset that covers diverse dynamics, we
implement rule-based expert policies to collect trajectories from 13 different
stationary environment conditions for each subtask, where each condition
corresponds to a specific kinematic configuration. We collect $3120$ episodes
(240 for each environment condition), and we annotate $24$ episodes (0.76% of
our dataset) with subtask-level language instructions. This annotated data is
only used to train the semantic skill sequence encoder, especially for S-OnIS
(in Section 4.1).

###  4.1 OnIS Implementation

We implement our OnIS framework using the opensource project Jax (Bradbury et
al., 2018). OnIS is composed of 4 modules: semantic skill sequence encoder
$\Phi_{enc}$, skill decoder $\Phi_{dec}$, dynamics encoder $\psi_{enc}$, and
skill transfer module $\pi_{tr}$. We use the Transformer-based CLIP pretrained
model (Radford et al., 2021) to implement the sequence encoder. Specifically,
we implement two versions of OnIS: one with supervised semantic skills
(S-OnIS) and the other with unsupervised semantic skills (U-OnIS). The former
S-OnIS corresponds to subtask-level instructions where contrastive learning
for semantic skills is driven by the supervision of subtask-level instructions
(with sub-trajectory and skill pairs). The latter U-OnIS corresponds to
episode-level instructions where contrastive learning for semantic skills is
driven by the supervision of episode-level instructions (without sub-
trajectory and skill pairs).

###  4.2 Baselines

  * •

BC-Z (Jang et al., 2022) is a state-of-the-art imitation learning framework,
which utilizes both video demonstrations and language instructions. To extract
semantic skill information from multi-modal data, BC-Z is trained to maximize
the similarity between a visual and language embedding encoded by the
pretrained language model.

  * •

Decision Transformer (DT) (Chen et al., 2021) is a Transformer-based model
tailored for imitation learning. Considering that DT is not originally meant
for multi-modal data, we modify DT to have state vectors in different
modalities, along with either a video demonstration or language instruction.
We use DT to compare OnIS with conventional imitation learning approaches
which do not consider skill semantics.

  * •

SPiRL (Pertsch et al., 2021) is a state-of-the-art unsupervised skill-based RL
algorithm, which embeds sub-trajectories into the latent skill space. Similar
to DT, we modify SPiRL to handle either a video demonstration or language
instruction in a single framework.

###  4.3 One-shot Imitation Performance

Tables 1 and 2 compare the one-shot imitation performance in task success
rates achieved by our framework (U-OnIS, S-OnIS) and other baselines (DT,
SPiRL, BC-Z), given two distinct scenarios in terms of modality with a single
input in either a video demonstration or language instruction. As shown in
Figure 3(c), all the learned modules in the framework freeze when being
deployed in a target environment, and they are immediately evaluated against
various demonstrations and time-varying dynamics conditions.

Specifically, we measure the one-shot imitation performance in task success
rates for multi-stage tasks with $K$ sequential objectives ($K=1,2,4$) upon
several dynamics change levels (stationary, low, medium, and high). Regarding
baselines, DT and SPiRL variants for leveraging video demonstrations are
denoted as V-DT and V-SPiRL, and those for leveraging language instructions
are denoted as L-DT and L-SPiRL. Due to its low performance, V-DT is not
included in the comparison.

Table 1: One-shot imitation performance for video demonstration: for multi-
stage Meta-world tasks, the performance in task success rates by our U-OnIS,
S-OnIS and other baselines is measured against various test conditions on $K$
sequential objectives in a task ($K=1,2,4$) and dynamics change levels
(stationary, low, medium, and high).

$K$ | Non-st. | V-SPiRL | BC-Z | U-OnIS | S-OnIS
---|---|---|---|---|---
1 | Stationary | 61.48% | 75.00% | 95.80% | 100.0%
Low | 48.23% | 48.33% | 86.28% | 94.55%
Medium | 45.22% | 44.29% | 84.09% | 94.34%
High | 39.09% | 37.63% | 81.10% | 90.49%
2 | Stationary | 61.75% | 71.08% | 67.50% | 100.0%
Low | 41.66% | 42.85% | 75.75% | 85.66%
Medium | 43.39% | 47.61% | 73.07% | 84.29%
High | 30.33% | 22.09% | 66.00% | 81.62%
4 | Stationary | 47.44% | 21.01% | 64.38% | 91.67%
Low | 27.31% | 14.38% | 50.03% | 74.95%
Medium | 20.82% | 12.79% | 49.22% | 74.89%
High | 15.54% | 11.13% | 49.82% | 70.19%

Table 2:  One-shot imitation performance for language instruction

$K$ | Non-st. | L-DT | L-SPiRL | BC-Z | U-OnIS | S-OnIS
---|---|---|---|---|---|---
1 | Stationary | 92.66% | 100.0% | 100.0% | 87.70% | 100.0%
Low | 62.49% | 54.29% | 96.61% | 90.40% | 95.00%
Medium | 53.84% | 50.71% | 75.32% | 81.29% | 89.04%
High | 20.19% | 47.75% | 68.91% | 82.19% | 88.50%
2 | Stationary | 58.33% | 60.08% | 76.25% | 74.50% | 95.00%
Low | 47.12% | 29.35% | 45.27% | 60.44% | 83.17%
Medium | 33.17% | 34.34% | 35.00% | 69.24% | 78.50%
High | 30.28% | 20.51% | 25.52% | 66.67% | 77.27%
4 | Stationary | 37.73% | 43.16% | 40.30% | 67.71% | 87.50%
Low | 22.11% | 22.83% | 20.23% | 59.64% | 76.60%
Medium | 20.43% | 15.50% | 20.60% | 56.68% | 71.39%
High | 8.90% | 16.39% | 11.13% | 54.82% | 66.92%

Overall performance. As shown in Table 1, our U-OnIS and S-OnIS yield higher
performance consistently than the baselines for all video demonstration cases.
Specifically, S-OnIS outperforms the most competitive baseline SPiRL by a
significant margin $38.25\%\sim 54.64\%$ on average. The performance gap
between S-OnIS and the baselines increases for larger $K=4$ where a task
consists of 4 subtasks, compared to $K=1,2$. This specifies that our approach
is competitive in learning on a relatively long-horizon demonstration that
performs multi-stage tasks.

Consistency in multi-modality. Similar to the video demonstration scenario, in
Table 2, U-OnIS and S-OnIS show competitiveness for the language instruction
scenario. When Tables 1 and 2 are compared, we observe that U-OnIS and S-OnIS
maintain high performance consistently for the two distinct scenarios, showing
$5.69\%$ and $3.17\%$ performance differences (between the language and video
scenarios) on average between them. Unlike OnIS, although BC-Z explores multi-
modality, it achieves inconsistent performance between the two scenarios,
showing a $16.01\%$ performance drop for the video demonstration scenario from
the language instruction scenario. This inconsistent performance pattern
between different modalities can be also found in (Jang et al., 2022). While
our approach exploits task compositionality with shareable subtasks and
semantic skills, BC-Z tends to transform a whole video input into an
individual task; thus, BC-Z can degrade significantly for new video
demonstrations.

Generalization to non-stationarity. As shown in Tables 1 and 2, our U-OnIS and
S-OnIS show more robust performance than the baselines for all non-stationary
settings (including low, medium, and high). Specifically, in Table 2, S-OnIS
achieves a relatively small drop of $9.49\%\sim 18.13\%$ on average in non-
stationary environments compared to the stationary. The most competitive
baseline, BC-Z shows a large drop of $19.72\%\sim 57.02\%$ on average in the
same condition. In OnIS, the separate embedding spaces for semantic skills and
dynamics are used to address time-varying dynamics. Without that decoupled
embedding strategy, BC-Z can experience low-performance skills, which are
entangled with certain dynamics in training environments.

U-OnIS and S-OnIS. As presented above, while both U-OnIS and S-OnIS outperform
the baselines, S-OnIS shows better performance than U-OnIS with an average
margin of $16.63\%$. In our framework with efficient prompt learning, the
ability of semantic skill learning can be improved by a small annotated
dataset (i.e., 0.76% of total trajectories of 3120 episodes) with subtask-
level instructions.

Figure 4: Semantic embedding correspondence for a demonstration and one-shot
imitation run

Semantic embedding correspondence. Given a video demonstration, in Figure 4,
we visualize its embeddings (as black-colored circles, Demo.) and the
embeddings by its respective episodic run by U-OnIS (as blue-colored circles,
OnIS) in a multi-modal semantic space. As observed, two streams of the
embeddings construct a similar path, indicating their temporal correspondence
in semantic skills in an episode. We also visualize semantic skill embeddings
(as rhombuses of different colors) of sub-trajectories in our dataset, which
are matched with the 4 subtasks in the demonstration. The path (OnIS) turns
out to follow these semantic skill embeddings. This correspondence specifies
that a policy achieved by one-shot imitation can perform skills, which are
required to complete the subtasks in the demonstration, in the correct
sequence.

###  4.4 Ablation Study

Temporal contrast. Table 3 evaluates our positive pair sampling strategy in
the dynamics encoder, where Fixed$\pm N$ and Random$\pm N$ denote sampling on
$N$-fixed interval and and random sampling within $N$ interval, respectively,
in a trajectory. Specifically, our framework employs Random$\pm$T, where T is
the trajectory length. With Random$\pm$T, the framework shows $3.75\%\sim
18.25\%$ higher performance than others. This is because pairs sampled in a
whole episode contain a wider range of state variations with respect to
dynamics.

Table 3: Effect by temporal contrast

Dynamics | Fixed$\pm$1 | Fixed$\pm$10 | Random$\pm$10 | Random$\pm$T
---|---|---|---|---
Unseen Stationary | 49.50% | 54.25% | 55.00 % | 63.25%
Non-Stationary | 39.75% | 44.50% | 63.25% | 62.50%

Annotation sample size. Figure 5 evaluates S-OnIS with respect to different
sample sizes (4, 8, 16, 24) used for training the prompt (in (5)). As
expected, the extremely low data regime of 4 samples rarely solves even a
single task ($K=1$), but the default S-OnIS implementation with 24 samples
achieves robust performance for all $K=1,2,4$. This shows the efficiency of
prompt-based contrastive learning used for the semantic skill sequence
encoder.

Figure 5: Effect by annotated sample size: the x-axis denotes the number of
annotated samples used for training the semantic skill sequence encoder in
S-OnIS, and the y-axis denotes the one-shot imitation performance by S-OnIS.

###  4.5 Use Cases

Noisy video demonstration. To evaluate the robustness of our framework on
noisy video demonstrations, we construct different scenarios: (1) Color
jitter, (2) Gaussian noise, (3) Cutout. Table 4 shows the performance of
4-stage tasks for the scenarios. The baselines are highly vulnerable to noisy
video demonstrations, but OnIS is robust to them. In the training phase, OnIS
employs video and text retrieval tasks with a vision-language model, tending
to extract semantic information from visual observations. This renders OnIS
robust to noisy demonstrations, especially when noise does not alter
underlying semantics.

Table 4: Noisy video demonstrations

Noise type | V-SPiRL | BC-Z | U-OnIS | S-OnIS
---|---|---|---|---
w/o noise | 47.44% | 21.01% | 64.38% | 91.67%
Color jitter | 0.00% | 2.78% | 35.19% | 68.05%
Gaussian noise | 0.00% | 0.00% | 25.42% | 20.83%
Cutout | 0.00% | 0.00% | 27.78% | 8.33%

Real-world video demonstration. Here, we test our framework with real-world
video demonstrations, where we use a few real-world video clips, e.g., slide
puck and push button at the top of Figure 6. The graph associated with the
real-world video clips specifies the inference for pretrained semantics skills
over timesteps, yielded by our semantic skill sequence encoder. The inference
results state that the encoder is able to extract correct semantic skills from
the input clips to perform respective robotic manipulation tasks, exploring
the domain-invariant semantic knowledge tuned on the pretrained CLIP visual
encoder. At the bottom of the figure, we present the semantic skill execution
over timesteps in the target environment, where each semantic skill for a
certain timestep is determined by the semantic skill decoder and is then
translated into an optimized action sequence through our skill transfer
module.

Figure 6: Real-world demonstrations

Instruction variation. We evaluate our framework with language instruction
variants, which are not used for training the semantic skill sequence encoder.
For this, we use several instruction variations with verb replacement,
modifier extension, and verbose instructions. As shown in Table 5, both U-OnIS
and S-OnIS achieve robust performance against unseen language instructions
with variations. This is because the sequence encoder does not directly map an
instruction in the embedding space through the CLIP language encoder, but
translates it to its respective semantic skill by using its similarity to
learned semantic skills.

Table 5: Language instruction variations

Variations | U-OnIS | S-OnIS
---|---|---
Seen instruction | 70.01% | 93.50%
Verb replacement | 70.01% | 93.50%
Modifier extension | 70.01% | 93.50%
Verbose instruction | 40.74% | 54.25%

##  5 Related Works

One-shot imitation learning. To adapt to a new task with a single
demonstration, one-shot imitation learning has been investigated by extracting
task information from state action sequences (Duan et al., 2017; Yu et al.,
2018a) or from video demonstrations (Yu et al., 2018b; Mandi et al., 2022; Xu
et al., 2021). For example, (Yu et al., 2018b) exploited human video
demonstrations to learn robotic manipulation tasks, focusing on the
generalization ability to unseen tasks in video demonstrations, and (Mandi et
al., 2022) adopted contrastive learning schemes for cross-domain one-shot
imitation. We share the same direction to one-shot imitation with these prior
works, but we address specific one-shot imitation issues for a non-stationary
environment, aiming to achieve one-shot imitation for unseen long-horizon
tasks and zero-shot adaptation across various dynamics.

Skill-based RL. To tackle long-horizon multi-stage tasks, skill-based RL
techniques have been explored along with hierarchical learning frameworks.
(Eysenbach et al., 2019) explored unsupervised skill learning in a
hierarchical framework without requiring explicit reward signals from the
environment. Recently, in (Pertsch et al., 2021; Nam et al., 2022), with the
availability of diverse expert trajectory datasets in multi-task environments,
skill embedding techniques in offline settings were further investigated to
accelerate online policy learning for long-horizon tasks. We also use skill
embedding techniques to tackle multi-stage long-horizon tasks, but adapt them
with a vision-language pretrained model to exploit semantics in skills and
demonstrations. Our OnIS is the first framework using semantic skill
embeddings in the CLIP space to enable one-shot imitation and zero-shot
adaptation.

Multi-modality for imitation learning. Recently, several imitation learning
approaches using large scale pretrained models have been introduced. (Guhur et
al., 2022) used a pretrained language model to tackle instruction-following
tasks in a vision-based environment, and (Liu et al., 2022; Shridhar et al.,
2022) used a pretrained multi-modal model to encode both instruction and
visual observations for better scaling and generalization ability. (Jang et
al., 2022) presented a vision-language joint learning method for one-shot
imitation learning which can utilize both video demonstrations and language
instructions. In the same vein as (Jang et al., 2022), we also use a multi-
modal model for one-shot imitation, but we focus on zero-shot adaptation for
different dynamics in a non-statinary environment.

##  6 Conclusion

We presented a skill-based imitation learning framework, OnIS that enables not
only one-shot imitation learning from a single demonstration but also zero-
shot adaptation of a learned policy to different dynamics in a non-stationary
environment. To this end, we leverage the compositionality of long-horizon
tasks, by which a task can be decomposed into a sequence of skills represented
in the vision-language aligned semantic space, as well as explore meta-
learning techniques to enable skill transfer upon various dynamics conditions.
Through experiments, we demonstrated the superiority of OnIS with several one-
shot imitation use cases in terms of one-shot imitation performance,
generalization ability to unseen tasks and dynamics, and extensibility to
different modalities in demonstrations. Our future direction includes
addressing the issue of learning complex tasks via imitation in the
environment with limited data, where sufficient expert data does not exist to
learn diverse skills.

##  7 Acknowledgement

We would like to thank anonymous reviewers for their valuable comments and
suggestions. This work was supported by Institute of Information &
communications Technology Planning & Evaluation (IITP) grant funded by the
Korea government (MSIT) (No. 2022-0-01045, 2022-0-00043, 2019-0-00421,
2020-0-01821) and by the National Research Foundation of Korea (NRF) grant
funded by the MSIT (No. RS-2023-00213118).

## References

  * Bradbury et al. (2018) Bradbury, J., Frostig, R., Hawkins, P., Johnson, M. J., Leary, C., Maclaurin, D., Necula, G., Paszke, A., VanderPlas, J., Wanderman-Milne, S., and Zhang, Q.  JAX: composable transformations of Python+NumPy programs, 2018\.  URL http://github.com/google/jax.
  * Brohan et al. (2022) Brohan, A., Chebotar, Y., Finn, C., Hausman, K., Herzog, A., Ho, D., Ibarz, J., Irpan, A., Jang, E., Julian, R., et al.  Do as i can, not as i say: Grounding language in robotic affordances.  In _Proceedings of the 6th Conference on Robot Learning (CoRL)_. PMLR, 2022.
  * Buslaev et al. (2020) Buslaev, A., Iglovikov, V. I., Khvedchenya, E., Parinov, A., Druzhinin, M., and Kalinin, A. A.  Albumentations: fast and flexible image augmentations.  _Information_ , 11(2):125, 2020.
  * Chen et al. (2021) Chen, L., Lu, K., Rajeswaran, A., Lee, K., Grover, A., Laskin, M., Abbeel, P., Srinivas, A., and Mordatch, I.  Decision transformer: Reinforcement learning via sequence modeling.  In _Proceedings of the 34th Conference on Neural Information Processing Systems (NeurIPS)_ , volume 34, pp. 15084–15097, 2021.
  * Duan et al. (2017) Duan, Y., Andrychowicz, M., Stadie, B., Jonathan Ho, O., Schneider, J., Sutskever, I., Abbeel, P., and Zaremba, W.  One-shot imitation learning.  In _Proceedings of the 30th Conference on Neural Information Processing Systems (NeurIPS)_ , volume 30, 2017.
  * Eysenbach et al. (2019) Eysenbach, B., Gupta, A., Ibarz, J., and Levine, S.  Diversity is all you need: Learning skills without a reward function.  In _Proceedings of the 7th International Conference on Learning Representations (ICLR)_ , 2019.
  * Finn et al. (2017) Finn, C., Yu, T., Zhang, T., Abbeel, P., and Levine, S.  One-shot visual imitation learning via meta-learning.  In _Proceedings of the 1st Annual Conference on Robot Learning (CoRL)_ , volume 78, pp. 357–368. PMLR, 2017.
  * Garg et al. (2022) Garg, D., Vaidyanath, S., Kim, K., Song, J., and Ermon, S.  LISA: Learning interpretable skill abstractions from language.  In _Proceedings of the 35th Conference on Neural Information Processing Systems (NeurIPS)_ , 2022.
  * Guhur et al. (2022) Guhur, P.-L., Chen, S., Pinel, R. G., Tapaswi, M., Laptev, I., and Schmid, C.  Instruction-driven history-aware policies for robotic manipulations.  In _Proceedings of the 6th Conference on Robot Learning (CoRL)_. PMLR, 2022.
  * Gupta et al. (2019) Gupta, A., Kumar, V., Lynch, C., Levine, S., and Hausman, K.  Relay policy learning: Solving long horizon tasks via imitation and reinforcement learning.  In _Proceedings of the 2nd Conference on Robot Learning (CoRL)_. PMLR, 2019.
  * Jang et al. (2022) Jang, E., Irpan, A., Khansari, M., Kappler, D., Ebert, F., Lynch, C., Levine, S., and Finn, C.  BC-Z: Zero-shot task generalization with robotic imitation learning.  In _Proceedings of the 5th Conference on Robot Learning (CoRL)_ , pp. 991–1002. PMLR, 2022.
  * Kumar et al. (2021) Kumar, A., Fu, Z., Pathak, D., and Malik, J.  RMA: rapid motor adaptation for legged robots.  In _Robotics: Science and Systems XVII, Virtual Event, July 12-16, 2021_ , 2021.
  * Lee et al. (2022) Lee, K.-H., Nachum, O., Yang, M., Lee, L., Freeman, D., Xu, W., Guadarrama, S., Fischer, I., Jang, E., Michalewski, H., and Mordatch, I.  Multi-game decision transformers, 2022.
  * Li et al. (2022) Li, J., Lu, T., Cao, X., Cai, Y., and Wang, S.  Meta-imitation learning by watching video demonstrations.  In _Proceedings of the 10th International Conference on Learning Representations (ICLR)_ , 2022.
  * Liu et al. (2022) Liu, H., Lee, L., Lee, K., and Abbeel, P.  Instruction-following agents with jointly pre-trained vision-language models.  _arXiv preprint arXiv:2210.13431_ , 2022.
  * Mandi et al. (2022) Mandi, Z., Liu, F., Lee, K., and Abbeel, P.  Towards more generalizable one-shot visual imitation learning.  In _Proceedings of the 39th International Conference on Robotics and Automation (ICRA)_ , pp. 2434–2444. IEEE, 2022.
  * Nam et al. (2022) Nam, T., Sun, S.-H., Pertsch, K., Hwang, S. J., and Lim, J. J.  Skill-based meta-reinforcement learning.  In _Proceedings of the 10th International Conference on Learning Representations (ICLR)_ , 2022.
  * Pertsch et al. (2021) Pertsch, K., Lee, Y., and Lim, J.  Accelerating reinforcement learning with learned skill priors.  In _Proceedings of the 4th Conference on robot learning (CoRL)_ , pp. 188–204. PMLR, 2021.
  * Pertsch et al. (2022) Pertsch, K., Desai, R., Kumar, V., Meier, F., Lim, J. J., Batra, D., and Rai, A.  Cross-domain transfer via semantic skill imitation.  In _Proceedings of the 6th Conference on Robot Learning (CoRL)_. PMLR, 2022.
  * Qi et al. (2022) Qi, H., Kumar, A., Calandra, R., Ma, Y., and Malik, J.  In-hand object rotation via rapid motor adaptation.  In _Conference on Robot Learning, CoRL 2022, 14-18 December 2022, Auckland, New Zealand_ , pp. 1722–1732, 2022.
  * Radford et al. (2021) Radford, A., Kim, J. W., Hallacy, C., Ramesh, A., Goh, G., Agarwal, S., Sastry, G., Askell, A., Mishkin, P., Clark, J., et al.  Learning transferable visual models from natural language supervision.  In _Proceedings of the 38th International Conference on Machine Learning (ICML)_ , pp. 8748–8763. PMLR, 2021.
  * Shridhar et al. (2022) Shridhar, M., Manuelli, L., and Fox, D.  CLIPort: What and where pathways for robotic manipulation.  In _Proceedings of the 5th Conference on Robot Learning (CoRL)_ , pp. 894–906. PMLR, 2022.
  * van den Oord et al. (2017) van den Oord, A., Vinyals, O., and kavukcuoglu, k.  Neural discrete representation learning.  In _Proceedings of the 33rd Conference on Neural Information Processing Systems (NeurIPS)_ , volume 30, 2017.
  * Woo et al. (2022) Woo, H., Yoo, G., and Yoo, M.  Structure learning-based task decomposition for reinforcement learning in non-stationary environments.  In _Thirty-Sixth AAAI Conference on Artificial Intelligence, AAAI ,February 22 - March 1, 2022_ , pp. 8657–8665, 2022.
  * Xie et al. (2021) Xie, A., Harrison, J., and Finn, C.  Deep reinforcement learning amidst continual structured non-stationarity.  In _Proceedings of the 38th International Conference on Machine Learning, ICML 2021, 18-24 July 2021, Virtual Event_ , pp. 11393–11403, 2021\.
  * Xu et al. (2021) Xu, C., Amato, C., and Wong, L.  Hierarchical robot navigation in novel environments using rough 2-D maps.  In _Proceedings of the 3rd Conference on Robot Learning (CoRL)_ , volume 155, pp. 1971–1991. PMLR, 2021.
  * Xu et al. (2018) Xu, D., Nair, S., Zhu, Y., Gao, J., Garg, A., Fei-Fei, L., and Savarese, S.  Neural task programming: Learning to generalize across hierarchical tasks.  In _Proceedings of the 35th International Conference on Robotics and Automation (ICRA)_ , pp. 3795–3802. IEEE, 2018.
  * Yu et al. (2018a) Yu, T., Abbeel, P., Levine, S., and Finn, C.  One-shot hierarchical imitation learning of compound visuomotor tasks.  _arXiv preprint arXiv:1810.11043_ , 2018a.
  * Yu et al. (2018b) Yu, T., Finn, C., Dasari, S., Xie, A., Zhang, T., Abbeel, P., and Levine, S.  One-shot imitation from observing humans via domain-adaptive meta-learning.  In _Proceedings of the Robotics: Science and Systems_ , 2018b.
  * Yu et al. (2020) Yu, T., Quillen, D., He, Z., Julian, R., Hausman, K., Finn, C., and Levine, S.  Meta-world: A benchmark and evaluation for multi-task and meta reinforcement learning.  In _Proceedings of the 3rd Conference on robot learning (CoRL)_ , pp. 1094–1100. PMLR, 2020.
  * Zheng et al. (2022) Zheng, K., Chen, X., Jenkins, O., and Wang, X. E.  VLMbench: A compositional benchmark for vision-and-language manipulation.  In _Proceedings of the 36th Conference on Neural Information Processing Systems (NeurIPS) Datasets and Benchmarks Track_ , 2022.
  * Zhou et al. (2022) Zhou, K., Yang, J., Loy, C. C., and Liu, Z.  Learning to prompt for vision-language models.  _International Journal of Computer Vision_ , 130(9):2337–2348, 2022.
  * Zintgraf et al. (2020) Zintgraf, L. M., Shiarlis, K., Igl, M., Schulze, S., Gal, Y., Hofmann, K., and Whiteson, S.  Varibad: A very good method for bayes-adaptive deep RL via meta-learning.  In _8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020_ , 2020.

##  Appendix A Environments and Dataset

###  A.1 Multi-Stage Meta-world

Multi-stage Meta-world is a robotic arm manipulation environment specifically
designed for the execution of multi-stage tasks. For each task, multiple
objects are positioned on a desk, where each object corresponds to existing
tasks in Meta-world (Yu et al., 2020). The successful completion of a task is
determined by the execution of $N$ subtasks in a predetermined sequence,
similar to the Franka Kitchen environment (Gupta et al., 2019).

The state space is 140-dimensional, consisting of 4-dimensional position
information of the robot arm, and 136-dimensional position (e.g., the location
of the drawer) and status (e.g., whether the drawer is opened or closed) of
objects. The action space is 4-dimensional, consisting of a 3-dimensional
directional vector applied to the end-effector of the robot arm, and an
1-dimensional torque vector applied to the gripper. An example including the
robot arm and several objects is illustrated in Figure 7. The reward function
in the multi-stage Meta-world environment follows the same success metrics of
Meta-world in a subtask level.

Figure 7: Multi-stage Meta-world

Non-stationarity. To emulate a non-stationary environment, we adopt a time-
varying kinematic parameter of the robot arm. Specifically, the environment
involves such a kinematic parameter $w_{t}$ that varies for each timestep $t$,
i.e.,

$\begin{split}&w_{t}=m+0.25b\times sin(\rho_{t})\\\
&\rho_{t}=w_{t-1}+0.75\pi\times z\end{split}$ (19)

where $w_{0}=l$, and $z$ is a standard normal distribution $\mathcal{N}(0,1)$.
Note that $b$, $l$ and $m$ are hyperparameters to represent various
configurations of environment dynamics. Accordingly, given action $a_{t}$, the
environment yields the next state which is sampled from the transition
probability $s_{t+1}^{\prime}\sim\mathcal{P}(\cdot|s_{t},a_{t}+w_{t})$.

###  A.2 Expert Data Collection

For expert data collection, we devise rule-based policies and execute them in
a stationary environment. Regarding environment hidden dynamics, we use 13
different stationary environment configurations with dynamics parameters $b=0$
and $m\in\\{-0.3,-0.25,-0.2,-0.15,-0.1,-0.05,0,0.05,0.1,0.15,\\\
0.2,0.25,0.3\\}$ in (19). Based on the initial location of objects and subtask
execution orders, we have 24 different tasks. We use 16 tasks for training and
use 8 tasks for evaluation, where for each task, 15 episodes are generated.
Accordingly, we build a dataset of 15 episodes for each of the 16 tasks in 13
different environment settings with distinct dynamics parameters, containing
3,120 trajectories (by $16\times 15\times 13$).

##  Appendix B Experiment Details

###  B.1 OnIS Implementation

Table 6: Hyperparameters for S-OnIS

Hyperparameter | Value
---|---
Batch size | 1024
Prompt $\theta_{p}$ | 16 vector in 512-dim
Frames $H$ of video clip | 30
Contrastive $\alpha$ (for (5)) | 1
Dynamics encoder $\psi_{enc}$ | 3 FC with LSTM
Activation function for $\psi_{enc}$ | Leaky ReLU
Inverse dynamics decoder $\psi_{dec}$ | 5 FC with 128 units
Activation function for $\psi_{dec}$ | Leaky ReLU
The number of codebooks $M$ for $\psi_{enc}$ | 20
Dimension of codebook | 10
Learning rate for $\psi_{enc}$ | 5.55e-4
Length of sub-trajectory $H_{0}$ | 3
Semantic skill decoder $\Phi_{dec}$ | 4 FC with 128 units
Activation function for $\Phi_{dec}$ | Leaky ReLU
Learning rate for $\Phi_{dec}$ | 1e-4
Learning rate for $\psi_{dec}$ | 1e-4
Skill transfer module $\pi_{tr}$ | 5 FC with 128 units
Activation function for $\pi_{tr}$ | Leaky ReLU
Learning rate for $\pi_{tr}$ | 1e-4

S-OnIS. S-OnIS involves 4-phases of training: (1) semantic skill sequence
encoder $\Phi_{enc}$, (2) semantic skill decoder $\Phi_{dec}$, (3) dynamics
encoder $\psi_{enc}$, and (4) skill transfer module $\pi_{tr}$.

During the training process of the semantic skill sequence encoder, the video
demonstration being passed through the encoder $\Phi_{enc}$ creates a physical
time bottleneck in the end-to-end training procedure of S-OnIS. To address
this issue, we assign the inferred skill $z_{t}=\Phi_{enc}(v_{t:t+H})$ to each
tuple $(s_{t},v_{t},l_{t},a_{t})$ in the dataset after training $\Phi_{enc}$.

During the training process of the semantic skill decoder $\Phi_{dec}$, we
have observed that the semantic skills within our dataset, consisting of long-
horizon episodes, change infrequently (i.e., a maximum of 4 times). As a
result, the semantic skill decoder should predominantly predict 0 for most
timesteps, leading to significantly lower performance in binary
classification. To address this issue, we assign the label skill done to each
timestep in the dataset. This label indicates the end of a skill for a
specific timestep as well as the preceding and subsequent 5-step intervals of
the skill. Then, we have

$\textit{skill
done}_{t}=\begin{cases}1&\text{if}\,\,\,\Phi_{enc}(v_{i:i+H})\neq\Phi_{enc}(v_{i+1:i+1+H})\\\
0&o.w\end{cases}$ (20)

where $i\in[t-5,t+5]$. The semantic skill decoder $\Phi_{dec}$ is trained to
minimize binary cross entropy

$\mathcal{L}_{tp}(\tau)=\text{BCE}(\textit{skill
done}_{t},\Phi_{dec}(s_{t_{0}},s_{t},z_{t}))$ (21)

in (12).

In the training procedure for the dynamics encoder $\psi_{enc}$, to enable
backpropagation through non-differentiable quantization operator $\mathbf{q}$,
we use the straight-through gradient estimator as explained in (van den Oord
et al., 2017).

The skill transfer module $\pi_{tr}$ is implemented in a simple MLP which is
trained via behavior cloning. The hyperparameter settings for S-OnIS are
summarized in Table 6.

Table 7: Hyperparameters for U-OnIS

Hyperparameter | Value
---|---
Batch size | 1024
Encoder $\theta_{v}$ | Causal transformer
Encoder $\theta_{l}$ | Causal transformer
The number of layers for $\theta_{v}$, $\theta_{l}$ | 1
The number of heads for $\theta_{v}$, $\theta_{l}$ | 4
The number of codebooks $K$ for $\theta_{l}$ | 20
Learning rate for $\theta_{v}$, $\theta_{l}$ | 5e-5
Dimension of codebook for $\theta_{l}$ | 10
Length of sub-trajectory $H$ | 20
Contrastive $\alpha$ (for (9)) | 1
Dynamics encoder $\psi_{enc}$ | 3 FC with LSTM
Activation function for $\psi_{enc}$ | Leaky ReLU
Inverse dynamics decoder $\psi_{dec}$ | 5 FC with 128 units
Activation function for $\psi_{dec}$ | Leaky ReLU
The number of codebook for $\psi_{enc}$ | 20
Dimension of codebook for $\psi_{enc}$ | 10
Learning rate for $\psi_{enc}$ | 5.55e-4
Learning rate for $\psi_{dec}$ | 1e-4
Length of sub-trajectory $H_{0}$ | 3
Semantic skill decoder $\Phi_{dec}$ | 4 FC with 128 units
Activation function for $\Phi_{dec}$ | Leaky ReLU
Learning rate for $\Phi_{dec}$ | 1e-4
Skill transfer module $\pi_{tr}$ | 5 FC with 128 units
Activation function for $\pi_{tr}$ | Leaky ReLU
Learning rate for $\pi_{tr}$ | 1e-4

U-OnIS. U-OnIS involves 4-phases of training: (1) semantic skill sequence
encoder $\Phi_{enc}$, (2) semantic skill decoder $\Phi_{dec}$, (3) dynamics
encoder $\psi_{enc}$, and (4) skill transfer module $\pi_{tr}$.

In the training procedure for the semantic skill sequence encoder
$\Phi_{enc}$, we use two additional encoders $\theta_{v}$ and $\theta_{l}$
which are similar to the skill predictor used in (Garg et al., 2022), as
explained in Section 3.2. Encoder $\theta_{l}$ takes $(\Phi_{l}(l),s_{t-H:t})$
as input, and maps to a quantized semantic skill $z_{t}$ as

$\theta_{l}=\mathbf{q}\circ\theta^{c}_{l}:(\Phi_{l}(l),s_{t-H:t})\mapsto
z_{t}$ (22)

where $\theta^{c}_{l}$ maps the input $(\Phi_{l}(l),s_{t-H:t})$ to a
continuous latent vector. As explained in (15), the quantization operator
$\mathbf{q}$ maps the output of $\theta^{c}_{l}$ to one of the vectors in the
skill codebook $\\{c_{1},\cdots,c_{K}\\}$. That is, quantized semantic skill
embedding $z_{t}$ at each timestep $t$ is defined as

$z_{t}=\underset{j\in\\{1,\cdots,K\\}}{\textnormal{argmin}}\\{\|\theta^{c}_{l}(\Phi_{l}(l),s_{t-H:t})-c_{j}\|\\}.$
(23)

Note that the contrastive loss presented in (9) is only responsible for
updating the encoder $\theta_{v}$, and there is no gradient from (9) to update
the encoder $\theta_{l}$. The whole procedure of U-OnIS is presented in
Algorithm 2 and the the hyperparameter settings for U-OnIS are summarized in
Table 7.

Algorithm 2 One-shot imitation of U-OnIS

1: Semantic skill sequence encoder $\Phi_{enc}$

2: Semantic skill decoder $\Phi_{dec}$

3: Skill transfer module $\pi_{tr}$, Dynamics encoder $\psi_{enc}$

4: Environment $env$, Demonstration $d$

5: done $\leftarrow$ False$,\ t\leftarrow 0,\ \tau_{t}=queue(size=H_{0})$

6: $s_{0}\leftarrow env.\text{reset}()$

7: while not done do

8: $z_{t}\leftarrow\Phi_{enc}(d,s_{t-H:t})$ using (7)

9: $h_{t}\leftarrow\psi_{enc}(\tau_{t-1})$ using (14)

10: $a_{t}\leftarrow\pi_{tr}(s_{t},h_{t},z_{t})$ using (13)

11: $s_{t+1},\text{done}\leftarrow env.\text{step}(a_{t})$

12: if $\Phi_{dec}(s_{start},s_{t},z_{t})=1$ then

13: $s_{start}\leftarrow s_{t}$

14: end if

15: $\tau_{t}.enqueue((s_{t},a_{t})),\ t\leftarrow t+1$

16: end while

###  B.2 Baseline Implementation

BC-Z. BC-Z is an imitation learning framework which exploits multi-modal
demonstrations. In BC-Z, a language regression network is contrastively
trained on an expert demonstration by maximizing the similarity between the
embeddings of video demonstration and the embeddings of language instruction.
Instead of training a video encoder, we use the embedding vectors extracted
from the vision-language pertained model CLIP. For a fair comparison, a
history of state-actions is combined with either the vision or the language
embedding as input to the policy. The policy network is trained via behavior
cloning, conditioned on video or language embeddings. The hyperparameter
settings for BC-Z are summarized in Table 8.

Table 8: Hyperparameters for BC-Z

Hyperparameter | Value
---|---
Batch size | 1024
Language regression network | 5 FC with 512 units
Length of expert sub-trajectory | 10
Learning rate | 1e-4

Decision Transformer (DT). DT is a Transformer-based behavior cloning model.
Specifically, each token in language instruction is embedded into the latent
vector using the pertained CLIP, and then it is concatenated with the state
vector for the input of DT. The hyperparameter settings for DT are summarized
in Table 9.

Table 9: Hyperparameters for DT Hyperparameter | Value
---|---
Batch size | 64
The number of layers | 3
Embedding dimension | 128
The number of heads | 4
Dropout | 0.1
Learning rate | 1e-4

SPiRL. SPiRL is an unsupervised skill learning framework, which embeds an
expert sub-trajectory into a latent vector skill and decodes it into a
sequence of actions via a primitive policy. Specifically, the CLIP embedding
vector of video is concatenated with a state and used as input of V-SPiRL, and
the CLIP embedding of language instruction is concatenated with a state and
used as input to L-SPiRL. The hyperparameter settings for SPiRL are summarized
in Table 10.

Table 10: Hyperparameters for SPiRL

Hyperparameter | Value
---|---
Batch size | 1024
Skill encoder | 1 FC with LSTM
Primitive policy | 7 FC with 128 units
Length of expert sub-trajectory | 10
Skill dimension | 20
Learning rate | 1e-4

Figure 8: T-SNE embedding maps: (i) a video demonstration performing 4
subtasks (Demo.), (ii) a video generated by OnIS, and (iii) a video generated
by OnIS without dynamics embedding (OnIS wo/h). Expert videos for each subtask
(e.g., Slide puck, Push button, Close drawer, Open door) are represented in
the same space.

###  B.3 Evaluation Details

The distribution of dynamics at test time assumes to be completely out of the
distribution of training time dynamics. Specifically, the parameter $m$ is
sparsely sampled (from 13 distinct values) according to uniform distribution
of $[-0.3,0.3]$ and $\omega_{t}$ is fixed in each trajectory at training time
with $b=0$ in (19). At test time, different hyperparameter values for $b,m$
are sampled at each episode from $[1,2]$ and $[-0.45,0.45]$, respectively.
Accordingly, $\omega_{t}$ is changed by (19) at each timestep.

###  B.4 Use Case Experiment Details

Noisy video demonstrations. We add noise to video demonstrations through three
different methods (Color Jitter, Gaussian Noise, and Cutout) as described in
(Buslaev et al., 2020).

  * •

Color jitter. In each video frame, adjustments are made to the brightness,
contrast, saturation, and hue, in a way that random noise vectors are sampled
from uniform distributions of $[-0.1,0.1]$, $[-0.1,0.1]$, $[-0.1,0.1]$, and
$[-0.025,0.025]$ and the noise vectors are then added to video demonstrations.

  * •

Gaussian noise. In each video frame, we sample a noise vector from a Gaussian
distribution with a mean of 1 and a variance ranging from 1 to 10. The sampled
noise is then added to video demonstrations.

  * •

Cutout. 16 square regions of size $4\times 4$ are randomly cropped in each
video frame.

Real-world video demonstration. We record various real-world video
demonstrations in $700\times 700$ resolution, which involve interactions with
objects that represent different stages in multi-stage Meta-world. Then, we
evaluate our framework with different noisy levels. We also conduct the
performance evaluation after fine-tuning the semantic skill sequence encoder
(in the row of w/ Fine-tune in Table 11), presuming the cases where fine-
tuning is feasible if a real-world video demonstration is given with subtask-
level instructions.

Table 11: Performance on real-world demonstrations

| wo/ noise | noisy | very noisy
---|---|---|---
Success Rate | 91.67% | 73.75% | 54.38%
Skills Matching Ratio | 100.0% | 72.08% | 52.92%
Success Rate (w/ Fine-tune) | 91.67% | 83.13% | 67.50%
Skills Matching Ratio (w/ Fine-tune) | 100.0% | 89.16% | 76.46%

As observed in Table 11, the skills matching ratio is positively correlated
with the success rate, specifying that if correct semantic skills are
predicted from the semantic skill sequence encoder, no performance degradation
is expected for noisy real-world settings. When a single real-world
demonstration is given with a subtask instruction, it is feasible to improve
the success rate by tuning the language prompt of the semantic skill sequence
encoder. As we test several real-world demonstrations with robotic
manipulation in multi-stage Meta-world scenarios, we observe that the
performance in the success rate can be rather dependent on the capability of
underlying pretrained vision-language models.

Instruction variation. Table 12 describes the examples of the instruction
variations in Table 5.

Table 12: Examples of instruction variations

Verb replacement
---
push button $\rightarrow$ press the button
open door $\rightarrow$ pull the door
close drawer $\rightarrow$ shut the drawer
slide puck $\rightarrow$ move the puck
Modifier extension
push button $\rightarrow$ push the red button
open door $\rightarrow$ open the black door
close drawer $\rightarrow$ close the green drawer
slide puck $\rightarrow$ slide the black puck
Verbose instruction
push button $\rightarrow$ press the red button on yellow box
open door $\rightarrow$ open the closed black door on the desk
close drawer $\rightarrow$ close the green drawer on the desk
slide puck $\rightarrow$ move the black plate to the white goal on the left
side

##  Appendix C Additional Experiments

###  C.1 Additional Experiments for Unseen Environments

Table 13 shows the performance of OnIS when dynamics parameter $m$ in the seen
task increases from $-0.6$ to $0$. As shown, S-OnIS yields the best
performance for all $K$ (the number of subtasks), implying that S-OnIS is
robust to unseen dynamics. This result corresponds to Section 4.3.
Furthermore, Figure 9 shows that OnIS achieves the most robust performance as
the dynamics level $|m|$ in (19) increases (e.g., to unseen dynamics) over
previously trained dynamics. This specifies that the dynamics embedding
procedure of OnIS is able to disentangle dynamics information from expert
trajectories.

Table 13: Performance in the success rate for multi-stage Meta-world with
language instructions

$K$ | L-DT | L-SPiRL | BC-Z | U-OnIS | S-OnIS
---|---|---|---|---|---
1 | 57.81% | 29.68% | 78.13% | 70.32% | 94.79%
2 | 33.59% | 37.50% | 26.57% | 66.57% | 77.04%
4 | 21.09% | 17.96% | 18.60% | 47.44% | 60.58%

Figure 9: Relative performance of unseen dynamics to previously trained
dynamics

###  C.2 Long-Horizon Multi-Stage Tasks

Table 14: Experiments for Long-Horizon Tasks |  Number of objects ($K$)
---|---
| 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10
OnIS(Language Instruction) | 99.33% | 76.0% | 74.66% | 69.33% | 65.33% | 58.0% | 50.66% | 44.0% | 36.0% | 26.0%
OnIS(Video Demonstration) | 99.33% | 72.0% | 64.0% | 60.33% | 42.0% | 29.0% | 10.03% | 4.0% | 1.33% | 0.0%

To investigate whether our approach can tackle long-horizon multi-stage tasks
with a large number of skills, we evaluate OnIS with manipulation tasks that
involve up to 10 objects. As shown in Table 14, OnIS is able to solve the
tasks with up to 10 objects ($K=1,\dots,10$). This clarifies that the semantic
skill sequence encoder scales well to predict semantic skills even when a task
involves multiple subtasks of up to 10.

###  C.3 Visualization of Semantic Reconstruction

Figure 8 represents the embedding maps of the expert demonstration (Demo.),
U-OnIS (OnIS), and OnIS trained without dynamics embedding (OnIS wo/h) for 6
different tasks. Similar to the experiment in Section 4 in the main
manuscript, the embeddings of episodic run by U-OnIS and demonstrations show
highly similar paths, meaning that their semantic skills are aligned. On the
other hand, the OnIS without dynamics embedding has low semantic skill
correspondence with demonstrations.

###  C.4 Additional Baseline: Multigame-DT

Table 15:  One-shot imitation performance for language instruction

$K$ | Non-st. | Multi-DT | L-DT | BC-Z | U-OnIS | S-OnIS
---|---|---|---|---|---|---
1 | Stationary | 100.0% | 92.66% | 100.0% | 87.70% | 100.0%
Low | 85.58% | 62.49% | 96.61% | 90.40% | 95.00%
Medium | 73.07% | 53.84% | 75.32% | 81.29% | 89.04%
High | 66.34% | 20.19% | 68.91% | 82.19% | 88.50%
2 | Stationary | 68.75% | 58.33% | 76.25% | 74.50% | 95.00%
Low | 53.84% | 47.12% | 45.27% | 60.44% | 83.17%
Medium | 52.40% | 33.17% | 35.00% | 69.24% | 78.50%
High | 42.79% | 30.28% | 25.52% | 66.67% | 77.27%
4 | Stationary | 31.25% | 37.73% | 40.30% | 67.71% | 87.50%
Low | 20.55% | 22.11% | 20.23% | 59.64% | 76.60%
Medium | 18.51% | 20.43% | 20.60% | 56.68% | 71.39%
High | 14.90% | 8.90% | 11.13% | 54.82% | 66.92%

We perform the experiment with the latest version of DT, a multigame-DT (Lee
et al., 2022). Table 15 shows the performance in multi-stage Meta-world
achieved by OnIS and baselines including Multi-DT. We observe that Multi-DT
achieves better than L-DT with an average gain of 11.73%, but shows worse
performance than S-OnIS with an average of loss of 31.74% For fair
comparisons, we use video and language embeddings from the pretrained CLIP
model as an additional input to Multi-DT, as well as other baselines. Since
Multi-DT uses rewards in offline datasets (unlike OnIS without any rewards),
we consider that Multi-DT has some advantage over the original DT in terms of
policy learning. Otherwise, Multi-DT becomes similar to the original DT. In
our problem setting of one-shot imitation for multi-stage tasks, Multi-DT is
yet limited in exploring task compositionality.

###  C.5 Ablation Studies

Table 16: Ablation: Contrastive Learning on Dynamics Encoder $K$ | Non-st. | w/ Contra, w/ VQ | w/ Contra, wo/ VQ | wo/ Contra, w/ VQ | wo/ Contra, wo/ VQ | wo/ Recon
---|---|---|---|---|---|---
1 | Stationary | 100.0% | 100.0% | 100.0% | 100.0% | 100.0%
Low | 95.01% | 93.84% | 92.30% | 89.42% | 89.10%
Medium | 94.55% | 92.11% | 83.65% | 77.24% | 84.40%
High | 88.50% | 89.42% | 72.62% | 61.21% | 81.76%
2 | Stationary | 100.0% | 100.0% | 68.75% | 79.16% | 95.10%
Low | 85.82% | 83.75% | 66.98% | 76.60% | 75.04%
Medium | 84.42% | 82.37% | 61.37% | 58.52% | 74.45%
High | 81.14% | 69.07% | 49.34% | 48.96% | 75.04%
4 | Stationary | 87.50% | 88.09% | 46.88% | 58.34% | 85.50%
Low | 80.11% | 74.45% | 45.97% | 53.50% | 71.96%
Medium | 75.24% | 75.58% | 43.50% | 44.31% | 67.04%
High | 66.92% | 63.19% | 34.39% | 34.13% | 65.32%

Contrastive learning on dynamics encoder. To investigate the effect of
contrastive learning-based dynamics embedding in our proposed OnIS framework,
we evaluate and compare several different implementations for dynamics
embedding as part of ablation. Table 16 illustrates those including the cases
of whether the contrastive loss is added to the typical dynamics embedding
loss or not (w/ Contra and wo/ Contra, respectively, in the column names), the
cases of whether the dynamics embedding vector $\psi_{enc}(\tau)$ is quantized
or not (w/ VQ and wo/ VQ, respectively, in the column names), and the case of
reconstruction loss (18) is not added (wo/ Recon, in the column name) to the
dynamics embedding loss.

Here, the typical dynamics embedding loss is based on reconstruction, as in
the formula in (18), similar to several related works such as (Zintgraf et
al., 2020; Woo et al., 2022; Xie et al., 2021):

$\mathcal{L}_{DE}=\mathbb{E}_{t-H_{0}<i\leq
t}\left[\|a_{i}-\psi_{dec}(s_{i},s_{i+1},\psi_{enc}(\tau_{t})\|^{2}\right]$
(24)

where $\psi_{enc}$ and $\psi_{dec}$ denote the dynamics encoder and decoder,
respectively.

As observed in Table 16, the contrastively trained dynamics encoder (w/
Contra) improves the performance by 19.8% in average success rates for multi-
stage Meta-world tasks, compared to the dynamics encoder trained without
contrastive learning (wo/ Contra).

Moreover, we also observe that quantizing the output of the dynamics encoder
(w/ VQ) improves the performance by 2.2% in average success rates, compared to
the dynamics embedding without quantization (wo/ VQ). Additionally, the
reconstruction loss in the training of dynamics encoder (w/ Contra, w/VQ)
improves the performance by 6.21%, compared to the dynamics encoder trained
without reconstruction loss (wo/ Recon).

These ablation results with such performance gains specify that contrastive
learning enables the dynamics encoder to disentangle task-invariant dynamics-
relevant features from expert demonstrations. The task information - what to
do, such as open door and push button - is extracted by the semantic skill
sequence encoder in prior, and then it is feasible to enforce different sub-
trajectories sharing the same dynamics have a close proximity in the dynamics
embedding space, regardless of which skills the sub-trajectories involve. As
two sub-trajectories in a trajectory are used as a positive pair in our
sampling strategy, the contrastive learning loss in (17) renders the
aforementioned desired property of the dynamics embedding space.