# Skill-based Model-based Reinforcement Learning

Lucy Xiaoyang Shi1 Joseph J. Lim2,3 Youngwoon Lee1
1University of Southern California 2KAIST 3NAVER AI Lab AI Advisor at NAVER AI
Lab

###### Abstract

Model-based reinforcement learning (RL) is a sample-efficient way of learning
complex behaviors by leveraging a learned single-step dynamics model to plan
actions in imagination. However, planning every action for long-horizon tasks
is not practical, akin to a human planning out every muscle movement. Instead,
humans efficiently plan with high-level skills to solve complex tasks. From
this intuition, we propose a Skill-based Model-based RL framework (SkiMo) that
enables planning in the skill space using a skill dynamics model, which
directly predicts the skill outcomes, rather than predicting all small details
in the intermediate states, step by step. For accurate and efficient long-term
planning, we jointly learn the skill dynamics model and a skill repertoire
from prior experience. We then harness the learned skill dynamics model to
accurately simulate and plan over long horizons in the skill space, which
enables efficient downstream learning of long-horizon, sparse reward tasks.
Experimental results in navigation and manipulation domains show that SkiMo
extends the temporal horizon of model-based approaches and improves the sample
efficiency for both model-based RL and skill-based RL. Code and videos are
available at https://clvrai.com/skimo.

> Keywords: Model-Based Reinforcement Learning, Skill Dynamics Model

##  1 Introduction

A key trait of human intelligence is the ability to plan abstractly for
solving complex tasks [1]. For instance, we perform cooking by imagining
outcomes of high-level skills like washing and cutting vegetables, instead of
planning every muscle movement involved [2]. This ability to plan with
temporally-extended skills helps to scale our internal model to long-horizon
tasks by reducing the search space of behaviors. To apply this insight to
artificial intelligence agents, we propose a novel skill-based and model-based
reinforcement learning (RL) method, which learns a model and a policy in a
high-level skill space, enabling accurate long-term prediction and efficient
long-term planning.

Typically, model-based RL involves learning a flat single-step dynamics model,
which predicts the next state from the current state and action. This model
can then be used to simulate “imaginary” trajectories, which significantly
improves sample efficiency over their model-free alternatives [3, 4]. However,
such model-based RL methods have shown only limited success in long-horizon
tasks due to inaccurate long-term prediction [5] and computationally expensive
search [6, 7, 8].

Skill-based RL enables agents to solve long-horizon tasks by acting with
multi-action subroutines (skills) [9, 10, 11, 12, 13, 14] instead of primitive
actions. This temporal abstraction of actions enables systematic long-range
exploration and allows RL agents to plan farther into the future, while
requiring a shorter horizon for policy optimization, which makes long-horizon
downstream tasks more tractable. Yet, on complex long-horizon tasks, skill-
based RL still requires a few million to billion environment interactions to
learn [13], which is impractical for real-world applications.

To combine the best of both model-based RL and skill-based RL, we propose
Skill-based Model-based RL (SkiMo), which enables effective planning in the
skill space using a skill dynamics model. Given a state and a skill to
execute, the skill dynamics model directly predicts the resultant state after
skill execution, without needing to model every intermediate step and low-
level action (Figure 1), whereas the flat dynamics model predicts the
immediate next state after one action execution. Thus, planning with skill
dynamics requires fewer predictions than flat dynamics, resulting in more
reliable long-term future predictions and plans.

Concretely, we first jointly learn the skill dynamics model and a skill
repertoire from large offline datasets collected across diverse tasks [15, 12,
16]. This joint training shapes the skill embedding space for easy skill
dynamics prediction and skill execution. Then, to solve a complex downstream
task, we train a high-level task policy that acts in the learned skill space.
For more efficient policy learning and better planning, we leverage the skill
dynamics model to simulate skill trajectories.

The main contribution of this work is to propose Skill-based Model-based RL
(SkiMo), a novel sample-efficient model-based hierarchical RL algorithm that
leverages task-agnostic data to extract not only a reusable skill set but also
a skill dynamics model. The skill dynamics model enables efficient and
accurate long-term planning for sample-efficient RL. Our experiments show that
our method outperforms the state-of-the-art skill-based and model-based RL
algorithms on long-horizon navigation and robotic manipulation tasks with
sparse rewards.

Figure 1:  Intelligent agents can use their internal models to imagine
potential futures for planning. Instead of planning out every primitive action
(black arrows in a), they aggregate action sequences into skills (red and blue
arrows in b). Further, they can leap directly to the predicted outcomes of
executing skills in sequence (red and blue arrows in c), which leads to better
long-term prediction and planning compared to predicting step-by-step
(blurriness of images represents the level of error accumulation in
prediction).

##  2 Related Work

Model-based RL leverages a (learned) dynamics model of the environment to plan
a sequence of actions that leads to the desired behavior. The dynamics model
predicts the future state of the environment after taking a specific action,
which enables simulating candidate behaviors in imagination instead of in the
physical environment. Then, these imaginary rollouts can be used for planning
[8, 4] through, e.g., CEM [17] and MPPI [18], as well as for policy
optimization [19, 3, 20, 4] to improve sample efficiency. Yet, due to the
accumulation of prediction error at each step and the increasing search space,
finding an optimal, long-horizon plan is inaccurate and computationally
expensive [6, 7, 8].

To facilitate learning of long-horizon behaviors, skill-based RL lets the
agent act over temporally-extended skills (i.e. options [21] or motion
primitives [22]), which can be represented as sub-policies or a coordinated
sequence of low-level actions. Temporal abstraction effectively reduces the
task horizon for the agent and enables directed exploration [23]. The reusable
skills can be manually defined [22, 24, 10, 11, 14], extracted from large
offline datasets [15, 25, 26, 27, 28], discovered online in an unsupervised
manner [29, 30], or acquired in the form of goal-reaching policies [31, 32,
33, 34, 20]. However, skill-based RL is still impractical for real-world
applications, requiring a few million to billion environment interactions
[13]. In this paper, we use model-based RL to guide the planning of skills to
improve the sample efficiency of skill-based approaches.

There have been attempts to plan over skills in model-based RL [29, 35, 5, 36,
37]. However, most of these approaches [29, 5, 36] still utilize the
conventional flat (single-step) dynamics model, which struggles at handling
long-horizon planning due to error accumulation. Wu et al. [35] proposes to
learn a temporally-extended dynamics model; however, it conditions on low-
level actions rather than skills and is only used for low-level planning. A
concurrent work, Shah et al. [37], is most similar to our work in that it
learns a skill dynamics model, but with a limited set of discrete, manually-
defined skills. To fully unleash the potential of temporally abstracted
skills, we extract the skill space from data and devise a skill-level dynamics
model to provide accurate long-term prediction, which is essential for solving
long-horizon tasks. To the best of our knowledge, SkiMo is the first work that
jointly learns skills and a skill dynamics model from data for model-based RL.

##  3 Method

To enable accurate long-term prediction and efficient long-horizon planning
for RL, we introduce SkiMo, a novel skill-based and model-based RL algorithm
that shares synergistic benefits from both frameworks. A key change to prior
model-based approaches is the use of a skill dynamics model that directly
predicts the outcome of a chosen skill, which enables efficient and accurate
long-term planning. As illustrated in Figure 2, our approach consists of two
phases: (1) learning the skill dynamics model and skills from an offline
dataset (Section 3.3) and (2) downstream task learning with the skill dynamics
model (Section 3.4).

Figure 2:  Our approach, SkiMo, combines model-based RL and skill-based RL for
sample efficient learning of long-horizon tasks. SkiMo consists of two phases:
(1) learn a skill dynamics model and a skill repertoire from offline task-
agnostic data, and (2) learn a high-level policy for the downstream task by
leveraging the learned model and skills. We omit the encoded latent state
$\mathbf{h}$ in the figure and directly write the observation $\mathbf{s}$ for
clarity, but most modules take the latent state $\mathbf{h}$ as input.

###  3.1 Preliminaries

##### RL

We formulate a problem as a Markov decision process [38], which is defined by
a tuple $(\mathcal{S},\mathcal{A},R,P,\rho_{0},\gamma)$ of the state space
$\mathcal{S}$, action space $\mathcal{A}$, reward $R(\mathbf{s},\mathbf{a})$,
transition probability $P(\mathbf{s}^{\prime}\lvert\mathbf{s},\mathbf{a})$,
initial state distribution $\rho_{0}$, and discounting factor $\gamma$. A
policy $\pi(\mathbf{a}|\mathbf{s})$ maps from a state $\mathbf{s}$ to an
action $\mathbf{a}$. RL aims to find the optimal policy that maximizes the
expected discounted return,
$\mathbb{E}_{\mathbf{s}_{0}\sim\rho_{0},(\mathbf{s}_{0},\mathbf{a}_{0},\dots,\mathbf{s}_{T_{i}})\sim\pi}\left[\sum_{t=0}^{T_{i}-1}\gamma^{t}R(\mathbf{s}_{t},\mathbf{a}_{t})\right]$,
where $T_{i}$ is the variable episode length.

##### Unlabeled Offline Data

We assume access to a reward-free task-agnostic dataset [15, 12], which is a
set of $N$ state-action trajectories,
$\mathcal{D}=\\{\tau_{1},\dots,\tau_{N}\\}$. Since it is task-agnostic, this
data can be collected from human teleoperation, unsupervised exploration, or
training data for other tasks. We do not assume this dataset contains
solutions for the downstream task; therefore, tackling the downstream task
requires re-composition of skills learned from diverse trajectories.

##### Skill-based RL

We define skills as a sequence of actions
$(\mathbf{a}_{0},\dots,\mathbf{a}_{H-1})$ with a fixed horizon111It is worth
noting that our method is compatible with variable-length skills [26, 39, 27]
and goal-conditioned skills [20] with minimal change; however, for simplicity,
we adopt fixed-length skills of $H=10$ in this paper. $H$ and parameterize
skills as a skill latent $\mathbf{z}$ and skill policy,
$\pi^{L}(\mathbf{a}|\mathbf{s},\mathbf{z})$, that maps a skill latent and
state to the corresponding action sequence. The skill latent and skill policy
can be trained using variational auto-encoder (VAE [40]), where a skill
encoder $q(\mathbf{z}|(\mathbf{s},\mathbf{a})_{0:H-1})$ embeds a sequence of
transitions into a skill latent $\mathbf{z}$, and the skill policy decodes it
back to the original action sequence. Following SPiRL [12], we also learn a
skill prior $p(\mathbf{z}|\mathbf{s})$, which is the skill distribution in the
offline data, to guide the downstream task policy to explore promising skills
over the large skill space.

###  3.2 SkiMo Model Components

SkiMo consists of three major model components: the skill policy
($\pi_{\theta}^{L}$), skill dynamics model ($D_{\psi}$), and task policy
($\pi_{\phi}$), along with auxiliary components for representation learning
and value estimation. A state encoder $E_{\psi}$ first encodes an observation
$\mathbf{s}$ into the latent state $\mathbf{h}$. Then, given a skill
$\mathbf{z}$, the skill dynamics $D_{\psi}$ predicts the skill effect in the
latent space. The task policy $\pi_{\phi}$, reward function $R_{\phi}$, and
value function $Q_{\phi}$ predict a skill, reward, and value on the (imagined)
latent state, respectively. The following is a summary of the notations of our
model components:

$\begin{split}\begin{array}[]{ll}\text{State
encoder:}&\mathbf{h}_{t}=E_{\psi}(\mathbf{s}_{t})\\\ \text{Observation
decoder:}&\hat{\mathbf{s}}_{t}=O_{\theta}(\mathbf{h}_{t})\\\ \text{Skill
prior:}&\hat{\mathbf{z}}_{t}\sim p_{\theta}(\mathbf{s}_{t})\\\ \text{Skill
encoder:}&\mathbf{z}_{t}\sim q_{\theta}((\mathbf{s},\mathbf{a})_{t:t+H-1})\\\
\text{Skill
policy:}&\hat{\mathbf{a}}_{t}=\pi^{L}_{\theta}(\mathbf{s}_{t},\mathbf{z}_{t})\\\
\end{array}\end{split}\;\begin{split}\begin{array}[]{ll}\text{Skill
dynamics:}&\hat{\mathbf{h}}_{t+H}=D_{\psi}(\mathbf{h}_{t},\mathbf{z}_{t})\\\
\text{Task policy:}&\hat{\mathbf{z}}_{t}\sim\pi_{\phi}(\mathbf{h}_{t})\\\
\text{Reward:}&\hat{r}_{t}=R_{\phi}(\mathbf{h}_{t},\mathbf{z}_{t})\\\
\text{Value:}&\hat{v}_{t}=Q_{\phi}(\mathbf{h}_{t},\mathbf{z}_{t})\\\
\end{array}\end{split}$ (1)

For convenience, we label the trainable parameters $\psi$, $\theta$, $\phi$ of
each component according to which phase they are trained on:

  1. 1.

Learned from offline data and finetuned in downstream RL
$(\psi=\\{\psi_{E},\psi_{D}\\})$: The state encoder ($E_{\psi}$) and the skill
dynamics model ($D_{\psi}$) are first trained on the offline task-agnostic
data and then finetuned in downstream RL to account for unseen states and
transitions.

  2. 2.

Learned only from offline data
$(\theta=\\{\theta_{O},\theta_{q},\theta_{p},\theta_{\pi^{L}}\\})$: The
observation decoder ($O_{\theta}$), skill encoder ($q_{\theta}$), skill prior
($p_{\theta}$), and skill policy ($\pi_{\theta}^{L})$ are learned from the
offline data.

  3. 3.

Learned in downstream RL $(\phi=\\{\phi_{Q},\phi_{R},\phi_{\pi}\\})$: The
value ($Q_{\phi}$) and reward ($R_{\phi}$) functions, and the task policy
($\pi_{\phi}$) are trained for the downstream task using environment
interactions.

###  3.3 Pre-Training Skill Dynamics Model and Skills from Task-agnostic Data

SkiMo consists of pre-training and downstream RL phases. In pre-training,
SkiMo leverages offline data to extract (1) skills for temporal abstraction of
actions, (2) skill dynamics for skill-level planning on a latent state space,
and (3) a skill prior [12] to guide exploration. Specifically, we jointly
learn a skill policy and skill dynamics model, instead of learning them
separately [35, 5, 36], in a self-supervised manner. The key insight is that
this joint training could shape the latent skill space $\mathbb{Z}$ and state
embedding in that the skill dynamics model can easily predict the future.

In contrast to prior works that learn models completely online [3, 41, 4], we
leverage existing offline task-agnostic datasets to pre-train a skill dynamics
model and skill policy. This offers the benefit that the model and skills are
agnostic to specific tasks so that they may be used in multiple tasks.
Afterwards in the downstream RL phase, the agent continues to finetune the
skill dynamics model to accommodate task-specific trajectories.

To learn a low-dimensional skill latent space $\mathbb{Z}$ that encodes action
sequences, we train a conditional VAE [40, 42] on the offline dataset that
reconstructs the action sequence through a skill embedding given a state-
action sequence as in SPiRL [12, 16]. Specifically, given $H$ consecutive
states and actions $(\mathbf{s},\mathbf{a})_{0:H-1}$, a skill encoder
$q_{\theta}$ predicts a skill embedding $\mathbf{z}$ and a skill decoder
$\pi_{\theta}^{L}$ (i.e. the low-level skill policy) reconstructs the original
action sequence from $\mathbf{z}$:

$\mathcal{L}_{\text{VAE}}=\mathbb{E}_{(\mathbf{s},\mathbf{a})_{0:H-1}\sim\mathcal{D}}\bigg{[}\frac{\lambda_{\text{BC}}}{H}\sum_{i=0}^{H-1}\underbrace{(\pi^{L}_{\theta}(\mathbf{s}_{i},\mathbf{z})-\mathbf{a}_{i})^{2}}_{\text{Behavioral
cloning}}+\beta\cdot\underbrace{KL\big{(}q_{\theta}(\mathbf{z}|(\mathbf{s},\mathbf{a})_{0:H-1})~{}\|~{}p(\mathbf{z})}_{\text{Embedding
regularization}}\big{)}\bigg{]},$ (2)

where $\mathbf{z}$ is sampled from $q_{\theta}$ and
$\lambda_{\text{BC}},\beta$ are weighting factors for regularizing the skill
latent $\mathbf{z}$ distribution to a prior of a $\tanh$-transformed unit
Gaussian distribution, $Z\sim\tanh(\mathcal{N}(0,1))$.

To ensure the latent skill space is suited for long-term prediction, we
jointly train a skill dynamics model with the VAE above. The skill dynamics
model learns to predict $\mathbf{h}_{t+H}$, the latent state $H$-steps ahead
conditioned on a skill $\mathbf{z}$, for $N$ sequential skill transitions
using the latent state consistency loss [4]. To prevent a trivial solution and
encode rich information from observations, we additionally train an
observation decoder $O_{\theta}$ using the observation reconstruction loss.
Altogether, the skill dynamics $D_{\psi}$, state encoder $E_{\psi}$, and
observation decoder $O_{\theta}$ are trained on the following objective:

$\mathcal{L}_{\text{REC}}=\mathbb{E}_{(\mathbf{s},\mathbf{a})_{0:NH}\sim\mathcal{D}}\bigg{[}\sum_{i=0}^{N-1}\Big{[}\underbrace{\lambda_{\text{O}}\|\mathbf{s}_{iH}-O_{\theta}(E_{\psi}(\mathbf{s}_{iH}))\|^{2}_{2}}_{\text{Observation
reconstruction}}+\underbrace{\lambda_{\text{L}}\|D_{\psi}(\hat{\mathbf{h}}_{iH},\mathbf{z}_{iH})-E_{\psi^{-}}(\mathbf{s}_{(i+1)H})\|^{2}_{2}}_{\text{Latent
state consistency}}\Big{]}\bigg{]},$ (3)

where $\lambda_{\text{O}},\lambda_{\text{L}}$ are weighting factors and
$\hat{\mathbf{h}}_{0}=E_{\psi}(\mathbf{s}_{0})$ and
$\hat{\mathbf{h}}_{(i+1)H}=D_{\psi}(\hat{\mathbf{h}}_{iH},\mathbf{z}_{iH})$
such that gradients are back-propagated through time. For stable training, we
use a target network whose parameter $\psi^{-}$ is slowly soft-copied from
$\psi$.

Furthermore, to guide the exploration for downstream RL, we also extract a
skill prior [12] from offline data that predicts the skill distribution for
any state. The skill prior is trained by minimizing the KL divergence between
output distributions of the skill encoder $q_{\theta}$ and the skill prior
$p_{\theta}$:

$\mathcal{L}_{\text{SP}}=\mathbb{E}_{(\mathbf{s},\mathbf{a})_{0:H-1}\sim\mathcal{D}}\bigg{[}\lambda_{\text{SP}}\cdot
KL\Big{(}\textbf{sg}(q_{\theta}(\mathbf{z}|\mathbf{s}_{0:H-1},\mathbf{a}_{0:H-1}))~{}\|~{}p_{\theta}(\mathbf{z}|\mathbf{s}_{0})\Big{)}\bigg{]},$
(4)

where $\lambda_{\text{SP}}$ is a weighting factor and sg denotes the stop
gradient operator.

Combining the objectives above, we jointly train the policy, model, and prior,
which leads to a well-shaped skill latent space that is optimized for both
skill reconstruction and long-term prediction:

$\mathcal{L}=\mathcal{L}_{\text{VAE}}+\mathcal{L}_{\text{REC}}+\mathcal{L}_{\text{SP}}$
(5)

###  3.4 Downstream Task Learning with Learned Skill Dynamics Model

To accelerate downstream RL with the learned skill repertoire, SkiMo learns a
high-level task policy $\pi_{\phi}(\mathbf{z}_{t}|\mathbf{h}_{t})$ that
outputs a latent skill embedding $\mathbf{z}_{t}$, which is then translated
into a sequence of $H$ actions using the pre-trained skill policy
$\pi^{L}_{\theta}$ to act in the environment [12, 16].

To further improve the sample efficiency, we propose to use model-based RL in
the skill space by leveraging the skill dynamics model. The skill dynamics
model and task policy can generate imaginary rollouts in the skill space by
repeating (1) sampling a skill,
$\mathbf{z}_{t}\sim\pi_{\phi}(\mathbf{h}_{t})$, and (2) predicting $H$-step
future after executing the skill,
$\mathbf{h}_{t+H}=D_{\psi}(\mathbf{h}_{t},\mathbf{z}_{t})$. Our skill dynamics
model requires only $1/H$ dynamics predictions and action selections of the
flat model-based RL approaches [3, 4], resulting in more efficient and
accurate long-horizon imaginary rollouts (see Appendix, Figure 10).

Following TD-MPC [4], we leverage these imaginary rollouts both for planning
(Algorithm 2) and policy optimization (Equation (7)), significantly reducing
the number of necessary environment interactions. During rollout, we perform
Model Predictive Control (MPC), which re-plans every step using CEM and
executes the first skill of the skill plan (see Appendix, Section C for more
details).

To evaluate imaginary rollouts, we train a reward function
$R_{\phi}(\mathbf{h}_{t},\mathbf{z}_{t})$ that predicts the sum of $H$-step
rewards222For clarity, we use $r_{t}$ to abbreviate the sum of $H$-step
environment rewards $\sum_{i=0}^{H-1}R(\mathbf{s}_{t+i},\mathbf{a}_{t+i})$.,
$r_{t}$, and a Q-value function $Q_{\phi}(\mathbf{h}_{t},\mathbf{z}_{t})$. We
also finetune the skill dynamics model $D_{\psi}$ and state encoder $E_{\psi}$
on the downstream task to improve the model prediction:

$\displaystyle\mathcal{L}^{\prime}_{REC}=\mathbb{E}_{\mathbf{s}_{t},\mathbf{z}_{t},\mathbf{s}_{t+H},r_{t}\sim\mathcal{D}}\bigg{[}$
$\displaystyle\underbrace{\lambda_{\text{L}}\|D_{\psi}(\hat{\mathbf{h}}_{t},\mathbf{z}_{t})-E_{\psi^{-}}(\mathbf{s}_{t+H})\|^{2}_{2}}_{\text{Latent
state
consistency}}+\underbrace{\lambda_{\text{R}}\|r_{t}-R_{\phi}(\hat{\mathbf{h}}_{t},\mathbf{z}_{t})\|^{2}_{2}}_{\text{Reward
prediction}}$ (6) $\displaystyle+\underbrace{\lambda_{\text{V}}\|r_{t}+\gamma
Q_{\phi^{-}}(\hat{\mathbf{h}}_{t+H},\pi_{\phi}(\hat{\mathbf{h}}_{t+H}))-Q_{\phi}(\hat{\mathbf{h}}_{t},\mathbf{z}_{t})\|^{2}_{2}}_{\text{Value
prediction}}\bigg{]}.$

Finally, we train a high-level task policy $\pi_{\phi}$ to maximize the
estimated $Q$-value while regularizing it to the pre-trained skill prior
$p_{\theta}$ [12], which helps the policy output plausible skills:

$\mathcal{L}_{RL}=\mathbb{E}_{\mathbf{s}_{t}\sim\mathcal{D}}\big{[}-Q_{\phi}(\hat{\mathbf{h}}_{t},\pi_{\phi}(\textbf{sg}(\hat{\mathbf{h}}_{t})))+\alpha\cdot
KL\big{(}\pi_{\phi}(\mathbf{z}_{t}|\textbf{sg}(\hat{\mathbf{h}}_{t}))~{}\|~{}p_{\theta}(\mathbf{z}_{t}|\mathbf{s}_{t})\big{)}\big{]}.$
(7)

The models $(\psi,\phi)$ are trained to minimize Equation (6) and Equation (7)
using back-propagation through time over $N$ consecutive skill-level
transitions, similar to Equation (3).

##  4 Experiments

In this paper, we propose SkiMo, a model-based RL approach that can
efficiently and accurately plan long-horizon trajectories by leveraging skills
and skill dynamics model. In our experiments, we aim to answer the following
questions: (1) Can the skill dynamics model improve the efficiency of RL for
long-horizon tasks? and (2) Is the joint training of skills and the skill
dynamics model essential for efficient model-based RL?

We compare SkiMo with prior model-based RL and skill-based RL methods on four
long-horizon tasks with sparse rewards: maze navigation, two kitchen
manipulation, and tabletop manipulation tasks, as illustrated in Figure 3.
More experimental details can be found in Appendix, Section C.

###  4.1 Tasks

(a) Maze

(b) Kitchen

(c) Mis-aligned Kitchen

(d) CALVIN

Figure 3:  We evaluate our method on four long-horizon, sparse reward tasks.
(a) The green point mass navigates the maze to reach the goal (red). (b, c)
The robot arm in the kitchen must complete four tasks in the correct order,
Microwave - Kettle - Bottom Burner - Light and Microwave - Light - Slide
Cabinet - Hinge Cabinet. (d) The robot arm needs to complete four tasks in the
correct order, Open Drawer - Turn on Lightbulb - Move Slider Left - Turn on
LED.

Maze We use the maze navigation task from Pertsch et al. [16], where a point
mass agent is randomly initialized near the green region and needs to reach
the fixed goal region in red (Figure 3(a)). The agent observes its 2D position
and 2D velocity, and controls its $(x,y)$-velocity. The agent receives a
sparse reward of 100 only when it reaches the goal. The task-agnostic dataset
[16] consists of 3,046 trajectories between randomly sampled initial and goal
positions.

Kitchen We use the FrankaKitchen tasks and 603 trajectories from D4RL [43].
The 7-DoF Franka Emika Panda arm needs to perform four sequential sub-tasks,
Microwave - Kettle - Bottom Burner - Light. In Mis-aligned Kitchen, we also
test another task sequence, Microwave - Light - Slide Cabinet - Hinge Cabinet,
which has a low sub-task transition probability in the offline data
distribution [16]. The agent observes 11D robot state and 19D object state,
and uses 9D joint velocity control. The agent receives a reward of 1 for every
sub-task completion in order.

CALVIN We adapt CALVIN [44] to have the target task, Open Drawer - Turn on
Lightbulb - Move Slider Left - Turn on LED, and 21D robot and object states.
It uses the Franka Panda arm with 7D end-effector pose control. The offline
data is from play data [44] consisting of 1,239 trajectories. The agent
receives a reward of 1 for every sub-task completion in the correct order.

###  4.2 Baselines

  * •

Dreamer [3] and TD-MPC [4] learn a flat (single-step) dynamics and train a
policy using latent imagination to achieve a high sample efficiency.

  * •

DADS [29] discovers skills and learns a dynamics model through unsupervised
learning.

  * •

LSP [36] plans in the skill space, but using a single-step dynamics model from
Dreamer [3].

  * •

SPiRL [12] learns skills and a skill prior, and guides a high-level policy
using the learned prior.

  * •

SPiRL + Dreamer and SPiRL + TD-MPC pre-train the skills using SPiRL and learn
a policy and model in the skill space, instead of the low-level action space,
using Dreamer and TD-MPC, respectively. In contrast to SkiMo, these baselines
do not jointly train the model and skills.

(a) Maze

(b) Kitchen

(c) Mis-aligned Kitchen

(d) CALVIN

Figure 4:  Learning curves of our method and baselines. All averaged over 5
random seeds.

###  4.3 Results

Maze Maze navigation poses a hard exploration problem due to the sparsity of
the reward: the agent only receives reward after taking 1,000+ steps to reach
the goal. Figure 4(a) shows that only SkiMo is able to consistently reach the
goal, whereas baselines struggle to learn a policy or an accurate model due to
the challenges in sparse feedback and long-term planning.

We qualitatively analyze the behavior of each agent in Appendix, Figure 9.
Dreamer and TD-MPC have a small coverage of the maze since it is challenging
to coherently explore for 1,000+ steps to reach the goal from taking primitive
actions. Similarly, DADS and LSP could not learn meaningful skills and never
find the goal. SPiRL is able to explore a large fraction of the maze, but it
does not learn to consistently find the goal due to difficult policy
optimization in long-horizon tasks. SPiRL + Dreamer and SPiRL + TD-MPC fail to
learn an accurate model and often collide with walls.

Kitchen Figure 4(b) demonstrates that SkiMo reaches the same performance
(above 3 sub-tasks) with 5x less environment interactions than SPiRL. In
contrast, Dreamer, TD-MPC, DADS, and LSP rarely succeed on the first sub-task
due to the sparse reward. SPiRL + Dreamer and SPiRL + TD-MPC perform better
than flat model-based RL by leveraging skills, yet the independently trained
model and policy are not accurate enough to consistently achieve more than two
sub-tasks.

Mis-aligned Kitchen The mis-aligned target task makes the downstream learning
harder because the skill prior, which reflects offline data distribution,
offers less meaningful regularization to the policy. However, Figure 4(c)
shows that SkiMo still performs well. This demonstrates that the skill
dynamics model is able to adapt to the new distribution of behaviors, which
might significantly deviate from the distribution in the offline dataset.

CALVIN One of the major challenges in CALVIN is that the offline data is very
task-agnostic: any particular sub-task transition has probability lower than
0.1% on average, resulting in a large number of plausible skills from any
state. Figure 4(d) shows that SkiMo can learn faster than the model-free
baseline, SPiRL, which supports the benefit of using our skill dynamics model.
Meanwhile, Dreamer performs better in CALVIN than in Kitchen because objects
in CALVIN are more compactly located and easier to manipulate with the end-
effector control; thus, it becomes viable to accomplish initial sub-tasks
through random exploration. However, it falls short in composing coherent
action sequences to achieve a longer task sequence due to the lack of
temporally-extended reasoning.

In summary, we show the synergistic benefit of temporal abstraction in both
the policy and dynamics model. SkiMo is the only method that consistently
solves the long-horizon tasks. Our results also demonstrate the importance of
algorithmic design choices (e.g. skill-level planning, joint training of a
model and skills) as naive combinations (SPiRL + Dreamer, SPiRL + TD-MPC) fail
to learn.

###  4.4 Ablation Studies

Model-based vs. Model-free In Figure 5, SkiMo achieves better asymptotic
performance and higher sample efficiency across all tasks than SkiMo + SAC,
which uses model-free RL (SAC [45]) to train the high-level policy to select
skills. The comparison suggests that the task policy can make more informative
decisions by leveraging accurate long-term predictions of the skill dynamics
model.

Joint training of skills and skill dynamics model SkiMo w/o joint training
learns the latent skill space using only the VAE loss in Equation (2). Figure
5 shows that the joint training is crucial in more challenging scenarios,
where the agent needs to generate accurate long-term plans (for Maze) or the
skills are very diverse (in CALVIN).

CEM planning As shown in Figure 5, SkiMo learns significantly better and
faster in Kitchen, Mis-aligned Kitchen, and CALVIN than SkiMo w/o CEM,
indicating that CEM planning can effectively find a better plan. On the other
hand, in Maze, SkiMo w/o CEM learns twice as fast. We find that action noise
for exploration in CEM leads the agent to get stuck at walls and corners. We
believe that with a careful tuning of action noise, SkiMo can solve Maze much
more efficiently.

Skill dynamics model SkiMo w/ single-step dynamics replaces the skill dynamics
model with a conventional single-step dynamics model. Similar to LSP, it
learns and plans on skills, but additionally pre-trains the skills and the
model on offline datasets for fair comparison. As shown in Figure 5, the
single-step model struggles at handling long-horizon planning on all tasks,
akin to the baseline results on DADS and LSP in Figure 4. In contrast, the
skill dynamics model can make accurate long-horizon predictions for planning
due to significantly less compounding errors.

For further ablations and discussion on skill horizon and planning horizon,
see Appendix, Section A.

(a) Maze

(b) Kitchen

(c) Mis-aligned Kitchen

(d) CALVIN

Figure 5:  Learning curves of our method and ablated methods. All averaged
over 5 random seeds.

###  4.5 Long-horizon Prediction with Skill Dynamics Model

To assess the accuracy of long-term prediction of our proposed skill dynamics
over flat dynamics, we visualize imagined trajectories in Appendix, Figure 10,
where the ground truth initial state and a sequence of 500 actions (50 skills
for SkiMo) are given. Dreamer struggles to make accurate long-horizon
predictions due to error accumulation. In contrast, SkiMo is able to reproduce
the ground truth trajectory with little prediction error even when traversing
through hallways and doorways. This confirms that SkiMo allows temporal
abstraction in the dynamics model, thereby enabling temporally-extended
prediction and reducing step-by-step prediction error.

##  5 Conclusion

We propose SkiMo, an intuitive instantiation of saltatory model-based
hierarchical RL [2], which combines skill-based and model-based RL approaches.
Our experiments demonstrate that (1) a skill dynamics model reduces the long-
term future prediction error via temporal abstraction in the dynamics model;
(2) without needing to plan step-by-step, downstream RL over the skill space
allows for efficient and accurate temporally-extended reasoning, improving the
performance of prior model-based RL and skill-based RL; and (3) joint training
of the skill dynamics and skills further improves the sample efficiency by
learning skills conducive to predict their consequences. We believe that the
ability to learn and utilize a skill-level model holds the key to unlocking
the sample efficiency and widespread use of RL agents for long-horizon tasks,
and our method takes a step toward this direction.

##### Limitations and future work

While our method extracts fixed-length skills from offline data, the lengths
of semantic skills may vary based on the contexts and goals. Future work can
learn variable-length semantic skills to improve long-term prediction and
planning. Further, although we only experimented on state-based inputs, SkiMo
is a general framework that can be extended to RGB, depth, and tactile
observations. Thus, extending our approach to real robots with high-
dimensional observations would be an interesting future work.

#### Acknowledgments

This work was supported by Institute of Information & communications
Technology Planning & Evaluation (IITP) grants (No.2019-0-00075, Artificial
Intelligence Graduate School Program, KAIST; No.2022-0-00077, AI Technology
Development for Commonsense Extraction, Reasoning, and Inference from
Heterogeneous Data) and National Research Foundation of Korea (NRF) grant
(NRF-2021H1D3A2A03103683), funded by the Korea government (MSIT). This work
was also partly supported by the Annenberg Fellowship from USC. We would like
to thank Ayush Jain and Grace Zhang for help on writing, Karl Pertsch for
assistance in setting up SPiRL and CALVIN, Kevin Xie for providing code of
LSP, and all members of the USC CLVR lab for constructive feedback.

## References

  * Legg and Hutter [2007] S. Legg and M. Hutter.  Universal intelligence: A definition of machine intelligence.  _Minds and machines_ , 17(4):391–444, 2007.
  * Botvinick and Weinstein [2014] M. Botvinick and A. Weinstein.  Model-based hierarchical reinforcement learning and human action control.  _Philosophical Transactions of the Royal Society B: Biological Sciences_ , 369(1655), 2014.
  * Hafner et al. [2019] D. Hafner, T. Lillicrap, J. Ba, and M. Norouzi.  Dream to control: Learning behaviors by latent imagination.  In _International Conference on Learning Representations_ , 2019.
  * Hansen et al. [2022] N. Hansen, X. Wang, and H. Su.  Temporal difference learning for model predictive control.  In _International Conference on Machine Learning_ , 2022.
  * Lu et al. [2021] K. Lu, A. Grover, P. Abbeel, and I. Mordatch.  Reset-free lifelong learning with skill-space planning.  In _International Conference on Learning Representations_ , 2021.
  * Lowrey et al. [2019] K. Lowrey, A. Rajeswaran, S. Kakade, E. Todorov, and I. Mordatch.  Plan online, learn offline: Efficient learning and exploration via model-based control.  In _International Conference on Learning Representations_ , 2019.  URL https://openreview.net/forum?id=Byey7n05FQ.
  * Janner et al. [2019] M. Janner, J. Fu, M. Zhang, and S. Levine.  When to trust your model: Model-based policy optimization.  In _Neural Information Processing Systems_ , 2019.
  * Argenson and Dulac-Arnold [2021] A. Argenson and G. Dulac-Arnold.  Model-based offline planning.  In _International Conference on Learning Representations_ , 2021.  URL https://openreview.net/forum?id=OMNB1G5xzd4.
  * Sutton et al. [1999] R. S. Sutton, D. Precup, and S. Singh.  Between mdps and semi-mdps: A framework for temporal abstraction in reinforcement learning.  _Artificial intelligence_ , 112(1-2):181–211, 1999.
  * Lee et al. [2019] Y. Lee, S.-H. Sun, S. Somasundaram, E. S. Hu, and J. J. Lim.  Composing complex skills by learning transition policies.  In _International Conference on Learning Representations_ , 2019.  URL https://openreview.net/forum?id=rygrBhC5tQ.
  * Lee et al. [2020] Y. Lee, J. Yang, and J. J. Lim.  Learning to coordinate manipulation skills via skill behavior diversification.  In _International Conference on Learning Representations_ , 2020.
  * Pertsch et al. [2020] K. Pertsch, Y. Lee, and J. J. Lim.  Accelerating reinforcement learning with learned skill priors.  In _Conference on Robot Learning_ , 2020.
  * Lee et al. [2021] Y. Lee, J. J. Lim, A. Anandkumar, and Y. Zhu.  Adversarial skill chaining for long-horizon robot manipulation via terminal state regularization.  In _Conference on Robot Learning_ , 2021.
  * Dalal et al. [2021] M. Dalal, D. Pathak, and R. Salakhutdinov.  Accelerating robotic reinforcement learning via parameterized action primitives.  In _Neural Information Processing Systems_ , 2021.
  * Lynch et al. [2020] C. Lynch, M. Khansari, T. Xiao, V. Kumar, J. Tompson, S. Levine, and P. Sermanet.  Learning latent plans from play.  In _Conference on Robot Learning_ , pages 1113–1132. PMLR, 2020.
  * Pertsch et al. [2021] K. Pertsch, Y. Lee, Y. Wu, and J. J. Lim.  Demonstration-guided reinforcement learning with learned skills.  In _Conference on Robot Learning_ , 2021.
  * Rubinstein [1997] R. Y. Rubinstein.  Optimization of computer simulation models with rare events.  _European Journal of Operational Research_ , 99(1):89–112, 1997.
  * Williams et al. [2015] G. Williams, A. Aldrich, and E. Theodorou.  Model predictive path integral control using covariance variable importance sampling.  _arXiv preprint arXiv:1509.01149_ , 2015.
  * Ha and Schmidhuber [2018] D. Ha and J. Schmidhuber.  World models.  _arXiv preprint arXiv:1803.10122_ , 2018.
  * Mendonca et al. [2021] R. Mendonca, O. Rybkin, K. Daniilidis, D. Hafner, and D. Pathak.  Discovering and achieving goals via world models.  In _Neural Information Processing Systems_ , 2021.
  * Sutton et al. [1999] R. S. Sutton, D. Precup, and S. Singh.  Between mdps and semi-mdps: A framework for temporal abstraction in reinforcement learning.  _Artificial intelligence_ , 112(1-2):181–211, 1999.
  * Pastor et al. [2009] P. Pastor, H. Hoffmann, T. Asfour, and S. Schaal.  Learning and generalization of motor skills by learning from demonstration.  In _IEEE International Conference on Robotics and Automation_ , pages 763–768, 2009.
  * Nachum et al. [2019] O. Nachum, H. Tang, X. Lu, S. Gu, H. Lee, and S. Levine.  Why does hierarchy (sometimes) work so well in reinforcement learning?  _arXiv preprint arXiv:1909.10618_ , 2019.
  * Mülling et al. [2013] K. Mülling, J. Kober, O. Kroemer, and J. Peters.  Learning to select and generalize striking movements in robot table tennis.  _The International Journal of Robotics Research_ , 32(3):263–279, 2013.
  * Shiarlis et al. [2018] K. Shiarlis, M. Wulfmeier, S. Salter, S. Whiteson, and I. Posner.  Taco: Learning task decomposition via temporal alignment for control.  In _International Conference on Machine Learning_ , pages 4654–4663. PMLR, 2018.
  * Kipf et al. [2019] T. Kipf, Y. Li, H. Dai, V. Zambaldi, A. Sanchez-Gonzalez, E. Grefenstette, P. Kohli, and P. Battaglia.  Compile: Compositional imitation learning and execution.  In _International Conference on Machine Learning_ , 2019.
  * Shankar and Gupta [2020] T. Shankar and A. Gupta.  Learning robot skills with temporal variational inference.  In _International Conference on Machine Learning_ , 2020.
  * Lu et al. [2021] Y. Lu, Y. Shen, S. Zhou, A. Courville, J. B. Tenenbaum, and C. Gan.  Learning task decomposition with ordered memory policy network.  In _International Conference on Learning Representations_ , 2021.
  * Sharma et al. [2020] A. Sharma, S. Gu, S. Levine, V. Kumar, and K. Hausman.  Dynamics-aware unsupervised discovery of skills.  In _International Conference on Learning Representations_ , 2020.
  * Eysenbach et al. [2018] B. Eysenbach, A. Gupta, J. Ibarz, and S. Levine.  Diversity is all you need: Learning skills without a reward function.  In _International Conference on Learning Representations_ , 2018.
  * Nachum et al. [2018] O. Nachum, S. S. Gu, H. Lee, and S. Levine.  Data-efficient hierarchical reinforcement learning.  In _Advances in Neural Information Processing Systems_ , pages 3303–3313, 2018.
  * Gupta et al. [2019] A. Gupta, V. Kumar, C. Lynch, S. Levine, and K. Hausman.  Relay policy learning: Solving long-horizon tasks via imitation and reinforcement learning.  _Conference on Robot Learning_ , 2019.
  * Mandlekar et al. [2020a] A. Mandlekar, D. Xu, R. Martín-Martín, S. Savarese, and L. Fei-Fei.  Gti: Learning to generalize across long-horizon tasks from human demonstrations.  In _Robotics: Science and Systems_ , 2020a.
  * Mandlekar et al. [2020b] A. Mandlekar, F. Ramos, B. Boots, L. Fei-Fei, A. Garg, and D. Fox.  Iris: Implicit reinforcement without interaction at scale for learning control from offline robot manipulation data.  In _IEEE International Conference on Robotics and Automation_ , 2020b.
  * Wu et al. [2021] B. Wu, S. Nair, L. Fei-Fei, and C. Finn.  Example-driven model-based reinforcement learning for solving long-horizon visuomotor tasks.  In _Conference on Robot Learning_ , 2021.
  * Xie et al. [2020] K. Xie, H. Bharadhwaj, D. Hafner, A. Garg, and F. Shkurti.  Latent skill planning for exploration and transfer.  In _International Conference on Learning Representations_ , 2020.
  * Shah et al. [2022] D. Shah, A. T. Toshev, S. Levine, and brian ichter.  Value function spaces: Skill-centric state abstractions for long-horizon reasoning.  In _International Conference on Learning Representations_ , 2022.  URL https://openreview.net/forum?id=vgqS1vkkCbE.
  * Sutton [1984] R. S. Sutton.  _Temporal credit assignment in reinforcement learning_.  PhD thesis, University of Massachusetts Amherst, 1984.
  * Shankar et al. [2020] T. Shankar, S. Tulsiani, L. Pinto, and A. Gupta.  Discovering motor programs by recomposing demonstrations.  In _International Conference on Learning Representations_ , 2020.
  * Kingma and Welling [2014] D. P. Kingma and M. Welling.  Auto-encoding variational bayes.  In _International Conference on Learning Representations_ , 2014.
  * Sekar et al. [2020] R. Sekar, O. Rybkin, K. Daniilidis, P. Abbeel, D. Hafner, and D. Pathak.  Planning to explore via self-supervised world models.  In _International Conference on Machine Learning_ , 2020.
  * Higgins et al. [2017] I. Higgins, L. Matthey, A. Pal, C. Burgess, X. Glorot, M. Botvinick, S. Mohamed, and A. Lerchner.  beta-VAE: Learning basic visual concepts with a constrained variational framework.  In _International Conference on Learning Representations_ , 2017.  URL https://openreview.net/forum?id=Sy2fzU9gl.
  * Fu et al. [2020] J. Fu, A. Kumar, O. Nachum, G. Tucker, and S. Levine.  D4rl: Datasets for deep data-driven reinforcement learning.  _arXiv preprint arXiv:2004.07219_ , 2020.
  * Mees et al. [2022] O. Mees, L. Hermann, E. Rosete-Beas, and W. Burgard.  Calvin: A benchmark for language-conditioned policy learning for long-horizon robot manipulation tasks.  _IEEE Robotics and Automation Letters_ , 2022.
  * Haarnoja et al. [2018] T. Haarnoja, A. Zhou, P. Abbeel, and S. Levine.  Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor.  In _International Conference on Machine Learning_ , pages 1856–1865, 2018.
  * Paszke et al. [2017] A. Paszke, S. Gross, S. Chintala, G. Chanan, E. Yang, Z. DeVito, Z. Lin, A. Desmaison, L. Antiga, and A. Lerer.  Automatic differentiation in PyTorch.  In _NIPS Autodiff Workshop_ , 2017.
  * Tassa et al. [2018] Y. Tassa, Y. Doron, A. Muldal, T. Erez, Y. Li, D. de Las Casas, D. Budden, A. Abdolmaleki, J. Merel, A. Lefrancq, T. P. Lillicrap, and M. A. Riedmiller.  Deepmind control suite.  _arXiv preprint arXiv:1801.00690_ , 2018.
  * Schaul et al. [2016] T. Schaul, J. Quan, I. Antonoglou, and D. Silver.  Prioritized experience replay.  In _International Conference on Learning Representations_ , 2016.

##  Appendix A Further Ablations

We include additional ablations on the Maze and Kitchen tasks to further
investigate the influence of skill horizon $H$, planning horizon $N$, and
dynamics model fine-tuning, which is important for skill learning and
planning.

###  A.1 Skill Horizon

(a) Maze

(b) Kitchen

Figure 6: Ablation analysis on skill horizon $H$.

In both Maze and Kitchen, we find that a too short skill horizon ($H=1,5$) is
unable to yield sufficient temporal abstraction. A longer skill horizon
($H=15,20$) has little influence in Kitchen, but it makes the downstream
performance much worse in Maze. This is because with long-horizon skills, a
skill dynamics prediction becomes inaccurate and stochastic, and composing
multiple skills can be not as flexible as short-horizon skills. The inaccurate
skill dynamics makes long-term planning harder, which is already a major
challenge in maze navigation.

###  A.2 Planning Horizon

(a) Maze

(b) Kitchen

Figure 7: Ablation analysis on planning horizon $N$.

In Figure 7(b), we see that short planning horizon makes learning slower in
the beginning, because it does not effectively leverage the skill dynamics
model to plan further ahead. Conversely, if the planning horizon is too long,
the performance becomes worse due to the difficulty in modeling every step
accurately. Indeed, the planning horizon 20 corresponds to 200 low-level
steps, while the episode length in Kitchen is 280, demanding the agent to make
plan for nearly the entire episode. The performance is not sensitive to
intermediate planning horizons. On the other hand, the effect of the planning
horizon differs in Maze due to distinct environment characteristics. We find
that very long planning horizon (eg. 20) and very short planning horizon (eg.
1) perform similarly in Maze (Figure 7(a)). This could attribute to the former
creates useful long-horizon plans, while the latter avoids error accumulation
altogether. We leave further investigation on planning horizon to future work.

###  A.3 Fine-Tuning Model

(a) Maze

(b) Kitchen

Figure 8: Ablation analysis on fine-tuning the model.

We freeze the skill dynamics model together with the state encoder to gauge
the effect of fine-tuning after pre-training. Figure 8 shows that without
fine-tuning the model, the agent performs worse due to the discrepancy between
distributions of the offline data and the downstream task. We hypothesize that
fine-tuning is necessary when the agent needs to adapt to a different task and
state distribution after pre-training.

##  Appendix B Qualitative Analysis on Maze

###  B.1 Exploration and Exploitation

(a) SkiMo (Ours)

(b) Dreamer

(c) TD-MPC

(d) SPiRL + Dreamer

(e) SPiRL + TD-MPC

(f) SPiRL

(g) DADS

(h) LSP

(i) SkiMo w/o joint training

(j) SkiMo + SAC

(k) SkiMo w/o CEM

(l) SkiMo w/ single-step dynamics

Figure 9:  Exploration and exploitation behaviors of our method and baseline
approaches. We visualize trajectories in the replay buffer at 1.5M training
steps in blue: light blue for early trajectories and dark blue for recent
trajectories. Our method shows wide coverage of the maze at the early stage of
training and fast convergence to the solution.

To gauge the agent’s ability of exploration and exploitation, we visualize the
replay buffer for each method in Figure 9. In this visualization, we represent
early trajectories in the replay buffer with light blue dots and recent
trajectories with dark blue dots. In Figure 9(a), the replay buffer of SkiMo
(ours) contains early explorations that span to most corners in the maze.
After it finds the goal, it exploits this knowledge and commits to paths that
are between the start location and the goal (in dark blue).

Dreamer and TD-MPC only explore a small fraction of the maze because they are
prone to get stuck at walls without guided exploration from skills and skill
priors. SPiRL + Dreamer, SPiRL + TD-MPC, and SkiMo w/o joint training explore
better than Dreamer and TD-MPC, but all fail to find the goal. This is because
without the joint training of the model and policy, the skill space is only
optimized for action reconstruction, not for planning, which makes long-
horizon exploration and exploitation harder.

On the other hand, SkiMo + SAC and SPiRL are able to explore the most portion
of the maze, but even after the agent finds the goal through exploration, it
continues to explore and does not exploit this experience to consistently
accomplish the task (darker blue). This could attribute to the difficult long-
horizon credit assignment problem which makes policy learning slow, and the
reliance on skill prior which encourages exploration. On the contrary, our
skill dynamics model effectively absorbs prior experience to generate goal-
achieving imaginary rollouts for the actor and critic to learn from, which
makes task learning more efficient.

We find the skill dynamics model useful in guiding the agent explore
coherently and exploit efficiently. Without a temporally-extended model, DADS,
LSP, and SkiMo w/ single-step dynamics fail to reach the goal. Even though
they likewise condition on the skill latent, they still need to roll out the
dynamics model step-by-step to predict future states. The single-step
prediction is prone to compounding error for long-horizon planning. As a
result, these agents do not collect sufficiently meaningful trajectories for
the policy to learn. Additionally, SkiMo w/o CEM performs as well as SkiMo,
indicating that CEM planning is not essential after the agent has already
learned a good policy in Maze. Nevertheless, these qualitative results
corroborate the effectiveness of our method.

###  B.2 Long-horizon Prediction

To compare the long-term prediction ability of the skill dynamics and flat
dynamics, we visualize imagined trajectories by sampling trajectory clips of
500 timesteps from the agent’s replay buffer (the maximum episode length in
Maze is 2,000), and predicting the latent state 500 steps ahead, which will be
decoded using the observation decoder, given the initial state and 500 ground-
truth actions (50 skills for SkiMo). The similarity between the imagined
trajectory and the ground truth trajectory can indicate whether the model can
make accurate predictions far into the future, producing useful imaginary
rollouts for policy learning and planning.

SkiMo is able to reproduce the ground truth trajectory with little prediction
error even when traversing through hallways and doorways while Dreamer
struggles to make accurate long-horizon predictions due to error accumulation.
This is mainly because SkiMo allows temporal abstraction in the dynamics
model, thereby enabling temporally-extended prediction and reducing step-by-
step prediction error.

(a) Dreamer

(b) SkiMo (Ours)

Figure 10: Prediction results of 500 timesteps using a flat single-step model
(a) and skill dynamics model (b), given the ground truth starting state and
500 actions (50 skills for SkiMo). The predicted states from the flat model
deviate from the ground truth trajectory quickly while the prediction of our
skill dynamics model has little error.

##  Appendix C Implementation Details

###  C.1 Computing Resources

Our approach and all baselines are implemented in PyTorch [46]. All
experiments are conducted on a workstation with an Intel Xeon E5-2640 v4 CPU
and a NVIDIA Titan Xp GPU. Pre-training of the skill policy and skill dynamics
model takes around 10 hours. Downstream RL for 2M timesteps takes around 18
hours. The policy and model update frequency is the same over all algorithms
but Dreamer [3] and TD-MPC [4]. Since Dreamer and TD-MPC train on primitive
actions, it has 10 times more frequent model and policy updates than skill-
based algorithms, which leads to slower training (about 52 hours).

###  C.2 Algorithm Implementation Details

For the baseline implementations, we use the official code for SPiRL, DADS,
and LSP. We re-implemented Dreamer and TD-MPC in PyTorch, which are verified
on DeepMind Control Suite [47]. The table below (Table 1) compares key
components of SkiMo with model-based and skill-based baselines and ablated
methods.

Table 1: Comparison to prior work and ablated methods. Method | Skill-based | Model-based | Joint training
---|---|---|---
Dreamer [3] and TD-MPC [4] | ✗ | ✓ | ✗
DADS [29] and LSP [36] | ✓ | ✓ | ✗
SPiRL [12] | ✓ | ✗ | ✗
SPiRL + Dreamer and SPiRL + TD-MPC | ✓ | ✓ | ✗
SkiMo w/o joint training | ✓ | ✓ | ✗
SkiMo + SAC | ✓ | ✗ | ✓
SkiMo (Ours) and SkiMo w/o CEM | ✓ | ✓ | ✓

Dreamer [3] We use the same hyperparameters with the official implementation.

TD-MPC [4] We use the same hyperparameters with the official implementation,
except that we do not use the prioritized experience replay [48]. The same
implementation is used for the SPiRL + TD-MPC baseline and our method with
only minor modification.

SPiRL [12] We use the official implementation of the original paper and use
the hyperparameters suggested in the official implementation.

SPiRL + Dreamer [12] We use our implementation of Dreamer and simply replace
the action space with the latent skill space of SPiRL. We use the same pre-
trained SPiRL skill policy and skill prior networks with the SPiRL baseline.
Initializing the high-level downstream task policy with the skill prior, which
is critical for downstream learning performance [12], is not possible due to
the policy network architecture mismatch between Dreamer and SPiRL. Thus, we
only use the prior divergence to regularize the high-level policy instead.
Directly pre-train the high-level policy did not lead to better performance,
but it might have worked better with more tuning.

SPiRL + TD-MPC [4] Similar to SPiRL + Dreamer, we use our implementation of
TD-MPC and replace the action space with the latent skill space of SPiRL. The
initialization of the task policy is also not available due to the different
architecture used for TD-MPC.

DADS [29] We use the official implementation and hyperparameters of the
original paper, except that we use DADS on a sparse reward setup since dense
reward is not available in our tasks.

LSP [36] We use the code provided by the authors and the default
hyperparameters in the code.

SkiMo (Ours) The skill-based RL part of our method is inspired by Pertsch et
al. [12] and the model-based component is inspired by Hansen et al. [4] and
Hafner et al. [3]. We elaborate our skill and skill dynamics learning in
Algorithm 1, planning algorithm in Algorithm 2, and model-based RL in
Algorithm 3. Table 2 lists the all hyperparameters that we used.

Algorithm 1 SkiMo (_skill and skill dynamics learning_)

0: $\mathcal{D}$: offline task-agnostic data

1: Randomly initialize $\theta,\psi$

2: $\psi^{-}\leftarrow\psi$ $\triangleright$ initialize target network

3: for each iteration do

4: Sample mini-batch $B=(\mathbf{s},\mathbf{a})_{(0:NH)}\sim\mathcal{D}$

5:
$[\theta,\psi]\leftarrow[\theta,\psi]-\lambda_{[\theta,\psi]}\nabla_{[\theta,\psi]}\mathcal{L}(B)$
$\triangleright$ $\mathcal{L}$ from Equation (5)

6: $\psi^{-}\leftarrow(1-\tau)\psi^{-}+\tau\psi$ $\triangleright$ update
target network

7: end for

8: return $\theta,\psi,\psi^{-}$

Algorithm 2 SkiMo (_CEM planning_)

0: $\theta,\psi,\phi:$ learned parameters, $\mathbf{s}_{t}$: current state

1: $\mathbf{\mu}^{0},\mathbf{\sigma}^{0}\leftarrow\mathbf{0},\mathbf{1}$
$\triangleright$ initialize sampling distribution

2: for $i=1,...,N_{\text{CEM}}$ do

3: Sample $N_{\text{sample}}$ trajectories of length $N$ from
$\mathcal{N}\big{(}\mu^{i-1},(\sigma^{i-1})^{2}\big{)}$ $\triangleright$
sample skill sequences from normal distribution

4: Sample $N_{\pi}$ trajectories of length $N$ using $\pi_{\phi},D_{\psi}$
$\triangleright$ sample skill sequences via imaginary rollouts

5: Estimate $N$-step returns of $N_{\text{sample}}+N_{\pi}$ trajectories using
$R_{\phi},Q_{\phi}$

6: Compute $\mu^{i},\sigma^{i}$ with top-k return trajectories
$\triangleright$ update parameters for next iteration

7: end for

8: Sample a skill
$\mathbf{z}\sim\mathcal{N}\big{(}\mu^{N_{\text{CEM}}},(\sigma^{N_{\text{CEM}}})^{2}\big{)}$

9: return $\mathbf{z}$

Algorithm 3 SkiMo (_downstream RL_)

0: $\theta,\psi,\psi^{-}:$ pre-trained parameters

1: $\mathcal{B}\leftarrow\emptyset$ $\triangleright$ initialize replay buffer

2: Randomly initialize $\phi$

3: $\phi^{-}\leftarrow\phi$ $\triangleright$ initialize target network

4: $\pi_{\phi}\leftarrow p_{\theta}$ $\triangleright$ initialize task policy
with skill prior

5: for not converged do

6: $t\leftarrow 0,s_{0}\sim\rho_{0}$ $\triangleright$ initialize episode

7: for episode not done do

8: $\mathbf{z}_{t}\sim$ CEM$(\mathbf{s}_{t})$ $\triangleright$ MPC with CEM
planning in Algorithm 2

9: $\mathbf{s},r_{t}\leftarrow\mathbf{s}_{t},0$

10: for $H$ steps do

11: $\mathbf{s},r\leftarrow$
ENV$(\mathbf{s},\pi^{L}_{\theta}(E_{\psi}(\mathbf{s}),\mathbf{z}_{t}))$
$\triangleright$ rollout low-level skill policy

12: $r_{t}\leftarrow r_{t}+r$

13: end for

14:
$\mathcal{B}\leftarrow\mathcal{B}\cup(\mathbf{s}_{t},\mathbf{z}_{t},r_{t})$
$\triangleright$ collect $H$-step environment interaction

15: $t\leftarrow t+H$

16: $\mathbf{s}_{t}\leftarrow\mathbf{s}$

17: Sample mini-batch $B=(\mathbf{s},\mathbf{z},r)_{(0:N)}\sim\mathcal{B}$

18:
$[\phi,\psi]\leftarrow[\phi,\psi]-\lambda_{[\phi,\psi]}\nabla_{[\phi,\psi]}\mathcal{L}^{\prime}_{\text{REC}}(B)$
$\triangleright$ $\mathcal{L}^{\prime}_{\text{REC}}$ from Equation (6)

19:
$\phi_{\pi}\leftarrow\phi_{\pi}-\lambda_{\phi}\nabla_{\phi_{\pi}}\mathcal{L}_{\text{RL}}(B)$
$\triangleright$ $\mathcal{L}_{\text{RL}}$ from Equation (7). Update only
policy parameters

20: $\psi^{-}\leftarrow(1-\tau)\psi^{-}+\tau\psi$ $\triangleright$ update
target network

21: $\phi^{-}\leftarrow(1-\tau)\phi^{-}+\tau\phi$ $\triangleright$ update
target network

22: end for

23: end for

24: return $\psi,\phi$

Table 2: SkiMo hyperparameters. Hyperparameter | Value
---|---
| Maze | FrankaKitchen | CALVIN
Model architecture
# Layers of $O_{\theta},p_{\theta},\pi^{L}_{\theta},E_{\psi},D_{\psi},\pi_{\phi},R_{\phi},Q_{\phi}$ | 5
Activation funtion | elu
Hidden dimension | 128 | 128 | 256
State embedding dimension | 128 | 256 | 256
Skill encoder ($q_{\theta}$) | 5-layer MLP | LSTM | LSTM
Skill encoder hidden dimension | 128
Pre-training
Pre-training batch size | 512
# Training mini-batches per update | 5
Model-Actor joint learning rate ($\lambda_{[\theta,\psi]}$) | 0.001
Encoder KL regularization ($\beta$) | 0.0001
Reconstruction loss coefficient ($\lambda_{\text{O}}$) | 1
Consistency loss coefficient ($\lambda_{\text{L}}$) | 2
Low-level actor loss coefficient ($\lambda_{\text{BC}}$) | 2
Planning discount ($\rho$) | 0.5
Skill prior loss coefficient ($\lambda_{\text{SP}}$) | 1
Downstream RL
Model learning rate | 0.001
Actor learning rate | 0.001
Skill dimension | 10
Skill horizon ($H$) | 10
Planning horizon ($N$) | 10 | 3 | 1
Batch size | 128 | 256 | 256
# Training mini-batches per update | 10
State normalization | True | False | False
Prior divergence coefficient ($\alpha$) | 1 | 0.5 | 0.1
Alpha learning rate | 0.0003 | 0 | 0
Target divergence | 3 | N/A | N/A
# Warm up step | 50,000 | 5,000 | 5,000
# Environment step per update | 500
Replay buffer size | 1,000,000
Target update frequency | 2
Target update tau ($\tau$) | 0.01
Discount factor ($\gamma$) | 0.99
Reward loss coefficient ($\lambda_{\text{R}}$) | 0.5
Value loss coefficient ($\lambda_{\text{Q}}$) | 0.1
CEM
CEM iteration ($N_{\text{CEM}}$) | 6
# Sampled trajectories ($N_{\text{sampled}}$) | 512
# Policy trajectories ($N_{\pi}$) | 25
# Elites ($k$) | 64
CEM momentum | 0.1
CEM temperature | 0.5
Maximum std | 0.5
Minimum std | 0.01
Std decay step | 100,000 | 25,000 | 25,000
Horizon decay step | 100,000 | 25,000 | 25,000

###  C.3 Environments and Offline Data

##### Maze [43, 16]

Since our goal is to leverage offline data collected from diverse tasks in the
same environment, we use a variant of the D4RL maze environment [43],
suggested in Pertsch et al. [16]. The maze is of size $40\times 40$; an
initial state is randomly sampled near a pre-defined region (the green circle
in Figure 3(a)); and the goal is fixed shown as the red circle in Figure 3(a).
The observation consists of the agent’s 2D position and velocity. The agent
moves around the maze by controlling the continuous value of its $(x,y)$
velocity. The maximum episode length is 2,000 but an episode is also
terminated if the agent reaches the circle around the goal with radius 2. The
reward of 100 is given at task completion. We use the offline data from
Pertsch et al. [16], consisting of 3,046 trajectories with randomly sampled
start and goal state pairs. Thus, the offline data and downstream task share
the same environment, but have different start and goal states (i.e. different
tasks). This data can be used to extract short-horizon skills like navigating
hallways or passing through narrow doors.

##### Kitchen [32, 43]

The 7-DoF Franka Emika Panda robot arm needs to perform four sequential tasks:
open microwave, move kettle, turn on bottom burner, and flip light switch. The
agent has a 30D observation space (11D robot proprioceptive state and 19D
object states), which removes a constant 30D goal state in the original
environment, and 9D action space (7D joint velocity and 2D gripper velocity).
The agent receives a reward of 1 for every sub-task completion. The episode
length is 280 and an episode also ends once all four sub-tasks are completed.
The initial state is set with a small noise in every state dimension. We use
603 trajectories collected by teleoperation from Gupta et al. [32] as the
offline task-agnostic data. The data involves interaction with all seven
manipulatable objects in the environment, but during downstream learning the
agent needs to execute an unseen sequence of four subtasks. Thus, the agent
can transfer a rich set of manipulation skills, but needs to recombine them in
new ways to solve the task.

##### Mis-aligned Kitchen [16]

The environment and task-agnostic data are the same with Kitchen but we use
the different downstream task: open microwave, flip light switch, slide
cabinet door, and open hinge cabinet, as illustrated in Figure 3(c). This task
ordering is not aligned with the sub-task transition probabilities of the
task-agnostic data, which leads to challenging exploration following the prior
from data. This is because the transition probabilities in the Kitchen human-
teleoperated dataset are not uniformly distributed; instead, certain
transitions are more likely than others. For example, the first transition in
our target task — from opening the microwave to flipping the light switch — is
very unlikely to be observed in the training data. This simulates the real-
world scenario where the large offline dataset may not be meticulously curated
for the target task.

##### CALVIN [44]

We adapt the CALVIN environment [44] for long-horizon learning with the state
observation. The CALVIN environment uses a Franka Emika Panda robot arm with
7D end-effector pose control (relative 3D position, 3D orientation, 1D gripper
action). The 21D observation space consists of the 15D proprioceptive robot
state and 6D object state. We use the teleoperated play data (Task
D$\rightarrow$Task D) of 1,239 trajectories from Mees et al. [44] as our task-
agnostic data. The agent receives a sparse reward of 1 for every sub-task
completion in the correct order: open drawer, turn on lightbulb, move slider
left, and turn on LED. The episode length is 360 and an episode also ends if
all four sub-tasks are completed. In data, there exist 34 available target
sub-tasks, and each sub-task can transition to any other sub-task, which makes
any transition probability lower than 0.1% on average.

(a) In pretraining, SkiMo leverages offline task-agnostic data to extract
skill dynamics and a skill repertoire. Unlike prior works that keep the model
and skill policy training separate, we propose to jointly train them to
extract a skill space that is conducive to plan upon.

(b) In downstream RL, we learn a high-level task policy in the skill space
(skill-based RL) and leverage the skill dynamics model to generate imaginary
rollouts for policy optimization and planning (model-based RL).

Figure 11: Illustration of our algorithm, SkiMo.

##  Appendix D Application to Real Robot Systems

Our algorithm is designed to be applied on real robot systems by improving
sample efficiency of RL using a temporally-abstracted dynamics model.
Throughout the extensive experiments in simulated robotic manipulation
environments, we show that our approach achieves superior sample efficiency
over prior skill-based and model-based RL, which gives us strong evidence for
the application to real robot systems. Especially in Kitchen and CALVIN, our
approach improves the sample efficiency of learning long-horizon manipulation
tasks with a 7-DoF Franka Emika Panda robot arm. Our approach consists of
three phases: (1) task-agnostic data collection, (2) skills and skill dynamics
model learning, and (3) downstream task learning. In each of these phases, our
approach can be applied to physical robots:

##### Task-agnostic data collection

Our approach is designed to fully leverage task-agnostic data without any
reward or task annotation. In addition to extracting skills and skill priors,
we further learn a skill dynamics model from this task-agnostic data.
Maximizing the utility of task-agnostic data is critical for real robot
systems as data collection with physical robots itself is very expensive. Our
method does not require any manual labelling of data and simply extracts
skills, skill priors, and skill dynamics model from raw states and actions,
which makes our method scalable.

##### Pre-training of skills and skill dynamics model

Our approach trains the skill policy, skill dynamics model, and skill prior
from the offline task-agnostic dataset, without requiring any additional real-
world robot interactions.

##### Downstream task learning

The goal of our work is to leverage skills and skill dynamics model to allow
for more efficient downstream learning, i.e., requires less interactions of
the agent with the environment for training the policy. This is especially
important on real robot systems where a robot-environment interaction is slow,
dangerous, and costly. Our approach directly addresses this concern by
learning a policy from imaginary rollouts rather than actual environment
interactions.

In summary, we believe that SkiMo can be applied to real-world robot systems
with only minor modifications.