$\hat{\textbf{B}}_{n}=\textbf{B}_{n}^{*}(\hat{\bm{\beta}}_{n})$,
$\hat{\textbf{A}}_{n}=\textbf{A}_{n}^{*}(\hat{\bm{\beta}}_{n})$.

The quasi-likelihood $\mathcal{L}_{n}^{*}(\textbf{y}^{*},\bm{\beta})$ can be
expanded about $\hat{\bm{\beta}}_{n}$ by Taylor’s theorem, we have, for any
$\bm{\beta}$,

$\mathcal{L}_{n}^{*}(\textbf{y}^{*},\bm{\beta})=\mathcal{L}_{n}^{*}(\textbf{y}^{*},\hat{\bm{\beta}}_{n})+(\bm{\beta}-\hat{\bm{\beta}}_{n})^{T}\Psi^{*}_{n}(\hat{\bm{\beta}}_{n})-\frac{1}{2}(\bm{\beta}-\hat{\bm{\beta}}_{n})^{T}\textbf{A}_{n}^{*}(\bm{\beta}^{\prime})(\bm{\beta}-\hat{\bm{\beta}}_{n})^{T},$

where $\bm{\beta}^{\prime}$ is on the lime segment between $\bm{\beta}$ and
$\hat{\bm{\beta}}_{n}$. Setting
$\textbf{u}=\delta^{-1}\hat{\textbf{B}}_{n}^{1/2}(\bm{\beta}-\hat{\bm{\beta}}_{n})$,
it follows:

$\mathcal{L}_{n}^{*}(\textbf{y}^{*},\bm{\beta})-\mathcal{L}_{n}^{*}(\textbf{y}^{*},\hat{\bm{\beta}}_{n})=\delta\textbf{u}^{T}\hat{\textbf{B}}_{n}^{-1/2}\Psi^{*}_{n}(\hat{\bm{\beta}}_{n})-\frac{1}{2}\delta^{2}\textbf{u}^{T}\hat{\textbf{T}}_{n}(\bm{\beta}^{\prime})\textbf{u}.$

where
$\hat{\textbf{T}}_{n}(\bm{\beta}^{\prime})=\hat{\textbf{B}}_{n}^{-1/2}\textbf{A}_{n}^{*}(\bm{\beta}^{\prime})\hat{B}_{n}^{-1/2}$.
Denote
$\partial\hat{H}_{n}(\delta)=\\{\bm{\beta}\in\mathbb{R}^{p}:\|\hat{\textbf{B}}_{n}^{1/2}(\bm{\beta}-\hat{\bm{\beta}}_{n})\|=\delta\\}$
as the boundary of the closed set
$\hat{H}_{n}(\delta)=\\{\bm{\beta}\in\mathbb{R}^{p}:\|\hat{\textbf{B}}_{n}^{1/2}(\bm{\beta}-\hat{\bm{\beta}}_{n})\|\leq\delta\\}$.
Note that $\bm{\beta}\in\partial\hat{H}_{n}(\delta)$ if and only if
$\|\textbf{u}\|=1$, and that $\bm{\beta}\in\partial\hat{H}_{n}(\delta)$
indicates that $\bm{\beta}^{\prime}\in\hat{H}_{n}(\delta)$ by the convexity of
$\hat{H}_{n}(\delta)$. We have
$\max_{\|\textbf{u}\|=1}\textbf{u}^{T}\hat{\textbf{B}}_{n}^{-1/2}\Psi^{*}_{n}(\hat{\bm{\beta}}_{n})=\|\hat{\textbf{B}}_{n}^{-1/2}\Psi^{*}_{n}(\hat{\bm{\beta}}_{n})\|$,
and according to the second part of condition C2, we have
$\min_{\|\textbf{u}\|=1}\textbf{u}^{T}\hat{\textbf{T}}_{n}(\bm{\beta}^{\prime})\textbf{u}\geq\min_{\bm{\beta}\in\hat{H}_{n}(\delta)}\lambda_{\min}\\{\hat{\textbf{T}}_{n}(\bm{\beta})\\}\geq
c$. Thus we have

$\max_{\bm{\beta}\in\partial\hat{H}_{n}(\delta)}\mathcal{L}_{n}^{*}(\textbf{y}^{*},\bm{\beta})-\mathcal{L}_{n}^{*}(\textbf{y}^{*},\hat{\bm{\beta}}_{n})\leq\delta\Big{(}\|\hat{\textbf{B}}_{n}^{-1/2}\Psi^{*}_{n}(\hat{\bm{\beta}}_{n})\|-\frac{1}{2}c\delta\Big{)},$

which entails that according to Markov’s inequality

$\begin{split}&P\Big{(}\max_{\bm{\beta}\in\partial\hat{H}_{n}(\delta)}\mathcal{L}_{n}^{*}(\textbf{y}^{*},\bm{\beta})<\mathcal{L}_{n}^{*}(\textbf{y}^{*},\hat{\bm{\beta}}_{n})\Big{)}\geq
P\Big{(}\|\hat{\textbf{B}}_{n}^{-1/2}\Psi^{*}_{n}(\hat{\bm{\beta}}_{n})\|^{2}\leq\frac{1}{4}c^{2}\delta^{2}\Big{)}\\\
&\geq
1-\frac{\mathbb{E}\Big{(}\|\hat{\textbf{B}}_{n}^{-1/2}\Psi^{*}_{n}(\hat{\bm{\beta}}_{n})\|^{2}\Big{)}}{c^{2}\delta^{2}/4}\geq
1-\frac{4p}{c^{2}\delta^{2}},\end{split}$

along with
$\mathbb{E}\Big{(}\|\hat{\textbf{B}}_{n}^{-1/2}\Psi^{*}_{n}(\hat{\bm{\beta}}_{n})\|^{2}\Big{)}=\mbox{tr}\\{\mathbb{E}(\Psi^{*}_{n}(\hat{\bm{\beta}}_{n})\Psi^{*}_{n}(\hat{\bm{\beta}}_{n})^{T})\hat{\textbf{B}}_{n}^{-1}\\}=\mbox{tr}(\textbf{I}_{p})=p$.
Thus, for any $0<\eta<1$, there exist $\delta=2p^{1/2}/(c\eta^{1/2})>0$, such
that

$P\Big{(}\max_{\bm{\beta}\in\partial\hat{H}_{n}(\delta)}\mathcal{L}_{n}^{*}(\textbf{y}^{*},\bm{\beta})<\mathcal{L}_{n}^{*}(\textbf{y}^{*},\hat{\bm{\beta}}_{n})\Big{)}\geq
1-\eta.$

In fact, if
$\\{\max_{\bm{\beta}\in\partial\hat{H}_{n}(\delta)}\mathcal{L}_{n}^{*}(\textbf{y}^{*},\bm{\beta})<\mathcal{L}_{n}^{*}(\textbf{y}^{*},\hat{\bm{\beta}}_{n})\\}$,
the continuous function $\mathcal{L}_{n}^{*}(\textbf{y}^{*},\cdot)$ has a
local maximum in the interior of $\hat{H}_{n}(\delta)$, and since the
condition C1 which implies that $\mathcal{L}_{n}^{*}(\textbf{y}^{*},\cdot)$ is
strictly concave, this maximum must be located at $\hat{\bm{\beta}}_{n}^{*}$.
Thus,

$P(\hat{\bm{\beta}}_{n}^{*}\in\hat{H}_{n}(\delta))=P(\|\hat{\textbf{B}}_{n}^{1/2}(\hat{\bm{\beta}}_{n}^{*}-\hat{\bm{\beta}}_{n})\|\leq\delta)\geq
1-\eta,$

for any $\eta\in(0,1)$, together with the first part of condition C2, proves
that $\hat{\bm{\beta}}_{n}^{*}-\hat{\bm{\beta}}_{n}=o_{p}(1)$.

The target “true” parameter of bootstrap model $\bm{\beta}_{n,0}^{*}$ is the
root of the expected score function

$\displaystyle\mathbb{E}\Big{(}\Psi_{n}^{*}(\bm{\beta}_{n,0}^{*})\Big{)}=\sum_{i=1}^{n}\textbf{x}_{i}D_{i}(\bm{\beta}_{n,0}^{*})V_{i}^{-1}(\bm{\beta}_{n,0}^{*})\Big{(}m_{i}^{*}-h_{i}(\bm{\beta}_{n,0}^{*})\Big{)}=0.$
(12)

Define $\textbf{C}_{n}^{*}=\textbf{B}_{n}^{*-1/2}\textbf{A}_{n}^{*}$, where
$\textbf{A}_{n}^{*}=\textbf{A}^{*}_{n}(\bm{\beta}_{n,0}^{*})$ and
$\textbf{B}_{n}^{*}=\textbf{B}^{*}_{n}(\bm{\beta}_{n,0}^{*})$, thus,
$\textbf{C}_{n}^{*}(\hat{\bm{\beta}}_{n}^{*}-\bm{\beta}_{n,0}^{*})\stackrel{{\scriptstyle
d}}{{\rightarrow}}N(0,\textbf{I}_{p})$ according to asymptotic normality for
QMLE. Observe that

$\textbf{C}_{n}^{*}(\hat{\bm{\beta}}_{n}^{*}-\hat{\bm{\beta}}_{n})=\textbf{C}_{n}^{*}(\hat{\bm{\beta}}_{n}^{*}-\bm{\beta}_{n,0}^{*})+\textbf{C}_{n}^{*}(\bm{\beta}_{n,0}^{*}-\hat{\bm{\beta}}_{n}).$

We have $\hat{\bm{\beta}}_{n}^{*}\rightarrow\hat{\bm{\beta}}_{n}$, and
$\hat{\bm{\beta}}_{n}^{*}\rightarrow\bm{\beta}_{n,0}^{*}$ according to the
consistency of QMLE, thus, we can obtain that
$|\bm{\beta}_{n,0}^{*}-\hat{\bm{\beta}}_{n}|\rightarrow 0$ conditional on
$\hat{\bm{\beta}}_{n}$. According to Slutsky’s lemma,
$\textbf{C}_{n}^{*}(\hat{\bm{\beta}}_{n}^{*}-\hat{\bm{\beta}}_{n})\stackrel{{\scriptstyle
d}}{{\rightarrow}}N(0,\textbf{I}_{p})$. Next we illustrate the consistency of
$\textbf{C}^{*}_{n}\rightarrow\textbf{C}_{n}$.

Note that the objective function
$\mathcal{L}_{n}(\bm{\beta})=\sum_{i=1}^{n}\Big{(}y_{i}\theta_{i}-b(\theta_{i})\Big{)}$.
Specifically, in probit binary model, both of
$\theta_{i}=\log(\frac{\mu_{i}}{1-\mu_{i}})$ and
$b(\theta_{i})=\log(1+e^{\theta_{i}})$ are the functions of only $\bm{\beta}$
give $\textbf{x}_{i}$. The pseudo-true parameter $\bm{\beta}_{n,0}^{\dagger}$
is to maximum the $\mathbb{E}(\mathcal{L}_{n}(\bm{\beta}))$, that is,

$\bm{\beta}_{n,0}^{\dagger}=\arg\max_{\bm{\beta}}\mathbb{E}(\mathcal{L}_{n}(\bm{\beta}))=\arg\max_{\bm{\beta}}\left\\{\sum_{i=1}^{n}\Big{(}m_{i}\theta_{i}-b(\theta_{i})\Big{)}\right\\}.$

While $\bm{\beta}_{n,0}^{*}$ is to maximum the
$\mathbb{E}(\mathcal{L}_{n}^{*}(\bm{\beta}))$, and

$\bm{\beta}_{n,0}^{*}=\arg\max_{\bm{\beta}}\mathbb{E}(\mathcal{L}_{n}^{*}(\bm{\beta}))=\arg\max_{\bm{\beta}}\left\\{\sum_{i=1}^{n}\Big{(}m_{i}^{*}\theta_{i}-b(\theta_{i})\Big{)}\right\\}.$

Because $m_{i}^{*}\rightarrow m_{i}$, for any $\bm{\beta}$ in parameter space,
we have
$\mathbb{E}(\mathcal{L}_{n}^{*}(\bm{\beta}))\rightarrow\mathbb{E}(\mathcal{L}_{n}(\bm{\beta}))$,
according to Lemma A in Newey and Powell (1987),
$\bm{\beta}_{n,0}^{*}\rightarrow\bm{\beta}_{n,0}^{\dagger}$. We also have
$m_{i}^{*}\rightarrow m_{i}$, and $\sigma_{i}^{*}\rightarrow\sigma_{i}$, the
results follows immediately,
$\textbf{A}_{n}^{*}\rightarrow\widetilde{\textbf{A}}_{n}$,
$\textbf{B}_{n}^{*}\rightarrow\widetilde{\textbf{B}}_{n}$, thus, we have
$\textbf{C}_{n}^{*}\rightarrow\textbf{C}_{n}$.

∎

Proof of Theorem 3

###### Proof.

Using Taylor expansion, we have
$\textbf{0}=\Psi_{n}^{*}(\hat{\bm{\beta}}_{n}^{*})=\Psi_{n}^{*}(\bm{\beta}_{n,0}^{*})+\textbf{A}_{n}^{*}(\bm{\beta}_{n,0}^{*})(\hat{\bm{\beta}}_{n}^{*}-\bm{\beta}_{n,0}^{*})+R_{n}(\hat{\bm{\beta}}_{n}^{*})$.
We can show that
$\mathbb{E}\left(\textbf{A}_{n}^{*}(\bm{\beta}_{n,0}^{*})(\bm{\beta}_{n,0}^{*}-\hat{\bm{\beta}}_{n}^{*})\right)=\mathbb{E}\left(\Psi_{n}^{*}(\bm{\beta}_{n,0}^{*})\right)+\mathbb{E}\left(R_{n}(\hat{\bm{\beta}}_{n}^{*})\right)$,
and

$\mathbb{E}\left(\textbf{A}_{n}^{*}(\bm{\beta}_{n,0}^{*})(\hat{\bm{\beta}}_{n}^{*}-\bm{\beta}_{n,0}^{*})(\hat{\bm{\beta}}_{n}^{*}-\bm{\beta}_{n,0}^{*})^{T}\textbf{A}_{n}^{*}(\bm{\beta}_{n,0}^{*})^{T}\right)=\mathbb{E}\left(\Psi_{n}^{*}(\bm{\beta}_{n,0}^{*})^{T}\Psi_{n}^{*}(\bm{\beta}_{n,0}^{*})\right)+\mathbb{E}\left(R_{n}^{2}(\hat{\bm{\beta}}_{n}^{*})\right),$

where
$R_{n}(\hat{\bm{\beta}}_{n}^{*})=o(\hat{\bm{\beta}}_{n}^{*}-\bm{\beta}_{n,0}^{*})=o(1)$
and $\mathbb{E}\left(\Psi_{n}^{*}(\bm{\beta}_{n,0}^{*})\right)=\textbf{0}$.

$\begin{split}\mbox{Var}\left(\hat{\bm{\beta}}_{n}^{*}\right)&=\mathbb{E}\left((\hat{\bm{\beta}}_{n}^{*}-\mathbb{E}(\hat{\bm{\beta}}_{n}^{*}))(\hat{\bm{\beta}}_{n}^{*}-\mathbb{E}(\hat{\bm{\beta}}_{n}^{*}))^{T}\right)\\\
&\leq\mathbb{E}\left((\hat{\bm{\beta}}_{n}^{*}-\bm{\beta}_{n,0}^{*})(\hat{\bm{\beta}}_{n}^{*}-\bm{\beta}_{n,0}^{*})^{T}\right)+(\mathbb{E}(\hat{\bm{\beta}}_{n}^{*})-\bm{\beta}_{n,0}^{*})(\mathbb{E}(\hat{\bm{\beta}}_{n}^{*})-\bm{\beta}_{n,0}^{*})^{T}\\\
&\leq\textbf{A}_{n}^{*-1}\textbf{B}_{n}^{*}\textbf{A}_{n}^{*-1}+\textbf{A}_{n}^{*-1}\mathbb{E}(R_{n}^{2}(\hat{\bm{\beta}}_{n}^{*}))\textbf{A}_{n}^{*-1}+\textbf{A}_{n}^{*-1}(\mathbb{E}(R_{n}(\hat{\bm{\beta}}_{n}^{*})))^{2}\textbf{A}_{n}^{*-1}.\end{split}$

Thus we have
$\mbox{Var}(\hat{\bm{\beta}}_{n}^{*})\rightarrow\textbf{A}_{n}^{*-1}\textbf{B}_{n}^{*}\textbf{A}_{n}^{*-1}$,
according to the step 3 in the proof of Theorem 2,
$\textbf{A}_{n}^{*-1}\textbf{B}_{n}^{*}\textbf{A}_{n}^{*-1}\rightarrow\widetilde{\textbf{A}}_{n}^{-1}\widetilde{\textbf{B}}_{n}\widetilde{\textbf{A}}_{n}^{-1}$.
This completes the proof.

∎


# Choreographer: Learning and
Adapting Skills in Imagination

Pietro Mazzaglia ${}^{\phantom{*}}$ 1 Tim Verbelen 1 Bart Dhoedt 1 Alexandre
Lacoste 2 Sai Rajeswar 2

1 Ghent University - imec 2 ServiceNow Research Correspondence to:
<EMAIL_ADDRESS>
###### Abstract

Unsupervised skill learning aims to learn a rich repertoire of behaviors
without external supervision, providing artificial agents with the ability to
control and influence the environment. However, without appropriate knowledge
and exploration, skills may provide control only over a restricted area of the
environment, limiting their applicability. Furthermore, it is unclear how to
leverage the learned skill behaviors for adapting to downstream tasks in a
data-efficient manner. We present Choreographer, a model-based agent that
exploits its world model to learn and adapt skills in imagination. Our method
decouples the exploration and skill learning processes, being able to discover
skills in the latent state space of the model. During adaptation, the agent
uses a meta-controller to evaluate and adapt the learned skills efficiently by
deploying them in parallel in imagination. Choreographer is able to learn
skills both from offline data, and by collecting data simultaneously with an
exploration policy. The skills can be used to effectively adapt to downstream
tasks, as we show in the URL benchmark, where we outperform previous
approaches from both pixels and states inputs. The learned skills also explore
the environment thoroughly, finding sparse rewards more frequently, as shown
in goal-reaching tasks from the DMC Suite and Meta-World.

Project website: https://skillchoreographer.github.io/

##  1 Introduction

Deep Reinforcement Learning (RL) has yielded remarkable success in a wide
variety of tasks ranging from game playing (Mnih et al., 2013; Silver et al.,
2016) to complex robot control (Smith et al., 2022; OpenAI et al., 2019).
However, most of these accomplishments are specific to mastering a single task
relying on millions of interactions to learn the desired behavior. Solving a
new task generally requires to start over, collecting task-specific data, and
learning a new agent from scratch.

Instead, natural agents, such as humans, can quickly adapt to novel situations
or tasks. Since their infancy, these agents are intrinsically motivated to try
different movement patterns, continuously acquiring greater perceptual
capabilities and sensorimotor experiences that are essential for the formation
of future directed behaviors (Corbetta, 2021). For instance, a child who
understands how object relations work, e.g. has autonomously learned to stack
one block on top of another, can quickly master how to create structures
comprising multiple objects (Marcinowski et al., 2019).

With the same goal, unsupervised RL (URL) methods aim to leverage intrinsic
motivation signals, used to drive the agent’s interaction with the
environment, to acquire generalizable knowledge and behaviors. While some URL
approaches focus on exploring the environment (Schmidhuber, 1991; Mutti et
al., 2020; Bellemare et al., 2016), _competence-based_ (Laskin et al., 2021)
methods aim to learn a set of options or skills that provide the agent with
the ability to control the environment (Gregor et al., 2016; Eysenbach et al.,
2019), a.k.a. empowerment (Salge et al., 2014).

Learning a set of options can provide an optimal set of behaviors to quickly
adapt and generalize to new tasks (Eysenbach et al., 2021). However, current
methods still exhibit several limitations. Some of these are due to the nature
of the skill discovery objective (Achiam et al., 2018), struggling to capture
behaviors that are natural and meaningful for humans. Another major issue with
current methods is the limited exploration brought by competence-based
methods, which tend to commit to behaviors that are easily discriminable but
that guarantee limited control over a small area of the environment. This
difficulty has been both analyzed theoretically (Campos et al., 2020) and
demonstrated empirically (Laskin et al., 2021; Rajeswar et al., 2022).

A final important question with competence-based methods arises when adapting
the skills learned without supervision to downstream tasks: how to exploit the
skills in an efficient way, i.e. using the least number of environment
interactions? While one could exhaustively test all the options in the
environment, this can be expensive when learning a large number of skills
(Eysenbach et al., 2019) or intractable, for continuous skill spaces (Kim et
al., 2021; Liu & Abbeel, 2021a).

In this work, we propose Choreographer, an agent able to discover, learn and
adapt unsupervised skills efficiently by leveraging a generative model of the
environment dynamics, a.k.a. world model (Ha & Schmidhuber, 2018).
Choreographer discovers and learns skills in imagination, thanks to the model,
detaching the exploration and options discovery processes. During adaptation,
Choreographer can predict the outcomes of the learned skills’ actions, and so
evaluate multiple skills in parallel in imagination, allowing to combine them
efficiently for solving downstream tasks.

Contributions. Our contributions can be summarized as follow:

  * •

We describe a general algorithm for discovering, learning, and adapting
unsupervised skills that is exploration-agnostic and data-efficient. (Section
3).

  * •

We propose a _code resampling_ technique to prevent the issue of index
collapse when learning high-dimensional codes with vector quantized
autoencoders (Kaiser et al., 2018), which we employ in the skill discovery
process (Section 3.2);

  * •

We show that Choreographer can learn skills both from offline data or
collecting data with exploration, and from both states and pixels inputs. The
skills are adaptable for multiple tasks, as shown in the URL benchmark, where
we outperform all baselines (Section 4.1);

  * •

We show the skills learned by Choreographer are effective for exploration,
discovering sparse rewards in the environment more likely than other methods
(Section 4.2), and we further visualize and analyze them to provide additional
insights (Section 4.3).

##  2 Preliminaries and Background

Figure 1: Unsupervised Reinforcement Learning. The agent should effectively
leverage the unsupervised phase, consisting of the data collection and the
pre-training (PT) stages, to efficiently adapt during the supervised phase,
where the agent is fine-tuned (FT) for a downstream task.

Reinforcement Learning (RL).  The RL setting can be formalized as a Markov
Decision Process, where we denote observations from the environment with
$x_{t}$, actions with $a_{t}$, rewards with $r_{t}$, and the discount factor
with $\gamma$. The objective of an RL agent is to maximize the expected
discounted sum of rewards over time for a given task, a.k.a. the return:
${G_{t}=\sum_{k=t+1}^{T}\gamma^{(k-t-1)}r_{k}}$. We focus on continuous-
actions settings, where a common strategy is to learn a parameterized function
that outputs the best action given a certain state, referred to as the actor
$\pi_{\theta}(a|x)$, and a parameterized model that estimates the expected
returns from a state for the actor, referred to as the critic $v_{\psi}(x)$.
The models for actor-critic algorithms can be instantiated as deep neural
networks to solve complex continuous control tasks (Haarnoja et al., 2018;
Lillicrap et al., 2016; Schulman et al., 2017).

Unsupervised RL (URL).  In our method, we distinguish three stages of the URL
process (Figure 1). During stage _(i) data collection_ , the agent can
interact with the environment without rewards, driven by intrinsic motivation.
During stage _(ii) pre-training_ , the agent uses the data collected to learn
useful modules (e.g. a world model, skills, etc.) to adapt faster to new
scenarios. Crucially, stage _(i)_ and _(ii)_ need to happen in parallel for
most of the previous competence-based methods (Eysenbach et al., 2019; Liu &
Abbeel, 2021a; Gregor et al., 2016; Achiam et al., 2018), as skills
discriminability is tied to the data collection process. Finally, in stage
_(iii) fine-tuning_ , the agent can interact with the environment receiving
task reward feedback and should leverage the pre-trained modules to adapt
faster for the given task.

Mutual Information Skill Learning.  Option discovery can be formalized as
learning a skill representation $z$ to condition skill policies $\pi(a|x,z)$.
The objective of these policies is to maximize the mutual information between
skills and trajectories of environment observations
$\tau_{x}=\\{x_{t},x_{t+1},...\\}$:

$\displaystyle I(\tau_{x};z)$
$\displaystyle=H(\tau_{x})-H(\tau_{x}|z)\quad\quad\textrm{\emph{// forward}}$
(1) $\displaystyle=H(z)\phantom{.}-H(z|\tau_{x})\quad\quad\textrm{\emph{//
reverse}}$ (2)

The skill policy $\pi(a|x,z)$ results in a family of diversified policies with
predictable behavior. Previous work uses the forward form, either considering
the entropy $H(\tau_{x})$ fixed (Campos et al., 2020) or estimating it through
sampling-based techniques (Sharma et al., 2020; Liu & Abbeel, 2021a; Laskin et
al., 2022), and learning a function mapping skills to states to estimate
conditional entropy. Other works use the reverse form, considering a fixed
distribution over skills, with maximal entropy $H(z)$, and learning a skill
discriminator that maps states to skills, to estimate conditional entropy
(Achiam et al., 2018; Gregor et al., 2016; Eysenbach et al., 2019; Hansen et
al., 2020).

##  3 Choreographer

Choreographer is a model-based RL agent designed for URL, leveraging a world
model for discovering, learning, and adapting skills in imagination. In
addition to a world model, the agent learns: a) a codebook of skill vectors,
clustering the model state space, b) skill-conditioned policies, maximizing
the mutual information between latent trajectories of model states and skill
codes, and c) a meta-controller, to coordinate and adapt skills for downstream
tasks, during fine-tuning. An overview of the approach is illustrated in
Figure 2 and a detailed algorithm is presented in Appendix F.

###  3.1 World Model

Choreographer learns a task-agnostic world model composed of the following
components:

$\begin{split}\textrm{Posterior:}\quad
q_{\phi}(s_{t}|s_{t-1},a_{t-1},x_{t}),\end{split}\quad\quad\begin{split}\textrm{Prior:}\quad
p_{\phi}(s_{t}|s_{t-1},a_{t-1}),\end{split}\quad\quad\begin{split}\textrm{Reconstruction:}\quad
p_{\phi}(x_{t}|s_{t}).\end{split}$

During the fine-tuning stage, the agent also learns a reward predictor
$p_{\phi}(r_{t}|s_{t})$. The model states $s_{t}$ have both a deterministic
component, modeled using the recurrent state of a GRU (Cho et al., 2014), and
a discrete stochastic component (Hafner et al., 2021). The remaining modules
are multi-layer perceptrons (MLPs) or CNNs for pixel inputs (LeCun et al.,
1999). The model is trained end-to-end by optimizing an evidence lower bound
on the log-likelihood of data:

$\mathcal{L}_{\textrm{wm}}=D_{\mathrm{KL}}[q_{\phi}(s_{t}|s_{t-1},a_{t-1},x_{t})\|p_{\phi}(s_{t}|s_{t-1},a_{t-1})]-\mathbb{E}_{q_{\phi}(s_{t})}[\log
p_{\phi}(x_{t}|s_{t})],$ (3)

where sequences of observations $x_{t}$ and actions $a_{t}$ are sampled from a
replay buffer.

(a) Skill discovery.

(b) Skill learning.

(c) Skill adaptation.

Figure 2: Choreographer. (a) The agent leverages representation learning to
learn a codebook of skill vectors, summarizing the model state space. (b)
Skill policies are learned in imagination, maximizing the mutual information
between latent trajectories of model states and skill codes. (c) The meta-
controller selects the skills to apply for downstream tasks, evaluating and
adapting skill policies in imagination by hallucinating the trajectories of
states and rewards they would obtain.

###  3.2 Skill Discovery

To discover skills that are exploration-agnostic, i.e. do not require the
agent to deploy them in order to evaluate their discriminability, we turn the
skill discovery problem into a representation learning one. We learn a skill
discrete representation $z$ on top of the model states using a VQ-VAE (van den
Oord et al., 2017). This representation aims to encode model states into the
$N$ codes of a codebook, learning $p(z|s)$, as well as the inverse relation,
decoding codes to model states, via $p(s|z)$.

In the skill auto-encoder, model states are first processed into embeddings
using a skill encoder $E(s)$. Then, each embedding is assigned to the closest
code in the codebook $z_{q}^{(i)},i\in 1...N$, through:

$\mathrm{quantize}(E(s))=z_{q}^{(k)}\quad\quad
k=\arg\min_{j}\|E(s)-z^{(j)}_{q}\|_{2}\vspace{-0.4em}$ (4)

The indices are mapped back to their corresponding vectors in the codebook,
from which the decoder $D(z_{q}^{k})$ reconstructs the corresponding model
states. The training objective is:

$\mathcal{L}_{z}=\|s-D(z_{q}^{(k)})\|^{2}_{2}+\beta\|\textrm{sg}(z_{q}^{(k)})-E(s)\|^{2}_{2},$
(5)

where $\textrm{sg}(\cdot)$ indicates to stop gradients. The first loss term is
the reconstruction error of the decoder, while the second term encourages the
embedded vectors to stay close to the codes in the codebook.

Crucially, as opposed to most previous work (van den Oord et al., 2017; Razavi
et al., 2019), we choose to learn a single-code representation, where each
model state gets assigned only one vector in the codebook. Our skill codes can
thus be seen as centroids partitioning the model state space into $N$ clusters
(Roy et al., 2018), so that intuitively each skill occupies one area of the
state space.

Code resampling. When training single-code representations with VQ-VAE, the
codebook may suffer from _index collapse_ (Kaiser et al., 2018), where only a
few of the code vectors get trained due to a preferential attachment process.
One way to mitigate this issue is to use sliced latent vectors, which assign
multiple smaller codes to the same embedding (Kaiser et al., 2018; Razavi et
al., 2019), but this would imply losing the one-to-one association between
codes and skills.

In order to mitigate the issue we propose a _code resampling_ technique which
consists of two steps:

  1. 1.

keeping track of the inactive codes during training, which is achieved by
simply looking at the codes to which no embedding is assigned for $M$
consecutive training batches;

  2. 2.

re-initializing the inactive codes, with values of embeddings $E(s)$ from the
latest training batch, selected with a probability
$\frac{d_{q}^{2}(E(s))}{\sum_{s}d^{2}_{q}(E(s))}$, where $d_{q}(\cdot)$ is the
euclidean distance from the closest code in the codebook, i.e.
$d_{q}(E(s))=\min_{i\in 1...N}\|E(s)-z_{q}^{(i)}\|_{2}^{2}$.

The second step ensures that every $M$ training batches all the codes become
active again. In Section 4.3, we show that our technique is necessary for
training large discrete skill codebooks, as it keeps all codes active, and
leads to lower training loss. We provide additional details in Appendix C.

###  3.3 Skill Learning

Choreographer learns a generative world model, which allows the agent to
imagine trajectories of future model states
$\tau_{s}=\\{s_{t},s_{t+1},...\\}$. A major novelty of our approach is that we
learn a skill representation that is linked to model states rather than to
environment observations. As a consequence, for skill learning, we define the
objective of Choreographer to maximize the mutual information between
trajectories of model states and skills $I(\tau_{s},z)$. In visual
environments, this is particularly useful as the model states $s_{t}$ can have
a much more compact representation than observations $x_{t}$.

We rewrite the mutual information objective in forward form as:

$I(\tau_{s},z)=H(\tau_{s})-H(\tau_{s}|z)=H(\tau_{s})+\mathbb{E}_{s\sim
d_{\pi_{z}}(s),z\sim\mathrm{Unif}(z)}[\log p(s|z)],$ (6)

where the second term uses $d_{\pi_{z}}(s)$ to express the stationary
distribution of model states visited by the skill policy conditioned on the
skill code $z$, and codes are sampled from a uniform distribution. In order to
maximize this objective with RL, we estimate the two terms to be part of an
overarching reward function. The entropy term is estimated using a K-NN
particle-based estimator (Singh et al., 2003), obtaining
$r_{\textrm{ent}}\propto\sum_{i=1}^{K}\log\|s-s_{i}^{\textrm{K-NN}}\|_{2}$.
For the second term, we can rollout skill-conditioned policies in imagination,
obtaining sampled trajectories of model states, and use the distance between
the states visited and the decoded state corresponding to the conditioning
skill code $z$ to estimate the likelihood (Goodfellow et al., 2016), obtaining
$r_{\textrm{code}}\propto-\|D(z)-s\|_{2}^{2}$.

Choreographer learns $N$ code-conditioned skill policies using an actor-critic
algorithm:

$\begin{split}r_{\textrm{skill}}=r_{\textrm{ent}}+r_{\textrm{code}},\end{split}\quad\quad\begin{split}\textrm{Skill
actor:}\quad\pi_{\textrm{skill}}(a_{t}|s_{t},z),\end{split}\quad\quad\begin{split}\textrm{Skill
critic:}\quad v_{\textrm{skill}}(s_{t},z),\end{split}$ (7)

where the policies are encouraged to reach states corresponding to their
skill-conditioning code $z$, by the $r_{\textrm{code}}$ term, but at the same
time are encouraged to be dynamic and visit different states in one
trajectory, by the $r_{\textrm{ent}}$ term. Given that the reward terms can be
computed using the agent’s model, the skill actor-critics are trained in the
agent’s imagination, backpropagating through the model dynamics (Hafner et
al., 2021). By doing so, we are able to disentangle exploration and skill
learning, as the two processes can happen in parallel or at different times.
Visualizations of the skills learned by Choreographer are available in Figure
5 and in the videos on the project website.

###  3.4 Skill Adaptation

Choreographer is able to learn several skill-conditioned policies in a
completely unsupervised fashion, leveraging representation and reinforcement
learning. This feature is particularly useful during the pre-training stage of
unsupervised RL, where the agent learns both a task-agnostic world model of
the environment and a set of reusable behaviors. During fine-tuning, the agent
needs to leverage such pre-trained components in order to adapt to a
downstream task in an efficient way.

Exhaustively deploying the skill policies, a.k.a. the agent’s options, in the
environment, to understand which ones correlate best with the downstream task,
is expensive and often unnecessary. Choreographer can again exploit its world
model to deploy the skill policies in imagination.

In order to coordinate the multiple options available, Choreographer learns a
_meta-controller_ , a policy over options (Precup, 2000) deciding which skill
codes use to condition the skill policy:

$\begin{split}r_{\textrm{meta}}=r_{\textrm{task}},\end{split}\quad\quad\begin{split}&\textrm{Meta-
ctrl
actor:}\quad\pi_{\textrm{meta}}(z|s_{t}),\end{split}\quad\quad\begin{split}&\textrm{Meta-
ctrl critic:}\quad v_{\textrm{meta}}(s_{t}).\end{split}$ (8)

The meta-controller is updated with rollouts in imagination using the task
rewards from a reward predictor, trained only at fine-tuning time. The outputs
of the meta-controller’s actor network are discrete (skill) vectors and we use
policy gradients to update the actor (Williams, 1992). In order to fine-tune
the pre-trained skill policies, we can backpropagate the gradients to the
skill actors as well (Kingma & Welling, 2013). As we discuss in Section 4.3,
this improves performance further as the pre-trained skills are not supposed
to perfectly fit the downstream task beforehand.

##  4 Experiments

We empirically evaluate Choreographer to answer the following questions:

  * •

Does Choreographer learn and adapt skills effectively for unsupervised RL?
(Section 4.1) We use the URL benchmark (Laskin et al., 2021) to show that,
after pre-training on exploration data, our agent can adapt to several tasks
in a data-efficient manner. We show this holds for both training on states and
pixels inputs and in both offline and online settings.

  * •

Do the skills learned by Choreographer provide a more exploratory
initialization? (Section 4.2) We zoom in on the Jaco sparse tasks from URLB
and use sparse goal-reaching tasks from Meta-World (Yu et al., 2019) to
evaluate the ability of our agent to find sparse rewards in the environment,
by leveraging its skills.

  * •

Which skills are learned by Choreographer? (Section 4.3) We visualize the
diverse skills learned, highlighting the effectiveness of the proposed
resampling technique. Finally, we discuss the ‘zero-shot’ adaptation
performance achievable when fine-tuning only the meta-controller, with no
gradients flowing to the skill policies.

Implementation details and hyperparameters are provided in Appendix B111The
PyTorch code for Choreographer and all the experiments will be open-sourced
upon publication..

Baselines. For state-based experiments, we adopt the baselines from the CIC
paper (Laskin et al., 2022), which can be grouped in three categories:
knowledge-based (blue in plots), which are ICM (Pathak et al., 2017), RND
(Burda et al., 2019), Disagreement (Pathak et al., 2019); data-based ( green
in plots), which are APT (Liu & Abbeel, 2021b) and ProtoRL (Yarats et al.,
2021); and competence-based (red/orange in plots) APS (Liu & Abbeel, 2021a),
SMM (Lee et al., 2019) and DIAYN (Eysenbach et al., 2019). For pixel-based
experiments, we compare to the strong baselines from (Rajeswar et al., 2022),
combining some of the above approaches with the Dreamer algorithm.

###  4.1 URL Benchmark

Figure 3: Offline state-based URLB. On the left, Choreographer (ours), pre-
trained with offline exploratory data, performs best against baselines, both
in terms of IQM and Optimality Gap. On the right, mean and standard deviations
across the different domains of URLB are detailed (10 seeds).

Offline data. Choreographer decouples the exploration and skill learning
processes, and can thus learn from offline pre-collected data. To show this,
we use the pre-collected datasets of exploratory data released in the ExORL
paper (Yarats et al., 2022). We take the first 2M steps from the datasets
corresponding to the state-based URLB domains and use these to pre-train
Choreographer offline, learning the world model and the skill-related
components. In Figure 3, we compare the fine-tuning performance of our method,
leveraging the meta-controller at fine-tuning time, to other methods. In
particular, we show the performance when pre-training Choreographer using the
data collected by the RND approach (Burda et al., 2019). In Appendix D, we
show results with different datasets.

Figure 4: Pixel-based URLB. Performance of Choreographer (ours) as a function
of pre-training steps. Scores are asymptotically normalized, and averaged
across tasks for each method. (10 seeds)

Following (Agarwal et al., 2021), we use the interquartile mean (IQM) and
optimality gap metrics, aggregated with stratified bootstrap sampling, as our
main evaluation metrics. We find that Choreographer overall performs best on
both metrics, particularly outperforming the other approaches on the Walker
and Quadruped tasks and performing comparably to CIC and ProtoRL in the Jaco
tasks. We highlight that Choreographer performs far better than the other
competence-based approaches, SMM, APS, and DIAYN, and also strongly
outperforms RND, which is the strategy used to collect its data, showing it’s
a stronger method to exploit the same exploratory data.

Parallel exploration. With no pre-collected data available, Choreographer can
learn an exploration policy in parallel to the skills (Mendonca et al., 2021),
to collect data and pre-train its components. This idea is detailed in
Appendix A. We use the pixel-based setup of URLB to see how the performance
evolves over time as the size of the pre-training data buffer grows.

In Figure 4, we show the performance of Choreographer, collecting data with an
exploration policy based on the LBS exploration approach (Mazzaglia et al.,
2021) and learning skills simultaneously. The agents are being tested at
different times during learning, taking snapshots of the agent at 100k, 500k,
1M, and 2M steps and fine-tuning them for 100k steps. Additional baselines
from (Rajeswar et al., 2022), are presented in Appendix D.

Overall, Choreographer’s performance is steadily on top, with similar
performance to APT, LBS, and ICM until 500k steps, and becoming the one best-
performing algorithm from 1M steps on. As also noticed in other works (Laskin
et al., 2021), other competence-based approaches APS and DIAYN struggle on the
benchmark. Furthermore, both LBS and APT’s performance in Jaco starts decaying
after 500k steps. As also discussed in (Rajeswar et al., 2022), this is due to
the bad initialization of the actor provided by the exploration approach. Our
approach overcomes this limitation by learning a set of diverse policies,
consistently providing a more general initialization that is easier to adapt.

Domain | Task | Choreo | Random | APS | DIAYN | APT | ICM | LBS | P2E | RND
---|---|---|---|---|---|---|---|---|---|---
Jaco | Bottom left | 1.0 | 1.0 | 0.67 | 0.67 | 0.67 | 0.33 | 0.0 | 0.67 | 0.67
Bottom right | 1.0 | 1.0 | 0.67 | 0.67 | 0.33 | 0.67 | 0.0 | 1.0 | 0.33
Top left | 1.0 | 1.0 | 0.67 | 0.67 | 0.0 | 0.67 | 0.33 | 0.33 | 0.0
Top right | 1.0 | 1.0 | 0.0 | 0.67 | 0.33 | 0.0 | 0.33 | 0.67 | 0.33
MetaWorld | Goal set 1 | 0.78 | 0.47 | 0.31 | 0.08 | 0.24 | 0.33 | 0.33 | 0.51 | 0.35
Goal set 2 | 0.80 | 0.49 | 0.18 | 0.02 | 0.16 | 0.37 | 0.18 | 0.41 | 0.43
Goal set 3 | 0.84 | 0.51 | 0.12 | 0.04 | 0.18 | 0.18 | 0.61 | 0.37 | 0.22

Table 1: Sparse rewards. Finding goals in the environment with sparse rewards
using the pre-trained behaviors. The table shows the fraction of goals found
by each method per task. (3 seeds)

###  4.2 Sparse Rewards

Finding sparse rewards in the environment is challenging, though pre-trained
exploratory behaviors may help (Parisi et al., 2021). Choreographer can
exploit the skill policies learned during pre-training to reach different
model states, efficiently skimming the environment to find rewards. We test
this hypothesis using sparse goal-reaching tasks in two pixel-based
environments: Jaco from the DMC Suite, and MetaWorld (Yu et al., 2019). In
order to solve these tasks, the agent needs to reach within a certain radius
of a target goal located in the environment with (nearly) no feedback.

Jaco. We zoom in on the performance on the four reaching tasks from URLB. In
Table 1, we show the fraction of runs reaching near the target, without any
update of the pre-trained policies, over 100 evaluation episodes. Only
Choreographer and random actions found the goals in all runs, the other
methods have to rely on entropic exploration to find the goals222For the
results in Figure 4 all the baselines (except Choreographer) had an action
entropy maximization term in the fine-tuning objective (Haarnoja et al.,
2018), helping them to find the goals.. The other skill-based methods, APS and
DIAYN, also struggled to find the targets consistently.

Meta-World. We sampled three sets of 50 goals, in the Sawyer environment from
MetaWorld. For each goal in the sets, we evaluate the policies pre-trained by
different methods, testing them for 100 evaluation episodes, without any fine-
tuning. The performance in Table 1 shows that Choreographer finds the most
goals in the environment. All the other approaches struggled in the task, with
random actions being the best alternative for higher chances of finding sparse
rewards for these goals.

###  4.3 Skills Analysis

Figure 5: Skills visualization. Highlighted in  green, reconstructions of the
skill codes conditioning the skill behavior on their right, sampling frames
over time. Animated results on the project website.

(a) Reconstruction of 64 codes, no code resampling.
Inactive, collapsed codes highlighted in  red.

Figure 6: Code resampling (CR). When using more than 32 codes, CR reduces the
number of unused codes ($\sim$0% after 10k steps) and leads to a lower
reconstruction loss. (3 seeds $\times$ 3 domains)

Skill visualization. In Figure 5, we visualize some of the skills learned by
Choreographer, in the pixel-based Walker and Quadruped environments. The codes
conditioning the skill policies are shown by first reconstructing the
corresponding model state, and next the image observation using the
reconstruction model. We note that Choreographer’s learned behaviors
accurately reflect the skill learning objective (Eq. 7): they tend to reach
the states corresponding to the conditioning skill codes, but they also
dynamically visit different states. For the Walker, this leads to walking-like
behaviors for standing-up skill codes (1st row), or behaviors that move
around, by hopping or flipping on the torso/legs (2nd row). For the Quadruped,
skills tend to learn tipping-over behaviors, where the agent reaches the skill
code state alternating rotating and overturning moves (3rd/4th rows). Further
visualizations are in Appendix E and on the project website.

Code resampling. Learning a diversified set of skill codes is important to
learn useful skill policies for adaptation. Because of index collapse (Kaiser
et al., 2018), VQ-VAEs tend to ignore some of the skill codes, choosing not to
use them during training. As Choreographer samples skills uniformly for skill
learning (Eq. 6), it would be inefficient to train policies that are
conditioned on inactive codes. As we show in Figure 6 in the pixel-based URLB
environments, code resampling circumvents this issue, making nearly all the
codes active during training, reducing the reconstruction loss, and avoiding
skill code collapse (see examples of collapsed codes in the Figure).

Zero-shot skill control. Pre-training skill policies is a general approach to
provide a better initialization of the agent’s behavior for fine-tuning to
downstream tasks (Eysenbach et al., 2021). As shown in (Eysenbach et al.,
2021), it is not possible to learn a set of skill policies able to solve any
task, without adaptation. We empirically reinforce this hypothesis with an
ablation on the pixel-based URLB, where we learn the meta-controller, which
interpolates the skill policies to solve the task, but we do not fine-tune the
skill policies. For space constraints, these results, showing a consistent gap
in performance versus the skill-adapting counterpart, are presented in
Appendix D. Combining multiple skills to solve complex long-term tasks with
little to no adaptation, similarly to hierarchical RL (Gehring et al., 2021;
Hafner et al., 2022), is a challenge we aim to address in future work.

##  5 Related Work

Unsupervised RL. Among URL methods, we can distinguish (Srinivas & Abbeel,
2021): knowledge-based methods, which aim to visit states of the environment
that maximize the agent’s model error/uncertainty (Schmidhuber, 1991; Pathak
et al., 2017; Burda et al., 2019), data-based methods, which aim to maximize
the entropy of the data collected (Hazan et al., 2019; Yarats et al., 2021;
Liu & Abbeel, 2021b), competence-based methods, which learn a set of options
or skills that allow the agent to control the environment (Gregor et al.,
2016; Eysenbach et al., 2019; Hansen et al., 2021). Knowledge and data-based
methods focus on expanding the available information, while competence-based
methods, such as Choreographer, aim to provide empowerment.

Exploration and skills. Exploration is a major challenge in skill learning
approaches. Previous work has attempted to solve this issue by using state
maximum entropy (Laskin et al., 2022; Liu & Abbeel, 2021a) or introducing
variations of the skill discovery objective (Kim et al., 2021; Park et al.,
2022). The idea of decoupling exploration and skill discovery has been less
explored and mostly applied in simpler environments (Campos et al., 2020;
Nieto et al., 2021). To the best of our knowledge, Choreographer represents
the first approach that enables scaling the idea to complex environments and
tasks successfully.

Behavior learning in imagination. The Dreamer algorithm (Hafner et al., 2019,
2021) introduced behavior learning in imagination, establishing a modern
implementation of the Dyna framework (Sutton, 1991). The idea has been applied
for exploration (Sekar et al., 2020; Rajeswar et al., 2022; Mazzaglia et al.,
2021), goal-directed behavior (Hafner et al., 2022; Mendonca et al., 2021),
and planning amortization Xie et al. (2020). In particular, Rajeswar et al.
(2022) showed the effectiveness of fine-tuning behavior in imagination on the
URL benchmark from pixels, combining unsupervised RL for exploration with
Dreamer. Choreographer not only yields better performance but also decouples
exploration and behavior pre-training, as the skill policies can be learned
offline.

##  6 Discussion

Competence-based approaches aim to learn a set of skills that allow adapting
faster to downstream tasks. Previous approaches have struggled to accomplish
this goal, mainly due to the impaired exploration capabilities of these
algorithms, and the impossibility of efficiently leveraging all the learned
skills during adaptation time. In this work, we introduced Choreographer, a
model-based agent for unsupervised RL, which is able to discover, learn and
adapt skills in imagination, without the need to deploy them in the
environment. The approach is thus exploration-agnostic and allows learning
from offline data or exploiting strong exploration strategies.

In order to decouple exploration and skill discovery, we learn a codebook of
skills, intuitively clustering the agent’s model state space into different
areas that should be pursued by the skill policies. While this proved to be
successful in several domains, it might not be the most effective skill
representation, and alternative ways for learning the skills could be
explored, such as identifying crucial states in the environment (Faccio et
al., 2022).

Choreographer can evaluate skill policies in imagination during the fine-
tuning stage and use a meta-controller to decide which policies should be
employed and adapted to the downstream task. While our method exploits a world
model for data-efficient adaptation, alternative strategies such as successor
representations (Hansen et al., 2020) or meta RL (Zintgraf et al., 2019) could
be considered.

The meta-controller allows evaluating all the skill policies in imagination
and fine-tuning them for downstream tasks efficiently, but currently doesn’t
allow combining the skills for longer-horizon complex tasks. We aim to further
investigate this in the future, leveraging hierarchical RL (Hafner et al.,
2022) and model learning, as established in the options framework (Precup,
2000).

Limitations. Choreographer is a data-driven approach to skill discovery and
learning. The skill representation depends on the data provided or collected
from the environment and on the world model learned. As a consequence, the
skills learned by Choreographer tend to reach multiple areas of the
environment and assume different behaviors, but struggle to capture the
naturalness that humans develop in their skills. We leave the idea of
incorporating human bias in the skill learning process for future work.

#### Reproducibility statement

We reported in appendix the detailed pseudo-code for both Choreographer
(Algorithm 2) and code-resampling (Algorithm 1). All the instructions and
training hyperparameters to implement and reproduce the results are in
Appendix B, along with a summary of other experiments we tried that could be
useful for future work. We will release our code and scripts.

## Acknowledgements

P.M., T.V., and B.D. received funding from the Flemish Government under the
“Onderzoeksprogramma Artificiële Intelligentie (AI) Vlaanderen” programme.

## References

  * Achiam et al. (2018) Joshua Achiam, Harrison Edwards, Dario Amodei, and Pieter Abbeel.  Variational option discovery algorithms.  _ArXiv_ , abs/1807.10299, 2018.
  * Agarwal et al. (2021) Rishabh Agarwal, Max Schwarzer, Pablo Samuel Castro, Aaron C Courville, and Marc Bellemare.  Deep reinforcement learning at the edge of the statistical precipice.  _Advances in Neural Information Processing Systems_ , 34, 2021.
  * Arthur & Vassilvitskii (2006) David Arthur and Sergei Vassilvitskii.  k-means++: The advantages of careful seeding.  Technical report, Stanford, 2006.
  * Bellemare et al. (2016) Marc Bellemare, Sriram Srinivasan, Georg Ostrovski, Tom Schaul, David Saxton, and Remi Munos.  Unifying count-based exploration and intrinsic motivation.  In D. Lee, M. Sugiyama, U. Luxburg, I. Guyon, and R. Garnett (eds.), _Advances in Neural Information Processing Systems_ , volume 29, 2016.
  * Bengio et al. (2013) Yoshua Bengio, Nicholas Léonard, and Aaron Courville.  Estimating or propagating gradients through stochastic neurons for conditional computation, 2013.  URL https://arxiv.org/abs/1308.3432.
  * Burda et al. (2019) Yuri Burda, Harrison Edwards, Amos J. Storkey, and Oleg Klimov.  Exploration by random network distillation.  _ICLR_ , 2019.
  * Campos et al. (2020) Víctor Campos, Alexander Trott, Caiming Xiong, Richard Socher, Xavier Giro i Nieto, and Jordi Torres.  Explore, discover and learn: Unsupervised discovery of state-covering skills.  In _ICML_ , 2020.
  * Cho et al. (2014) Kyunghyun Cho, Bart van Merrienboer, Dzmitry Bahdanau, and Yoshua Bengio.  On the properties of neural machine translation: Encoder–decoder approaches.  In _SSST@EMNLP_ , 2014.
  * Corbetta (2021) Daniela Corbetta.  Perception, action, and intrinsic motivation in infants’ motor-skill development.  _Current Directions in Psychological Science_ , 30(5):418–424, 2021.  doi: 10.1177/09637214211031939.  URL https://doi.org/10.1177/09637214211031939.
  * Eysenbach et al. (2019) Benjamin Eysenbach, Abhishek Gupta, Julian Ibarz, and Sergey Levine.  Diversity is all you need: Learning skills without a reward function.  In _ICLR_ , 2019.
  * Eysenbach et al. (2021) Benjamin Eysenbach, Ruslan Salakhutdinov, and Sergey Levine.  The information geometry of unsupervised reinforcement learning, 2021\.  URL https://arxiv.org/abs/2110.02719.
  * Faccio et al. (2022) Francesco Faccio, Aditya Ramesh, Vincent Herrmann, Jean Harb, and Jürgen Schmidhuber.  General policy evaluation and improvement by learning to identify few but crucial states, 2022.  URL https://arxiv.org/abs/2207.01566.
  * Gehring et al. (2021) Jonas Gehring, Gabriel Synnaeve, Andreas Krause, and Nicolas Usunier.  Hierarchical skills for efficient exploration, 2021.  URL https://arxiv.org/abs/2110.10809.
  * Goodfellow et al. (2016) Ian Goodfellow, Yoshua Bengio, and Aaron Courville.  _Deep learning_.  MIT press, 2016.
  * Gregor et al. (2016) Karol Gregor, Danilo Jimenez Rezende, and Daan Wierstra.  Variational intrinsic control, 2016.  URL https://arxiv.org/abs/1611.07507.
  * Ha & Schmidhuber (2018) David Ha and Jürgen Schmidhuber.  Recurrent world models facilitate policy evolution.  In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnett (eds.), _Advances in Neural Information Processing Systems_. Curran Associates, Inc., 2018.
  * Haarnoja et al. (2018) Tuomas Haarnoja, Aurick Zhou, P. Abbeel, and Sergey Levine.  Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor.  In _ICML_ , 2018.
  * Hafner et al. (2019) Danijar Hafner, Timothy Lillicrap, Jimmy Ba, and Mohammad Norouzi.  Dream to control: Learning behaviors by latent imagination.  In _ICLR_ , 2019.
  * Hafner et al. (2021) Danijar Hafner, Timothy P Lillicrap, Mohammad Norouzi, and Jimmy Ba.  Mastering atari with discrete world models.  In _ICLR_ , 2021.
  * Hafner et al. (2022) Danijar Hafner, Kuang-Huei Lee, Ian Fischer, and Pieter Abbeel.  Deep hierarchical planning from pixels, 2022.  URL https://arxiv.org/abs/2206.04114.
  * Hansen et al. (2020) Steven Hansen, Will Dabney, Andre Barreto, David Warde-Farley, Tom Van de Wiele, and Volodymyr Mnih.  Fast task inference with variational intrinsic successor features.  In _International Conference on Learning Representations_ , 2020.
  * Hansen et al. (2021) Steven Stenberg Hansen, Guillaume Desjardins, Kate Baumli, David Warde-Farley, Nicolas Heess, Simon Osindero, and Volodymyr Mnih.  Entropic desired dynamics for intrinsic control.  In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan (eds.), _Advances in Neural Information Processing Systems_ , 2021.
  * Hazan et al. (2019) Elad Hazan, Sham M. Kakade, Karan Singh, and Abby Van Soest.  Provably efficient maximum entropy exploration.  _ArXiv_ , abs/1812.02690, 2019.
  * Kaiser et al. (2018) Łukasz Kaiser, Aurko Roy, Ashish Vaswani, Niki Parmar, Samy Bengio, Jakob Uszkoreit, and Noam Shazeer.  Fast decoding in sequence models using discrete latent variables, 2018\.  URL https://arxiv.org/abs/1803.03382.
  * Kim et al. (2021) Jaekyeom Kim, Seohong Park, and Gunhee Kim.  Unsupervised skill discovery with bottleneck option learning.  In Marina Meila and Tong Zhang (eds.), _Proceedings of the 38th International Conference on Machine Learning_ , volume 139 of _Proceedings of Machine Learning Research_ , pp. 5572–5582. PMLR, 18–24 Jul 2021.  URL https://proceedings.mlr.press/v139/kim21j.html.
  * Kingma & Welling (2013) Diederik P Kingma and Max Welling.  Auto-encoding variational bayes, 2013.  URL https://arxiv.org/abs/1312.6114.
  * Laskin et al. (2021) Michael Laskin, Denis Yarats, Hao Liu, Kimin Lee, Albert Zhan, Kevin Lu, Catherine Cang, Lerrel Pinto, and Pieter Abbeel.  Urlb: Unsupervised reinforcement learning benchmark, 2021.  URL https://arxiv.org/abs/2110.15191.
  * Laskin et al. (2022) Michael Laskin, Hao Liu, Xue Bin Peng, Denis Yarats, Aravind Rajeswaran, and Pieter Abbeel.  Cic: Contrastive intrinsic control for unsupervised skill discovery, 2022\.  URL https://arxiv.org/abs/2202.00161.
  * LeCun et al. (1999) Yann LeCun, Patrick Haffner, Léon Bottou, and Yoshua Bengio.  Object recognition with gradient-based learning.  In _Shape, contour and grouping in computer vision_ , pp. 319–345. Springer, 1999.
  * Lee et al. (2019) Lisa Lee, Benjamin Eysenbach, Emilio Parisotto, Eric Xing, Sergey Levine, and Ruslan Salakhutdinov.  Efficient exploration via state marginal matching, 2019.  URL https://arxiv.org/abs/1906.05274.
  * Lillicrap et al. (2016) Timothy P. Lillicrap, Jonathan J. Hunt, Alexander Pritzel, Nicolas Heess, Tom Erez, Yuval Tassa, David Silver, and Daan Wierstra.  Continuous control with deep reinforcement learning.  In Yoshua Bengio and Yann LeCun (eds.), _ICLR_ , 2016.
  * Liu & Abbeel (2021a) Hao Liu and P. Abbeel.  Aps: Active pretraining with successor features.  In _ICML_ , 2021a.
  * Liu & Abbeel (2021b) Hao Liu and Pieter Abbeel.  Behavior from the void: Unsupervised active pre-training.  In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan (eds.), _Advances in Neural Information Processing Systems_ , volume 34, pp. 18459–18473. Curran Associates, Inc., 2021b.
  * Marcinowski et al. (2019) E. C. Marcinowski, E. Nelson, J. M. Campbell, and G. F. Michel.  The Development of Object Construction from Infancy through Toddlerhood.  _Infancy_ , 24(3):368–391, 2019.
  * Mazzaglia et al. (2021) Pietro Mazzaglia, Ozan Çatal, Tim Verbelen, and B. Dhoedt.  Curiosity-driven exploration via latent bayesian surprise.  _ArXiv_ , abs/2104.07495, 2021.
  * Mendonca et al. (2021) Russell Mendonca, Oleh Rybkin, Kostas Daniilidis, Danijar Hafner, and Deepak Pathak.  Discovering and achieving goals via world models.  In _Advances in Neural Information Processing Systems_ , 2021.
  * Mnih et al. (2013) Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Alex Graves, Ioannis Antonoglou, Daan Wierstra, and Martin A. Riedmiller.  Playing atari with deep reinforcement learning.  _CoRR_ , 2013.
  * Mutti et al. (2020) Mirco Mutti, Lorenzo Pratissoli, and Marcello Restelli.  Task-agnostic exploration via policy gradient of a non-parametric state entropy estimate, 2020.  URL https://arxiv.org/abs/2007.04640.
  * Nieto et al. (2021) Juan José Nieto, Roger Creus, and Xavier Giro-i Nieto.  Unsupervised skill-discovery and skill-learning in minecraft, 2021.  URL https://arxiv.org/abs/2107.08398.
  * OpenAI et al. (2019) OpenAI, Ilge Akkaya, Marcin Andrychowicz, Maciek Chociej, Mateusz Litwin, Bob McGrew, Arthur Petron, Alex Paino, Matthias Plappert, Glenn Powell, Raphael Ribas, Jonas Schneider, Nikolas Tezak, Jerry Tworek, Peter Welinder, Lilian Weng, Qiming Yuan, Wojciech Zaremba, and Lei Zhang.  Solving rubik’s cube with a robot hand, 2019.  URL https://arxiv.org/abs/1910.07113.
  * Parisi et al. (2021) Simone Parisi, Victoria Dean, Deepak Pathak, and Abhinav Gupta.  Interesting object, curious agent: Learning task-agnostic exploration, 2021.  URL https://arxiv.org/abs/2111.13119.
  * Park et al. (2022) Seohong Park, Jongwook Choi, Jaekyeom Kim, Honglak Lee, and Gunhee Kim.  Lipschitz-constrained unsupervised skill discovery, 2022.  URL https://arxiv.org/abs/2202.00914.
  * Pathak et al. (2017) Deepak Pathak, Pulkit Agrawal, Alexei Efros, and Trevor Darrell.  Curiosity-driven exploration by self-supervised prediction.  _ICML_ , 2017.
  * Pathak et al. (2019) Deepak Pathak, Dhiraj Gandhi, and Abhinav Gupta.  Self-supervised exploration via disagreement.  In _ICML_ , 2019.
  * Precup (2000) Doina Precup.  _Temporal abstraction in reinforcement learning_.  University of Massachusetts Amherst, 2000.
  * Rajeswar et al. (2022) Sai Rajeswar, Pietro Mazzaglia, Tim Verbelen, Alexandre Piché, Bart Dhoedt, Aaron Courville, and Alexandre Lacoste.  Unsupervised model-based pre-training for data-efficient control from pixels, 2022.  URL https://arxiv.org/abs/2209.12016.
  * Razavi et al. (2019) Ali Razavi, Aaron van den Oord, and Oriol Vinyals.  Generating diverse high-fidelity images with vq-vae-2, 2019.  URL https://arxiv.org/abs/1906.00446.
  * Roy et al. (2018) Aurko Roy, Ashish Vaswani, Arvind Neelakantan, and Niki Parmar.  Theory and experiments on vector quantized autoencoders, 2018.  URL https://arxiv.org/abs/1805.11063.
  * Salge et al. (2014) Christoph Salge, Cornelius Glackin, and Daniel Polani.  Changing the environment based on empowerment as intrinsic motivation.  _Entropy_ , 16(5):2789–2819, may 2014.  doi: 10.3390/e16052789.  URL https://doi.org/10.3390%2Fe16052789.
  * Schaul et al. (2015) Tom Schaul, Daniel Horgan, Karol Gregor, and David Silver.  Universal value function approximators.  In Francis Bach and David Blei (eds.), _Proceedings of the 32nd International Conference on Machine Learning_ , volume 37 of _Proceedings of Machine Learning Research_ , pp. 1312–1320, Lille, France, 07–09 Jul 2015\. PMLR.
  * Schmidhuber (1991) J. Schmidhuber.  Curious model-building control systems.  In _[Proceedings] 1991 IEEE International Joint Conference on Neural Networks_ , pp. 1458–1463 vol.2, 1991.  doi: 10.1109/IJCNN.1991.170605.
  * Schulman et al. (2017) John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov.  Proximal policy optimization algorithms.  _ArXiv_ , abs/1707.06347, 2017.
  * Sekar et al. (2020) Ramanan Sekar, Oleh Rybkin, Kostas Daniilidis, Pieter Abbeel, Danijar Hafner, and Deepak Pathak.  Planning to explore via self-supervised world models.  In _ICML_ , 2020.
  * Sharma et al. (2020) Archit Sharma, Shixiang Gu, Sergey Levine, Vikash Kumar, and Karol Hausman.  Dynamics-aware unsupervised discovery of skills.  In _International Conference on Learning Representations_ , 2020.
  * Silver et al. (2016) David Silver, Aja Huang, Chris J Maddison, Arthur Guez, Laurent Sifre, George Van Den Driessche, Julian Schrittwieser, Ioannis Antonoglou, Veda Panneershelvam, and Marc Lanctot.  Mastering the game of go with deep neural networks and tree search.  _nature_ , 2016.
  * Singh et al. (2003) Harshinder Singh, Neeraj Misra, Vladimir Hnizdo, Adam Fedorowicz, and Eugene Demchuk.  Nearest neighbor estimates of entropy.  _American journal of mathematical and management sciences_ , 23(3-4):301–321, 2003.
  * Smith et al. (2022) Laura Smith, Ilya Kostrikov, and Sergey Levine.  A walk in the park: Learning to walk in 20 minutes with model-free reinforcement learning, 2022.  URL https://arxiv.org/abs/2208.07860.
  * Srinivas & Abbeel (2021) Aravind Srinivas and Pieter Abbeel.  Unsupervised learning for reinforcement learning.  2021\.  URL https://icml.cc/media/icml-2021/Slides/10843_QHaHBNU.pdf.
  * Sutton (1991) Richard S. Sutton.  Dyna, an integrated architecture for learning, planning, and reacting.  _SIGART Bull._ , 2(4):160–163, jul 1991.  ISSN 0163-5719.  doi: 10.1145/122344.122377.  URL https://doi.org/10.1145/122344.122377.
  * Tassa et al. (2018) Yuval Tassa, Yotam Doron, Alistair Muldal, Tom Erez, Yazhe Li, Diego de Las Casas, David Budden, Abbas Abdolmaleki, Josh Merel, Andrew Lefrancq, Timothy P. Lillicrap, and Martin A. Riedmiller.  Deepmind control suite.  _CoRR_ , abs/1801.00690, 2018.
  * van den Oord et al. (2017) Aaron van den Oord, Oriol Vinyals, and koray kavukcuoglu.  Neural discrete representation learning.  In I. Guyon, U. Von Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett (eds.), _Advances in Neural Information Processing Systems_ , volume 30. Curran Associates, Inc., 2017.  URL https://proceedings.neurips.cc/paper/2017/file/7a98af17e63a0ac09ce2e96d03992fbc-Paper.pdf.
  * Williams (1992) Ronald J. Williams.  Simple statistical gradient-following algorithms for connectionist reinforcement learning.  _Machine Learning_ , 8(3):229–256, May 1992.  ISSN 1573-0565.  doi: 10.1007/BF00992696.  URL https://doi.org/10.1007/BF00992696.
  * Xie et al. (2020) Kevin Xie, Homanga Bharadhwaj, Danijar Hafner, Animesh Garg, and Florian Shkurti.  Latent skill planning for exploration and transfer, 2020.  URL https://arxiv.org/abs/2011.13897.
  * Yarats et al. (2021) Denis Yarats, Rob Fergus, Alessandro Lazaric, and Lerrel Pinto.  Reinforcement learning with prototypical representations.  2021\.
  * Yarats et al. (2022) Denis Yarats, David Brandfonbrener, Hao Liu, Michael Laskin, P. Abbeel, Alessandro Lazaric, and Lerrel Pinto.  Don’t change the algorithm, change the data: Exploratory data for offline reinforcement learning.  _ArXiv_ , abs/2201.13425, 2022.
  * Yu et al. (2019) Tianhe Yu, Deirdre Quillen, Zhanpeng He, Ryan Julian, Karol Hausman, Chelsea Finn, and Sergey Levine.  Meta-world: A benchmark and evaluation for multi-task and meta reinforcement learning.  In _Conference on Robot Learning (CoRL)_ , 2019.  URL https://arxiv.org/abs/1910.10897.
  * Zintgraf et al. (2019) Luisa Zintgraf, Kyriacos Shiarlis, Maximilian Igl, Sebastian Schulze, Yarin Gal, Katja Hofmann, and Shimon Whiteson.  Varibad: A very good method for bayes-adaptive deep rl via meta-learning, 2019.  URL https://arxiv.org/abs/1910.08348.

## Appendix

##  Appendix A Data Collection

In previous skill discovery methods (Eysenbach et al., 2019; Liu & Abbeel,
2021a; Achiam et al., 2018), skills are primarily used as a means to discover
the environment by maximizing the diversity of skills. In contrast,
Choreographer’s skill discovery is exploration-agnostic, in that it discovers
and learns skills using its model. For this reason, Choreographer can be
adopted in two ways:

  * •

using a pre-collected dataset used to learn the world model and skill-related
components;

  * •

using an exploration policy, which aims to collect diverse data to pre-train.

When not using pre-collected data, we learn an exploration actor-critic in
imagination that maximizes intrinsic rewards (Mendonca et al., 2021; Sekar et
al., 2020). In particular, we shape rewards following LBS (Mazzaglia et al.,
2021), which has proven to be a reliable exploration method when using world
models (Rajeswar et al., 2022). Exploration rewards approximate the
information gain between observations and model states, computed by measuring
the KL divergence between the posterior and the prior of the model. The
exploration actor-critic is defined as follows:

$\begin{split}\textrm{Expl.\
actor:}\quad\pi_{\textrm{expl}}(a_{t}|s_{t}),\end{split}\quad\quad\begin{split}\textrm{Expl.\
critic:}\quad v_{\textrm{expl}}(s_{t}),\end{split}$ (9)
$r_{\textrm{expl}}=D_{\mathrm{KL}}[q_{\phi}(s_{t}|s_{t-1},a_{t-1},x_{t})\|p_{\phi}(s_{t}|s_{t-1},a_{t-1})],$

##  Appendix B Implementation

###  B.1 Agent

World Model. We follow the architecture of the DreamerV2 agent’s world model
(Hafner et al., 2021). Model states have both a deterministic and a stochastic
component: the deterministic component is the 200-dimensional output of a GRU
(Cho et al., 2014), with a 200-dimensional hidden layer; the stochastic
component consists of 32 categorical distributions with 32 classes each. For
states-input experiments, the encoder and decoder are 4-layers MLP with a
dimensionality of 400. For pixels-based experiments, the encoder and decoder
follow the architecture of DreamerV2, taking $64\times 64$ RGB images as
inputs.

Skill auto-encoder. The encoder and decoder networks are 4-layer MLP with a
dimensionality of 400. The input and target of the auto-encoder are the
deterministic part of the world model states. The codebook consists of 64
codes of dimensionality 16. The gradients of the skill auto-encoder objective
are back-propagated through the decoder and the encoder using straight-through
gradients (Bengio et al., 2013). The codes of the codebooks are updated using
an exponential moving average of their assigned embeddings (Razavi et al.,
2019). Code resampling happens every $M=200$ training batches.

Both the world model and the skill-autoencoder networks are updated sampling
batches of 50 sequences of 50 timesteps, using Adam with learning rate
$3e^{-4}$ for the updates, and clipping gradients norm to 100.

Skill actor-critic. Skill actor-critic takes as input both model states and
skill vectors (Schaul et al., 2015). We use the Dreamer algorithm to update
them in latent imagination (Hafner et al., 2019), imagining sequences of 15
steps with the model, starting from the world model training batch states. The
actor and critic networks are instantiated as 4-layer MLP with a
dimensionality of 400 and updated using Adam with a learning rate of $8e^{-5}$
and gradients clipped to 100. The actor’s output is a truncated normal
distribution and its gradients are backpropagated using the reparameterization
trick Kingma & Welling (2013).

Rewards are computed in latent imagination using:

$r_{\textrm{skill}}(s,z)=\underbrace{\frac{1}{K}\sum_{i=1}^{K}\|s-s_{i}^{\textrm{K-NN}}\|_{2}}_{r_{\textrm{entropy}}}\underbrace{-\|D(z)-s\|_{2}\vphantom{\frac{1}{K}\sum_{i=1}^{K}}}_{r_{\textrm{code}}}.$

For the entropy term, where we approximate a particle-based entropy estimator
(Singh et al., 2003), we follow Liu & Abbeel (2021b) in averaging the distance
from the K-nearest neighbors ($K=30$) and remove the $\log$. For the code
term, we take the L2-norm directly instead of the squared L2-norm. With these
minor modifications, the two terms are on the same scale and we found it
unnecessary to add any scaling coefficient to balance the two reward terms.

Meta-controller actor-critic. The meta-controller actor-critic is instantiated
and trained in the same way as the skill actor-critic, using the environment
rewards provided by the model’s reward predictor. While the meta-controller is
updated using policy gradients (Williams, 1992), we can backpropagate the
meta-controller loss to the skill actor-critic to fine-tune them, with the
reparameterization trick.

One issue we found is that when the rewards are too small (or zero) to
predict, the predictions of the reward predictor and the critic are noisy
around zero, and tend to collapse the skill behaviors during fine-tuning. In
order to avoid that, we introduce a reward smoothing mechanism, which makes
sure the critic and reward predictions are exactly zero, until a reward $\geq
1e^{-4}$ is found. By doing so, we ensure that, until an actual reward is
encountered, the skill policies remain the ones learned during pre-training
and can be used to explore the environment, looking for rewards. To provide
consistent exploration, skills are sampled and followed for an entire
trajectory, until rewards are encountered.

###  B.2 Experiments

URLB. The URL benchmark (Laskin et al., 2021) is designed to compare different
URL strategies. The experimental setup consists of two phases: a longer data
collection/pre-training phase, where the agent has up to 2M environment steps
to interact with the environment without rewards, and a shorter fine-tuning
phase, where the agent has 100k steps to interact with a task-specific version
of the environment, where it should both discover rewards and solve a
downstream task. In the URLB, there are three domains from the DM Control
Suite (Tassa et al., 2018), Walker, Quadruped, and Jaco, with four downstream
tasks to solve per each domain.

For state-based inputs, we used an action repeat of 1. For pixel-based inputs,
we used an action repeat of 2. For the state-based experiments, we note that
the Jaco environment in ExORL slightly differs from the one in URLB, as the
position of the target is not provided to the agent in ExORL. Other than that,
the Jaco tasks remain the same and we see that, nonetheless, Choreographer is
able to find the targets without information about their position, performing
comparably to the top-performing baselines.

For the sparse-rewards Jaco, we consider the target reached when the agent
finds a reward $\geq 1e^{-4}$.

MetaWorld. In MetaWorld, the task can be adapted to be sparse, by setting
rewards equal to the success metrics in the environment. For the reach tasks,
success is when the agent reaches the target. In order to build the three goal
sets, we adapted the code to sample 50 goals in the intervals:
$\textrm{goal\\_low}=(-0.45,0.3,0.1),\textrm{goal\\_high}=(0.45,0.9,0.5)$,
with seeds $(1,2,3)$.

For both sparse experiments, the agents are pre-trained on 2M frames collected
using their exploration policy (Choreographer uses LBS as detailed in Appendix
A). Then, the pre-trained behavior is deployed for evaluation. For the skill-
based approaches, which are DIAYN, APS and Choreographer, skills are randomly
re-sampled every 50 steps.

###  B.3 Other Things We Tried

Skill Discovery. We tried alternative representations other than the skill
auto-encoder, for skill discovery. In particular, we tried contrastive
learning and soft clustering using the Sinkhorn-Knopp algorithm (Yarats et
al., 2021). While both of them worked, we found the skill vectors found by the
autoencoder to be more diverse and explainable (thanks to the reconstruction).
We also tried using a standard discrete VAE, using straight-through gradients
(Bengio et al., 2013), compared to a VQ-VAE; however, the training was way
more unstable.

Skill learning. We attempted to learn different numbers of skills, ranging
from 16 up to 512. We found that the best performance was obtained when using
a number of skills comprised between 64 and 128, and we stick to 64 for our
experiments. Too few skills were sometimes unable to find useful behaviors in
the Quadruped or to find sparse targets in the Jaco tasks. Too many skills,
under the same computational budget, generally led to skills that didn’t fully
converge or were harder to fine-tune. Though, likely, with a higher
computational budget, more skills might perform better.

Skill adaptation. The meta-controller generally samples skills at a time-
frequency of 1. Other works have focused on learning skills that last several
timesteps (Sharma et al., 2020), using them for hierarchical control (Gehring
et al., 2021). Choreographer, instead, tends to evaluate all skills in
imagination and narrow down its choice to a restricted set of skills that are
adapted for the downstream tasks. We found that sampling skills with lower
frequency (we tried every 3,5,15,50 steps) does not lead to better adaptation
performance, but often doesn’t harm performance either. We believe that
leveraging the skills for hierarchical problems is an interesting direction
for future work.

##  Appendix C Code Resampling

In RL, we mostly face datasets that grow arbitrarily over time, as the agent
collects new data. These datasets may become large (in our case, up to
$2e^{6}$ environment transitions) and thus require a long time to run
clustering algorithms, such as k-means. Furthermore, as the data increases,
the algorithm may need to be run again, several times, to comply with the new
data.

By using a VQ-VAE, we are able to perform clustering of the model state space,
as well as to decode the cluster centers into their corresponding model states
(which is necessary for Choreographer skill learning objective). However, VQ-
VAE suffers from index collapse.

We propose a code resampling technique that is inspired by the k-means++
initialization (Arthur & Vassilvitskii, 2006). This technique is generally
applied with fixed datasets to reduce the k-means convergence time, by
providing a better initialization. Instead, we use the code re-initialization
procedure to ‘resample’ codes every $M$ training batches, making sure that any
collapsed index/code is reactivated.

Our technique scales well for the RL setting, as it uses only the data in the
previous batch to resample the code. Of course, this may lead to a suboptimal
code choice, which can then be adjusted at following code resampling
iterations, if necessary. We showed in Figure 6 that in practice all the codes
tend to stay active over training.

⬇

if train_step mod M == 0:

for code in inactive_codes:

distance_matrix = l2norm_distance(current_batch, current_codes)

probabilities = min(distance_matrix, axis=code_axis) \

/ sum(min(distance_matrix, axis=code_axis))

code = current_batch[sample_index(probs=probabilities)]

Algorithm 1 Code Resampling

##  Appendix D Additional Results

In Table 2, we present detailed results and baselines on the state-based URLB.
Choreographer performs best on all the Walker and Quadruped tasks and
comparably to CIC in the Jaco tasks. The normalization scores for the results
in the paper are the Expert results.

Pre-trainining for $2\times 10^{6}$ environment steps |  |
---|---|---
Domain | Task | Expert | DDPG | ICM | Disagreement | RND | APT | ProtoRL | SMM | DIAYN | APS | CIC | Choreographer
Walker | Flip | 799 | 538$\pm$27 | 417$\pm$16 | 346$\pm$13 | 474$\pm$39 | 544$\pm$14 | 456$\pm$12 | 450$\pm$24 | 319$\pm$17 | 465$\pm$20 | 631 $\pm$ 34 | 960 $\pm$ 3
Run | 796 | 325$\pm$25 | 247$\pm$21 | 208$\pm$15 | 406$\pm$30 | 392$\pm$26 | 306$\pm$13 | 426$\pm$26 | 158$\pm$8 | 134$\pm$16 | 486 $\pm$ 25 | 635 $\pm$ 10
Stand | 984 | 899$\pm$23 | 859$\pm$23 | 746$\pm$34 | 911$\pm$5 | 942$\pm$6 | 917$\pm$27 | 924$\pm$12 | 695$\pm$46 | 721$\pm$44 | 959 $\pm$ 2 | 980 $\pm$ 1
Walk | 971 | 748$\pm$47 | 627$\pm$42 | 549$\pm$37 | 704$\pm$30 | 773$\pm$70 | 792$\pm$41 | 770$\pm$44 | 498$\pm$27 | 527$\pm$79 | 885 $\pm$ 28 | 967 $\pm$ 1
Quadruped | Jump | 888 | 236$\pm$48 | 178$\pm$35 | 389$\pm$62 | 637$\pm$12 | 648$\pm$18 | 617$\pm$44 | 96$\pm$7 | 660$\pm$43 | 463$\pm$51 | 595 $\pm$ 42 | 830 $\pm$ 9
Run | 888 | 157$\pm$31 | 110$\pm$18 | 337$\pm$30 | 459$\pm$6 | 492$\pm$14 | 373$\pm$33 | 96$\pm$6 | 433$\pm$29 | 281$\pm$17 | 505 $\pm$ 47 | 779 $\pm$ 11
Stand | 920 | 392$\pm$73 | 312$\pm$68 | 512$\pm$89 | 766$\pm$43 | 872$\pm$23 | 716$\pm$56 | 123$\pm$11 | 851$\pm$43 | 542$\pm$53 | 761 $\pm$ 54 | 946 $\pm$ 5
Walk | 866 | 229$\pm$57 | 126$\pm$27 | 293$\pm$37 | 536$\pm$39 | 770$\pm$47 | 412$\pm$54 | 80$\pm$6 | 576$\pm$81 | 436$\pm$79 | 723 $\pm$ 43 | 921 $\pm$ 5
Jaco | Reach bottom left | 193 | 72$\pm$22 | 111$\pm$11 | 124$\pm$7 | 110$\pm$5 | 103$\pm$8 | 129$\pm$8 | 45$\pm$7 | 39$\pm$6 | 76$\pm$8 | 138 $\pm$ 9 | 137 $\pm$ 10
Reach bottom right | 203 | 117$\pm$18 | 97$\pm$9 | 115$\pm$10 | 117$\pm$7 | 100$\pm$6 | 132$\pm$8 | 46$\pm$11 | 38$\pm$5 | 88$\pm$11 | 145 $\pm$ 7 | 147 $\pm$ 10
Reach top left | 191 | 116$\pm$22 | 82$\pm$14 | 106$\pm$12 | 99$\pm$6 | 73$\pm$12 | 123$\pm$9 | 36$\pm$3 | 19$\pm$4 | 68$\pm$6 | 153 $\pm$ 7 | 164 $\pm$ 7
Reach top right | 223 | 94$\pm$18 | 103$\pm$11 | 139$\pm$7 | 100$\pm$6 | 90$\pm$10 | 159$\pm$7 | 47$\pm$6 | 28$\pm$6 | 76$\pm$10 | 163 $\pm$ 4 | 169 $\pm$ 10

Table 2: Mean and standard error performance of Choreographer on state-based
URLB after pre-training for 2M steps and fine-tuning 100k steps. Baseline
scores from (Laskin et al., 2022).

In Figure 3, we showed the results of Choreographer pre-trained with 2M steps
of exploratory data from the RND approach. In Figure 7, we compare to the
results when pre-training with the data collected by other knowledge-based
approaches, i.e. ICM, Disagreement, or random actions data. We see that the
quality of the model and the skill learned, and so the fine-tuning
performance, depends on the data used. Choreographer is able to learn skills
that widely explore the model state space. However, if the state space
explored is little, then the skills will be also limited to a restricted area.
This is what generally happens for the case of the random actions in the
Figure, but also in the Quadruped for ICM (where we visually inspected the
data from ExORL, verifying that ICM explores poorly in the Quadruped
environment).

Figure 7: Different data. Training Choreographer’s world model and skills
using data from different exploration strategies.

Figure 8: Zero-shot skill control. We attempt to solve the task using the pre-
trained skill policies, without any fine-tuning.

In Figure 4, we showed the performance of Choreographer in the pixel-based
setting over time. This performance was achieved by learning the meta-
controller and adapting the skill policies during fine-tuning. In Figure 8, we
show the impact of fine-tuning the skill policies learned during pre-training,
after 2M steps. This can be seen as restricting the agent’s action space to a
set of N state-dependent actions (=64 skills in this case), which the agent
can use to solve the downstream task.

The performance on the Walker and Quadruped tasks is weaker, but the agent
still performs comparably to some weaker baselines in Figure 4, such as APS
and DIAYN. Instead, the performance on Jaco is much worse than the one with
skill adaptation, that’s because the task requires precision, to reach the
target block, and adaptability, in case the block moves.

In Figure 9, we show additional results, including more baselines from
(Rajeswar et al., 2022), on the pixel-based URLB. Overall, Choreographer
clearly outperforms all methods, while LBS is still competitive on the
Quadruped.

Figure 9: Pixel-based URLB. Additional baselines for the pixel-based URLB
experiments.

##  Appendix E Additional Visualizations

Figure 10: Walker. Sets of three skill codes and respective behaviors
visualized for Walker. Figure 11: Quadruped. Sets of three skill codes and
respective behaviors visualized for Quadruped. Figure 12: Jaco. Sets of three
skill codes and respective behaviors visualized for Jaco. Figure 13:
MetaWorld. Sets of three skill codes and respective behaviors visualized for
MetaWorld.

##  Appendix F Algorithm

Initialize agent modules: world model, skill autoencoder, skill policies ;

if  _no replay buffer available_ then

Initialize replay buffer. ;

Initialize exploration policy. ;

should_explore = True.

end if

while  _is unsupervised phase_ do

// Data collection

if  _should_explore_ then

Update model state with the posterior $s_{t}\sim
q(s_{t}|s_{t-1},a_{t-1},x_{t})$. ;

Sample exploratory action $a_{t}\sim\pi_{\textrm{expl}}(a_{t}|s_{t})$. ;

Collect new observation from the environment $x_{t+1}$ ;

Add $(x_{t},a_{t},x_{t+1})$ to the replay buffer. ;

end if

// Pre-training

if  _$t\operatorname{mod}train\\_every\\_n\\_steps=0$_ then

Draw a batch of past trajectories $\\{x,a\\}$ from the replay buffer. ;

Update world model on batch (Eq. 3) and get states $\\{s\\}$. ;

if  _should_explore_ then

Imagine exploration trajectories $\\{s,a,z\\}$, sampling actions from
$\pi_{\textrm{expl}}(a_{t}|s_{t})$. ;

Compute exploration rewards $r_{\textrm{expl}}$ (Eq. 9). ;

Update exploration actor-critic. ;

end if

// Skill discovery

Update skill autoencoder using the model states $\\{s\\}$ (Eq. 5). ;

// Skill learning

Sample skill codes from uniform distribution $z\sim p(z)$. ;

Imagine skill trajectories $\\{s,a,z\\}$, sampling actions from
$\pi_{\textrm{skill}}(a_{t}|s_{t},z)$. ;

Compute skill rewards $r_{\textrm{skill}}$ (Eq. 7). ;

Update skill actor-critic. ;

end if

end while

// Fine-tuning

Reset replay buffer and initialize meta-controller. ;

while  _is supervised phase_ do

Update model state with the posterior $s_{t}\sim
q(s_{t}|s_{t-1},a_{t-1},x_{t})$. ;

Sample skill $z\sim\pi_{\textrm{meta}}(z|s_{t})$ and action
$a_{t}\sim\pi_{\textrm{skill}}(a_{t}|s_{t},z)$. ;

Collect new task observation from the environment $x_{t+1},r_{t+1}$ ;

Add $(x_{t},a_{t},x_{t+1},r_{t+1})$ to the replay buffer. ;

if  _$t\operatorname{mod}train\\_every\\_n\\_steps=0$_ then

Draw batch of past trajectories $\\{x,a,r\\}$ from the replay buffer. ;

Update world model on batch (Eq. 3) and get states $\\{s\\}$. ;

// Skill adaptation

Imagine traj. $\\{s,a,z\\}$, sampling skills from
$\pi_{\textrm{meta}}(z|s_{t})$ and actions from
$\pi_{\textrm{skill}}(a_{t}|s_{t},z)$.;

Predict task rewards $r_{\textrm{task}}$ ;

Update meta-controller actor-critic and propagate gradients to skill actors. ;

end if

end while

Algorithm 2 Choreographer