# Temporal Abstractions-Augmented Temporally Contrastive Learning:
An Alternative to the Laplacian in RL

Akram Erraqabi  Mila  Université de Montréal  Marlos C. Machado  Amii
University of Alberta  CIFAR AI Chair  Mingde Zhao  Mila  McGill University
Sainbayar Sukhbaatar  Meta AI
Alessandro Lazaric  Meta AI  Ludovic Denoyer  Meta AI  Yoshua Bengio  Mila
Université de Montréal  CIFAR Fellow

###### Abstract

In reinforcement learning, the graph Laplacian has proved to be a valuable
tool in the task-agnostic setting, with applications ranging from skill
discovery to reward shaping. Recently, learning the Laplacian representation
has been framed as the optimization of a temporally-contrastive objective to
overcome its computational limitations in large (or continuous) state spaces.
However, this approach requires uniform access to all states in the state
space, overlooking the exploration problem that emerges during the
representation learning process. In this work, we propose an alternative
method that is able to recover, _in a non-uniform-prior setting_ , the
expressiveness and the desired properties of the Laplacian representation. We
do so by combining the representation learning with a skill-based covering
policy, which provides a better training distribution to extend and refine the
representation. We also show that a simple augmentation of the representation
objective with the learned temporal abstractions improves dynamics-awareness
and helps exploration. We find that our method succeeds as an alternative to
the Laplacian in the non-uniform setting and scales to challenging continuous
control environments. Finally, even if our method is not optimized for skill
discovery, the learned skills can successfully solve difficult continuous
navigation tasks with sparse rewards, where standard skill discovery
approaches are no so effective.

##  1 Introduction

With††Correspondance to AE<EMAIL_ADDRESS>the advent of deep
reinforcement learning [Mnih et al., 2015], representation learning [c.f.
Bengio et al., 2013] has become one of the main topics of interest in
reinforcement learning (RL). In fact, learning in environments with rich
observations and complex dynamics [e.g., Bellemare et al., 2020] has motivated
recent works on learning representations, e.g. controllable or contingent
features [Bengio et al., 2017, Choi et al., 2019] on top of which one can
potentially learn latent models in the perspective of planning [Hafner et al.,
2019, Nasiriany et al., 2019, Schrittwieser et al., 2020] and control [Watter
et al., 2015, Banijamali et al., 2018, Hafner et al., 2020].

Figure 1: Our representation is trained to encode the area that the agent has
learned to cover. Skills are continuously trained on the representation to
discover new areas where novel data is collected to refine the representation,
progressively extending its coverage. Similar incremental discovery is at the
core of other works [Ecoffet et al., 2021, Pong et al., 2020, Machado, 2019].

In this work, we are interested in the task-agnostic setting in which an RL
agent first interacts with the environment to build a representation, $\phi$,
of the state space, $\mathcal{S}$, without relying on any specific reward
signal. This representation can later be used to solve a task posed in the
environment in the form of a reward function. In this setting, the environment
dynamics are the only informative interaction channel available to the agent.
This has motivated the use of graph Laplacian-based methods where the graph
vertices correspond to the states and its edges to the transitions
probabilities. The Laplacian’s eigenvectors can been leveraged as a holistic
state representation, termed the Laplacian representation, which captures the
environment dynamics [Mahadevan, 2005, Mahadevan and Maggioni, 2007].

Wu et al. [2019] have recently proposed an efficient approximation of the
Laplacian representation (Lap-rep) by framing the graph drawing objective as a
temporally-contrastive loss (see Section 2.2). This formulation works around
potentially prohibitive eigendecompositions and extends the Laplacian’s
applicability to large (and continuous) state spaces. However, it assumes
access to a uniform sampling prior over all states in the state space. In
practice, this translates in the ability to reset the agent to any state in
the environment, which artificially alleviates the exploration problem. As we
show in Section 4, this assumption is crucial for the quality of the learned
representation. In the absence of the uniform prior privilege, such sampling
is not trivial to achieve since the agent has to first explore and learn about
the state space to be able to access arbitrary states. In effect, one must
handle the exploration along the representation learning in order to preserve
the quality of the representation. In this work, we propose TATC, an
alternative representation learning framework to Lap-rep that extends a
similar temporally-contrastive approach to a non-uniform-prior setting while
preserving the desired properties and quality.

In practice, the representation is trained on data collected with a uniformly
random policy. However, without a uniform access to the state space, the
collected data is concentrated around accessible starting states. To achieve
better data collection, we tie the representation learning problem to that of
learning a covering strategy. Briefly, our method consists in using the
available representation to train a skill-based covering policy that is in
turn used to discover yet unseen parts of the state space, providing novel
data to refine and expand the representation. Our approach, illustrated in
Figure 1 shares a similar motivation with several previous works [Machado,
2019, Machado et al., 2017, 2018, Jinnai et al., 2020]. In addition to the
aforementioned virtuous learning cycle between representation and skills, we
propose to integrate the temporal abstractions captured by the skills in the
contrastive representation learning objective. This augmentation contributes
to a better temporally-extended exploration and enforce the representation’s
_dynamics-awareness_ , i.e. how representative the representation-induced
metric is of distances in the state space.

We empirically show our agent’s ability to progressively explore the state
space and consistently extend the domain covered by representation in a non-
uniform-prior setting. We show that our representation leads to better value
predictions than Lap-rep, and that it recovers the representation quality
expected from a uniform prior. We also evaluate our representation in shaping
rewards for goal-achieving tasks, and we show it outperforms Lap-rep,
confirming both its superior ability in capturing dynamics and in scaling to
_larger_ environments. Finally, the skills learned with our framework also
prove to be successful at difficult continuous navigation tasks with sparse
rewards, where other standard skill discovery methods have limited efficacy.

##  2 Preliminaries

###  2.1 Task-agnostic RL

We describe a task-agnostic reinforcement learning (RL) environment as a task-
agnostic Markov decision process (MDP)
$\mathcal{M}=(\mathcal{S},\mathcal{A},P,\gamma,d_{0})$ where $\mathcal{S}$ is
the state space, $\mathcal{A}$ the action space,
$P:\mathcal{S}\times\mathcal{A}\rightarrow{\Delta(\mathcal{S})}$ is the
transition dynamics defining the next state distribution given the current
state and action taken, $\gamma\in[0,1)$ the discount factor, and $d_{0}$ is
the initial state distribution. A policy
$\pi:\mathcal{S}\rightarrow\Delta(\mathcal{A})$ maps states $s\in\mathcal{S}$
to distributions over actions. We denote by $\Delta$ the probability simplex.

Knowledge acquired from task-agnostic interactions with the environment (e.g.,
a representation or a policy) can then be leveraged for specific tasks. A task
is instantiated with a reward function, $R:\mathcal{S}\rightarrow\mathbb{R}$,
which is combined with the task-agnostic MDP. The task objective is to find
the optimal policy maximizing the expected discounted return,
$\mathbb{E}_{\pi,d_{0}}\Big{[}\sum_{t}\gamma^{t}R(s_{t},a_{t})\Big{]}$,
starting from state $s_{0}\sim d_{0}$ and acting according to
$a_{t}\sim\pi(\cdot|s_{t})$.

###  2.2 The Laplacian Representation

The Laplacian representation (Lap-rep), as proposed by Wu et al. [2019], can
be learned with the following contrastive objective:

$\mathcal{L}_{Lap}(\phi;\mathcal{D}_{\pi_{\mu}})=\mathbb{E}_{(u,v)\sim\mathcal{D}_{\pi_{\mu}}}\Big{[}\|\phi(u)-\phi(v)\|^{2}_{2}\Big{]}+\\\
\beta\leavevmode\nobreak\ \leavevmode\nobreak\
\mathbb{E}_{\begin{subarray}{c}u\sim\mathcal{D}_{\pi_{\mu}}\\\
v\sim\mathcal{D}_{\pi_{\mu}}\end{subarray}}\Big{[}(\phi(u)^{\top}\phi(v))^{2}-\|\phi(u)\|^{2}_{2}-\|\phi(v)\|^{2}_{2}\Big{]},$
(1)

where $\beta$ is a hyperparameter, $\phi:\mathcal{S}\rightarrow\mathbb{R}^{d}$
is a $d$-dimensional representation, $\pi_{\mu}$ the uniformly random policy
(random walk trajectories), $\mathcal{D}_{\pi_{\mu}}$ a set of trajectories
from $\pi_{\mu}$ (random walks). We use $(u,v)\sim\mathcal{D}_{\pi_{\mu}}$ to
denote the sampling of a random transition from $\mathcal{D}_{\pi_{\mu}}$, and
similarly $u\sim\mathcal{D}_{\pi_{\mu}}$ for a random state. Wu et al. [2019]
showed the competitiveness of the Laplacian representation when provided with
a uniform prior over $\mathcal{S}$ during the collection of
$\mathcal{D}_{\pi_{\mu}}$. This objective is a _temporally-contrastive_ loss:
it is comprised of an attractive term that forces temporally close states to
have similar representations, and of a repulsive term that keeps far apart
temporally far states’ representations. Here, the repulsive term was
specifically derived from the orthonormality constraint of the Laplacian
eigenvectors.

###  2.3 The non-uniform prior setting

In RL, representation learning is deeply coupled to the problem of
exploration. Indeed, the induced state distribution defines the
representation’s training distribution. For instance, Lap-rep [Wu et al.,
2019] has been learned in the specific _uniform prior_. In this setting,
$\mathcal{D}_{\pi_{\mu}}$, from Equation 1, is a collection of trajectories
with uniformly random starting states, which provides a uniform training
distribution to the representation learning objective. In the case of a non-
uniform prior, the induced visitation distribution can be quite concentrated
around the starting states distribution when solely relying on random walks to
collect data, hence the need for a better exploration strategy in order to
achieve a better training distribution for $\phi$.

To study the problem described above, we investigate the setting in which the
environment has a fixed predefined state $s_{0}$ to which it resets with a
probability $p_{r}$ every $K$ steps; with $K$ of the order of diameter of
$\mathcal{S}$. With a uniformly random behavior policy, this setting is
equivalent to a initial state distribution that is concentrated around $s_{0}$
and whose density decays exponentially away from it. We will refer to this
setting as the _non-uniform-prior_ (non-$\mu$) setting, as opposed to the
_uniform-prior_ ($\mu$) setting where the agent has uniformly access to the
state space.

##  3 Temporal abstractions augmented representation learning

In this section, we present Temporal Abstractions-augmented Temporally-
Contrastive learning (TATC), a representation learning approach in which the
representation works in tandem with a skill-based covering policy for a better
representation learning in the non-uniform prior setting.

Before presenting the components of TATC and how they are trained, we first
provide an intuitive description of how it operates in the non-uniform prior
setting. A detailed description of the algorithm can be found in Appendix A.

###  3.1 TATC: a sketch of the algorithm

In the non-uniform-prior setting, the agent is reset to a fixed state $s_{0}$
after each episode with some probability $p_{r}$. At the beginning of each
episode (not necessarily at $s_{0}$, due to the probabilistic resetting), our
agent can choose (with some probability $p_{rw}$) either to follow the uniform
policy $\pi_{\mu}$ or to act according to an exploratory skill-based policy.
Random walk data are, similarly to Wu et al. [2019], collected to train the
representation, while the skills are trained to extend the area of
$\mathcal{S}$ that the agent is able to efficiently reach. In order to
leverage the rich compositionality of skills, the agent executes a sequence of
$L$ consecutive skills each time it decides to call the skill-based policy.
Initially, trajectories from $\pi_{\mu}$ cover the vicinity of $s_{0}$, making
the representation reliable in that area, i.e. representative of its dynamics.
Inevitably, the skills trained on this representation benefit from its
emerging structure and progressively gain _behavioral_ structure: they allow
the agent to travel efficiently across this explored area. In other words, the
agent becomes capable of reaching the frontier of the explored regions faster,
and it is able to collect, using $\pi_{\mu}$, novel data for the
representation. The latter is hence refined, and its coverage extended. With a
refined representation, the skills are able to reach even further areas. This
process emerges as a virtuous collaboration between the representation and the
skills, where both benefit from improving each other. Eventually, by acquiring
more knowledge from unexplored areas, the agent helps overcoming the loss in
the representation’s expressiveness that we observed when solely relying on
$\pi_{\mu}$ in the non-uniform-prior setting.

In the remainder of this section we first propose a generic alternative
objective to Equation 1 that suits the non-uniform prior setting, then we
describe the exploratory policy training. Finally, we introduce an
augmentation of the proposed objective based on the learned temporal
abstractions, to improve exploration and enforce the dynamics-awareness of the
representation.

###  3.2 Temporally-Contrastive Representation Objective

As mentioned in Section 2.2, the repulsive term in Lap-rep’s objective, in
Equation 1, derives from the eigenvectors’ orthonormality constraint. However,
because the environment is expected to be progressively covered in the non-
uniform prior setting, the orthonormality constraint can make online
representation learning highly non-stationary.111In general, even within a
given matrix’s perturbation neighborhood, its eigenvectors can show a highly
nonlinear sensitivity [Trefethen and Bau, 1997]. For this reason, we adopt the
following objective with a generic repulsive term that is more amenable to
online learning:

$\mathcal{L}_{cont}(\phi;\mathcal{D}_{\pi_{\mu}})\triangleq\mathbb{E}_{(u,v)\sim\mathcal{D}_{\pi_{\mu}}}\left[\|\phi(u)-\phi(v)\|^{2}_{2}\right]+\\\
\beta\leavevmode\nobreak\ \leavevmode\nobreak\
\mathbb{E}_{\begin{subarray}{c}u\sim\mathcal{D}_{\pi_{\mu}}\\\
v\sim\mathcal{D}_{\pi_{\mu}}\end{subarray}}\left[\exp(-\|\phi(u)-\phi(v)\|_{2})\right].$
(2)

In the following, we describe our representation-based skills training
framework. Beyond addressing the exploration need, these skills will later
(Section 3.4) be used to augment the objective above to obtain TATC’s
representation learning objective.

###  3.3 Representation-based Skills Training

In the non-uniform-prior setting, exploration is required to provide the
representation with a better training distribution. To this purpose, we adopt
a hierarchical RL approach to leverage the exploration efficiency of skills,
also known as options [Sutton et al., 1999]. Let
$\phi:\mathcal{S}\rightarrow\mathbb{R}^{d}$ be our $d$-dimensional
representation. The agent acts according to a bi-level policy
$(\pi_{\text{hi}},\pi_{\text{low}})$. The high-level policy
$\pi_{\text{hi}}:\mathcal{S}\to\Delta(\Omega)$ defines, at each state $s$, a
distribution over a set $\Omega$ of directions (unit vectors) in the
representation space ($\Omega=\\{\bm{\delta}\leavevmode\nobreak\
|\leavevmode\nobreak\
\bm{\delta}\in\mathbb{R}^{d},\|\bm{\delta}\|_{2}{=}1\\}$). Each direction
corresponds to a fixed length skill encoded by the low-level policy
$\pi_{\text{low}}:\mathcal{S}\times\Omega\rightarrow\Delta(\mathcal{A})$.
These skills are trained to travel _in the representation space_ along the
directions instructed by $\pi_{\text{hi}}$. In short, given a sampled
direction $\pi_{\text{hi}}(\cdot|s)\sim\bm{\delta}\in\Omega$, the low-level
policy executes the skill $\pi_{\text{low}}(\cdot|s,\bm{\delta})$ for a fixed
number of steps $c$ before $\pi_{\text{hi}}$ is called again.

Now, we describe the intrinsic rewards used to train the policies
$\pi_{\text{low}}$ and $\pi_{\text{hi}}$.

##### Low-level Policy.

$\pi_{\text{low}}$ is trained to follow directions picked by $\pi_{\text{hi}}$
in the representation space. For a given
$\bm{\delta}\in\Omega\subset\mathbb{R}^{d}$, the corresponding skill
$\pi_{\text{low}}(\cdot|s,\bm{\delta})$ is trained to maximize the reward
function:

$r^{\bm{\delta}}(s,s^{\prime})\triangleq\frac{(\phi(s^{\prime})-\phi(s))^{\top}\bm{\delta}}{\|\phi(s^{\prime})-\phi(s)\|_{2}}\leavevmode\nobreak\
\leavevmode\nobreak\ \leavevmode\nobreak\ ,$ (3)

where $(s,s^{\prime})$ is an observed state transition, and $\phi$ the
representation being learned. We use the cosine similarity as a way to
encourage learning diverse directional skills. Indeed, skills co-
specialization is avoided by rewarding the agent for the steps induced along
the instructed direction $\bm{\delta}$ regardless of their magnitudes.

##### Connection to Behavioral Mutual Information.

It is worth noting that our reward design can be interpreted as a mutual-
information-based intrinsic control. We provide more details on this
connection in Appendix E.

##### High-level Policy.

The high-level policy guides the covering strategy. It does so by sampling the
skills of the most promising directions in terms of the amount of exploration,
measured by the travelled distance in the representation space. For this
purpose, we design a reward function defined over a sequence of $L$
consecutive skills. Let $\\{s^{\text{hi}}_{k}\\}_{k=1}^{L}$ be the sequence of
their initial states and their respective sampled directions,
$\bm{\delta}_{k}\sim\pi_{\text{hi}}(\cdot|s^{\text{hi}}_{k})$. Since $\phi$ is
trained to capture the dynamics, the travelled distance in $\phi$’s space
_expresses_ how far the choices made by $\pi_{\text{hi}}$ eventually brought
the agent in the environment. Therefore, for a given high-level trajectory,
$\tau^{\text{hi}}=(s^{\text{hi}}_{1},s^{\text{hi}}_{2},...,s^{\text{hi}}_{L},s^{\text{hi}}_{f})$,
with $s^{\text{hi}}_{f}$ the final state reached by the last skill, the high-
level policy is trained to maximize the following:

$\forall
k\in\\{1,...,L\\},R^{\text{hi}}(s^{\text{hi}}_{k},\bm{\delta}_{k})\triangleq\|\phi(s^{\text{hi}}_{1})-\phi(s^{\text{hi}}_{f})\|_{2}\leavevmode\nobreak\
\leavevmode\nobreak\ ,$ (4)

where $\bm{\delta}_{k}\sim\pi_{\text{hi}}(\cdot|s^{\text{hi}}_{k})$ is the
direction sampled at $s^{\text{hi}}_{k}$. From the policy optimization
perspective, each of these quantities plays the role of the return cumulated
along the sampled high-level trajectory and _not_ just a single (high-level)
step reward. This term looks at reaching $s^{\text{hi}}_{f}$ as the result of
a sequential collaboration of $L$ skills, rewarding them equally. It values
how far this sequence of skills eventually brought the agent.

In the following section, we show how the skills trajectories can be used to
augment the representation with the learned temporal abstractions.

###  3.4 Augmenting Representation Learning with Temporal Abstractions

A skill abstracts a temporally-extended behavior in a single (high-level)
action. As a _temporal abstraction_ , it represents a factorized knowledge of
the environment dynamics in the form of a policy. Here, we propose to
integrate these temporal abstractions to the representation objective to
better capture the environment dynamics. In order to preserve the temporal
contrast of the base objective (2), we augment it with the following
contracting term along skills trajectories:

$\mathcal{B}(\phi;\mathcal{D}_{s})\triangleq\leavevmode\nobreak\
\leavevmode\nobreak\ \leavevmode\nobreak\ \leavevmode\nobreak\
\leavevmode\nobreak\ \leavevmode\nobreak\
\mathop{\mathbb{E}}_{\mathclap{\begin{subarray}{c}\tau_{\bm{\delta}}\sim\mathcal{D}_{s}\\\
\tau_{\bm{\delta}}=(s_{0},...,s_{c})\end{subarray}}}\leavevmode\nobreak\
\leavevmode\nobreak\ \leavevmode\nobreak\ \leavevmode\nobreak\
\leavevmode\nobreak\ \leavevmode\nobreak\ \leavevmode\nobreak\
\leavevmode\nobreak\
\left[\sum_{k=0}^{c-1}\|\phi(s_{k})-\phi(s_{k+1})\|_{2}\right],$ (5)

where $\mathcal{D}_{s}$ is a set of collected skills trajectories. By
minimizing this term, $\phi$ integrates temporally-extended dynamics: areas
connected by skills are brought closer in the representation space. This term
will be referred to as the _boredom_ term. Its exploratory virtue is discussed
in the following.

##### How does boredom help exploration ?

The interplay between the high-level policy reward function (4) and this
boredom term (5) induces a progressive exploration mechanism. In effect,
$\pi_{\text{hi}}$ tends to sample skills that travel further, i.e. with larger
$R^{\text{hi}}$. The more often a skill is sampled, the less rewarding it
becomes to $\pi_{\text{hi}}$ due to the minimization of $\mathcal{B}(\phi)$.
This will favor sampling the remaining (under-sampled) skills, hence
encouraging the exploration of less visited parts of the state space. This
mechanism dynamically fights what can be considered as accumulated _boredom_
along over-sampled skills trajectories which increases the agent curiosity and
urge it to explore.

Finally, the proposed objective to train the representation $\phi$ consists in
the base objective (2 augmented with the boredom term (5), and can be written
as

$\mathcal{L}_{\text{TATC}}(\phi;\mathcal{D}_{s},\mathcal{D}_{\pi_{\mu}})\triangleq\mathcal{L}_{cont}(\phi;\mathcal{D}_{\pi_{\mu}})+\beta^{\prime}\mathcal{B}(\phi;\mathcal{D}_{s})\leavevmode\nobreak\
,$ (6)

with $\beta^{\prime}$ a hyperparameter controlling the strength of boredom
term.

##  4 Experiments

In this section, we investigate the behavior of TATC in two types of
environments: gridworld environments with discrete state and action spaces,
and continuous navigation environments for continuous state spaces (MuJoCo,
Todorov et al. [2012]).

For our method, we learn a 2D representation ($d=2$), and define $\Omega$ as a
set of $8$ unit vectors equally spaced on the unit sphere. These directions
are concatenated to the input state of $\pi_{\text{low}}$. Implementation
details of all the experiments in this section can be found in Appendix C.

###  4.1 GridWorld

We evaluate our approach in three gridworld domains: U-Maze, T-Maze and
4-rooms. These environments, visualized in Figure 2, raise different
explorations challenges. U-Maze is a simple but relevant environment to test
the dynamics-awareness of the representations;222The presence of the wall
makes L2-distance in xy-coordinates deceptive. The L2-distance in a dynamics-
aware representation space should correct for that. T-Maze raises the
challenge of splitting the exploration focus at an intersection while
maintaining the covering in both corridors; 4-rooms is similar to
U-Maze333Note that there is no door between the first and the fourth room.,
but requires learning more controlled skills for a useful exploration.

(a) U-Maze

(b) T-maze

(c) 4-rooms

Figure 2: Gridworld domains. The fixed initial state $s_{0}$ is highlighted in
red.

Figure 3: TATC representations learned throughout the training. For each
domain, all the states are mapped with $\phi$ and represented at different
stages of the training procedure. Axes scales were equalized for reliable
visual appreciation. Top row: U-Maze. Middle row: T-Maze. Bottom row: 4-rooms.
The colors reflect the distances in terms of the dynamics. They can be seen as
quantities proportional to the length of the shortest path from $s_{0}$
(marked in red) to the represented state.

####  4.1.1 Progressive Representation Learning

Figure 3 shows the evolution of the representations throughout training. The
agent progressively explores the environment starting around $s_{0}$, and
builds the representation by continuously integrating newly discovered parts.

U-Maze & 4-rooms. The agent starts from the bottom left corner, or room, of
the maze. Figure 3 shows how the representation progressively expands away
from $s_{0}$ along the corridor, or the rooms sequence. Note that while the
agent learns to reach and represent further areas, the full domain
representation _flattens_ , indicating the representation’s success in
capturing the maze dynamics.

T-Maze. The agent starts from the bottom left corner of the maze. As in the
U-Maze, it starts learning to travel along the corridor until it reaches the
intersection. There, the exploration focus is shared between both possible
paths whose representations are progressively disentangled. Eventually, the
agent fully explores both corridors and finalizes its representation. Note
that, the discovery of one of the corridors did not hinder finishing the
discovery of the other. The boredom term proved to be important for such
property (see Appendix B).

Importance of Boredom. Appendix B provides an ablation study showing the
importance of the boredom term for the agent’s exploratory behavior and the
representation’s dynamics-awareness.

####  4.1.2 Evaluating the Learned Representation

We now compare our representation against Lap-rep [Wu et al., 2019] in the
non-uniform-prior setting. First, to appreciate the sensitivity of Lap-rep to
the uniformity of said prior, we trained Lap-rep in two settings: (i) the
uniform-prior setting where the agent can be set to any arbitrary state, as
done by Wu et al. [2019], and (ii) the non-uniform-prior setting defined in
Section 2.3. In the following, we show that Lap-rep is sensitive to this
change in distribution while TATC recovers the expressive potential that a
uniform prior provides.

Prediction. To evaluate the learned representations, we first consider how
well they linearly approximate a given task’s optimal value function. To do
so, we train an actor-critic agent [Mnih et al., 2016] with a linear critic on
top of each representation. In Figure 5, we note a significant loss in the
representational power of Lap-rep when the access to the state space is not
uniformly distributed anymore. This figure also shows that TATC outperforms
Lap-rep in the non-uniform-prior setting, and succeeds in recovering Lap-rep’s
expressive power when it is learned with the unrealistic uniform prior.

Control. We also compare the representations from the perspective of control,
by training a deep actor-critic agent on top of each representation to solve a
goal-reaching task in the same domains as above. The agent is only rewarded
upon reaching the goal state ($r=1$). Figure 6 shows that TATC consistently
outperforms Lap-rep, which confirms the competitive quality of our
representation.

###  4.2 Continuous Control

The second set of experiments focuses on continuous state and action spaces.
Here, we consider two mazes for the MuJuCo Ant agent, as shown in Figure 4.
These are similar in shape to the ones from Wu et al. [2019], but are larger
and thus more challenging. More details in Appendix C.

To visualize our learned representations in these environments, Figure 4
depicts a grid of positional states in each environment domain and their
mapped representations. Similarly to U-Maze and 4-rooms, the learned
representation translates the environment dynamics, and specifically the
presence of walls, by mapping the original state space to a flatter manifold.

(a) AntMaze-1

(b) AntMaze-2

Figure 4: Learned TATC representations visualized on grids of positional
states. Colors reflect the distance in the representation space from the
initial state, highlighted in red. Axes scales were equalized. We can visually
appreciate how the continuous state domains are mapped to a flatter manifold
reflecting the presence of the walls. Figure 5: Learned representation’s
ability to approximate the value function. Lap-rep was learned in the same
non-uniform-prior setting (non-$\mu$) with $d=2$ and $d=3$ (no improvement was
observed for higher values). The dashed line gives the performance of Lap-rep
in the uniform-prior setting ($\mu$). TATC outperforms Lap-rep in non-$\mu$
setting, and succeeds in recovering its expressive power when learned from the
uniform prior. Performances were averaged over 5 different runs. $95\%$
confidence intervals are shaded. Figure 6: Control performance (Episode
Reward) in the non-uniform-prior setting. Performances were averaged over 5
different runs. $95\%$ confidence intervals are shaded.

####  4.2.1 Reward Shaping with Learned Representations

We demonstrate how TATC representation is able to improve an RL agent’s
performance when the distances in the representation space are used for reward
shaping, the same setting in which Wu et al. [2019] evaluated Lap-rep. We
define a goal-achieving task by setting a goal state $g$ at the end of the
corridor (visualized in Appendix C). The objective is to learn to navigate to
a state $s$ close enough to the goal area ($\|s-g\|_{2}\leq\epsilon$). We
define the reward function based on the distance in representation space (TATC
and Lap-rep). More specifically, we train a soft actor-critic (SAC) agent
[Haarnoja et al., 2018] to reach the goal with a dense reward defined as
$r^{dense}_{t}=-\|\phi(s_{t+1})-\phi(g)\|_{2}$. Similarly to Wu et al. [2019],
we also compare against the half-half mix of the dense reward and the sparse
reward $r^{mix}_{t}=0.5\cdot
r^{dense}_{t}+0.5\cdot\mathds{1}\left[\|s_{t+1}-g\|_{2}\leq\epsilon\right]$.

For this evaluation, Lap-rep was learned, unlike our representation, with a
_uniform_ prior over $\mathcal{S}$ as in Wu et al. [2019], and $d=2$.444Our
attempts with $d=20$ did not succeed at these reward shaping tasks. Figure 7
shows that our representation is effective in reward shaping, with both mix
and dense variants , and enjoys a comparable if not superior dynamics-
awareness to Lap-rep. This result stands while TATC is learned from a non-
uniform prior which is less advantageous than the uniform prior used to train
Lap-rep. Note that Lap-rep with a non-uniform prior is unable to guide the
agent to the goal.

Finally, these results further confirm the conclusions drawn from the
gridworld experiments and positions TATC as a competitive alternative to Lap-
rep in this challenging setting.

(a) AntMaze-1

(b) AntMaze-2

Figure 7: Reward shaping using learned representations: performances were
averaged over 5 different runs. $95\%$ confidence intervals are shaded. Curves
were exponentially smoothed (0.9) for better visualization.

####  4.2.2 Evaluating the Learned Skills

To evaluate the exploratory potential of TATC’s skills, we compare the learned
skills in AntMaze-1 against 2 task-agnostic skill discovery methods, DIAYN
[Eysenbach et al., 2019], and DCO [Jinnai et al., 2020]. DIAYN is a mutual
information-based approach to learn diverse set of skills, while DCO skills
are based on a temporally-contrastive representation, similarly to TATC. More
specifically, DCO requires a pretrained Lap-rep which approximates the
Laplacian’s second eigenvector; also called the Fiedler vector. We train the
required representation, as well as DCO, with the advantage of data collected
from a _uniform_ prior over $\mathcal{S}$. For a fair comparison, we train $8$
skills for both methods (DCO and DIAYN). Once trained, the skills learned by
each method are fixed and used to train a discrete high-level policy that can
select across the available skills to solve a goal-reaching task with a sparse
reward function $r_{t}=\mathds{1}\left[\|s_{t+1}-g\|_{2}\leq\epsilon\right]$.

The sparsity of the reward naturally poses a challenge as no additional signal
can guide the agent towards the goal, unlike the evaluation setting of DCO and
DIAYN by Jinnai et al. [2020]. Figure 8 shows that the skills learned by TATC
quickly assist to complete the task while the skills learned with DCO and
DIAYN do not. DIAYN’s limited performance in difficult sparse-reward
navigation tasks was also confirmed by Kamienny et al. [2021]. These results
suggest that in order to succeed, DCO and DIYAN skills may require a richer
signal like the distance-based dense reward used by Jinnai et al. [2020] to
evaluate both of them – and where they show similar performances.

Figure 8: Skills Evaluation: performance gathered from $5$ independent runs.
$95\%$ confidence intervals are shaded. Curves were exponentially smoothed
(0.9) for better visualization.

##  5 Related Work

Our work is related to self-supervised learning [Bromley et al., 1993, Chopra
et al., 2005], which brought recent advances in representation learning
[Bachman et al., 2019, He et al., 2020, Chen et al., 2020, Grill et al., 2020,
Caron et al., 2020]. These techniques have naturally been adapted to RL,
especially contrastive methods. While some of these benefited from visually
contrasting observations [Laskin et al., 2020, Yarats et al., 2021], others
leveraged temporal contrasts to learn useful representations [Mazoure et al.,
2020, Stooke et al., 2021, Li et al., 2021], which fall closer to our work.

We designed our covering policy as a hierarchical agent. This has actually
been the default setting to model temporally-extended actions [Sutton et al.,
1999]. Our work shares the same motivation as Vezhnevets et al. [2017] for
training skills to follow latent directions. Among the large body of work on
skill discovery, the eigenoptions framework [Machado et al., 2017] and its
extensions [Machado et al., 2018, 2021, Jinnai et al., 2020] are probably the
closest to our skill training scheme. Eigenoptions also fit in the directional
skills definition as they are trained to travel along the directions defined
by the eigenvectors of the Laplacian. These vectors have a dimensionality of
$|\mathcal{S}|$ which can be very large. To contrast, our directional skills
are defined by an arbitrary and diverse set of directions in the learned
representation space of adjustable dimensionality, which offers more
tractability. TATC skills and eigenoptions also share an interesting
connection to mutual information based intrinsic control methods [Hansen et
al., 2020], which we discuss in Appendix E. Indeed, a similar discussion can
be adapted to eigenoptions’ intrinsic reward function.

The exploration mechanism in TATC emerges from the interplay between the
representation and the covering policy rewarding scheme. As discussed in
Section 3.4, the boredom term prevents the high-level policy from focusing
exploration resources on previously over-sampled skills. This penalization
induces _optimism_ towards the remaining skills, in the hope of deviating from
previously explored regions. Zhang et al. [2021] share a similar motivation
and derive an exploration bonus that aims at directly maximizing this
deviation in terms of the policy occupancy. Both methods are inspired from the
_optimism-in-the-face-of-uncertainty_ principle.

The incremental discovery paradigm has been previously adopted, either for
exploration [Ecoffet et al., 2021], incremental skill discovery [Jinnai et
al., 2020, Pong et al., 2020], or even state abstraction [Misra et al., 2020].
Finally, we use learned skills to penalize boredom [Schmidhuber, 1991, Oudeyer
and Kaplan, 2009] in the representation space and encourage exploration. The
idea of using skills to foster curiosity has also been investigated by Bougie
and Ichise [2020].

##  6 Conclusion

The Laplacian representation as proposed by Wu et al. [2019] made the benefits
of spectral methods affordable in large state spaces. Unfortunately, the
quality of this representation is strongly tied to the uniformity of its
training data distribution, as shown is Section 4. This has motivated the
method proposed in this work where we reconcile a similar temporally-
contrastive representation with exploration demanding settings. Our approach
leverages the practical skill training framework that such representations
allow. The learned skills are used to better cover the state space and hence
learn a better representation. We validate our method in tabular as well as
continuous environments. Our representation learned in a non-uniform-prior
setting shows a comparable representational power to the one acquired by a
Laplacian representation from a uniform prior, and our skills proved to be
competitive in hard continuous control tasks. With these results, we hope to
bring such representations’ applicability one step closer to realistic
contexts.

We have proposed to augment the representation objective with temporal
abstractions captured in the acquired skills. This benefits exploration by
inducing a boredom-fighting mechanism, and enforces the representation’s
dynamics-awareness. Intuitively, this augmentation can be seen as bringing
temporally close regions (connected through skills) closer in the
representation space. This observation may motivate further investigations on
how temporal abstractions can improve the representational potential and
benefit task-agnostic representations in RL.

## Acknowledgement

AE would like to thank Emmanuel Bengio, Amy Zhang, and Ahmed Touati for the
insightful discussions, and Jad Kabbara and Riashat Islam for their feedback
on an earlier version of this paper. The authors acknowledge the financial
support of CIFAR, NSERC and Mila.

## References

  * Bachman et al. [2019] Philip Bachman, R Devon Hjelm, and William Buchwalter.  Learning representations by maximizing mutual information across views.  In _NeurIPS_ , 2019.
  * Banijamali et al. [2018] Ershad Banijamali, Rui Shu, Hung Bui, Ali Ghodsi, et al.  Robust locally-linear controllable embedding.  In _AISTATS_ , 2018.
  * Bellemare et al. [2020] Marc G. Bellemare, Salvatore Candido, Pablo Samuel Castro, Jun Gong, Marlos C. Machado, Subhodeep Moitra, Sameera S. Ponda, and Ziyu Wang.  Autonomous navigation of stratospheric balloons using reinforcement learning, 2020.
  * Bengio et al. [2017] Emmanuel Bengio, Valentin Thomas, Joelle Pineau, Doina Precup, and Yoshua Bengio.  Independently controllable features.  _arXiv:1703.07718_ , 2017.
  * Bengio et al. [2013] Yoshua Bengio, Aaron Courville, and Pascal Vincent.  Representation learning: A review and new perspectives.  _IEEE Trans. Pattern Anal. Mach. Intell._ , 2013.
  * Bougie and Ichise [2020] Nicolas Bougie and Ryutaro Ichise.  Skill-based curiosity for intrinsically motivated reinforcement learning.  In _Machine Learning_ , 2020.
  * Bromley et al. [1993] Jane Bromley, Isabelle Guyon, Yann LeCun, Eduard Säckinger, and Roopak Shah.  Signature verification using a "siamese" time delay neural network.  In _NeurIPS_ , 1993.
  * Caron et al. [2020] Mathilde Caron, Ishan Misra, Julien Mairal, Priya Goyal, Piotr Bojanowski, and Armand Joulin.  Unsupervised learning of visual features by contrasting cluster assignments.  In _NeurIPS_ , 2020.
  * Chen et al. [2020] Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton.  A simple framework for contrastive learning of visual representations.  In _ICML_ , 2020.
  * Choi et al. [2019] Jongwook Choi, Yijie Guo, Marcin Moczulski, Junhyuk Oh, Neal Wu, Mohammad Norouzi, and Honglak Lee.  Contingency-aware exploration in reinforcement learning.  In _ICLR_ , 2019.
  * Chopra et al. [2005] Sumit Chopra, Raia Hadsell, and Yann LeCun.  Learning a similarity metric discriminatively, with application to face verification.  In _CVPR_ , 2005.
  * Ecoffet et al. [2021] Adrien Ecoffet, Joost Huizinga, Joel Lehman, Kenneth O. Stanley, and Jeff Clune.  First return then explore.  In _Nature_ , 2021.
  * Eysenbach et al. [2019] Benjamin Eysenbach, Julian Ibarz, Abhishek Gupta, and Sergey Levine.  Diversity is all you need: Learning skills without a reward function.  In _ICLR_ , 2019.
  * Grill et al. [2020] Jean-Bastien Grill, Florian Strub, Florent Altché, Corentin Tallec, Pierre Richemond, Elena Buchatskaya, Carl Doersch, Bernardo Avila Pires, Zhaohan Guo, Mohammad Gheshlaghi Azar, Bilal Piot, koray kavukcuoglu, Remi Munos, and Michal Valko.  Bootstrap your own latent - a new approach to self-supervised learning.  In _NeurIPS_ , 2020.
  * Haarnoja et al. [2018] Tuomas Haarnoja, Aurick Zhou, Pieter 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, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson.  Learning latent dynamics for planning from pixels.  In _ICML_ , 2019.
  * Hafner et al. [2020] Danijar Hafner, Timothy Lillicrap, Jimmy Ba, and Mohammad Norouzi.  Dream to control: Learning behaviors by latent imagination.  In _ICLR_ , 2020.
  * 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 _ICLR_ , 2020.
  * He et al. [2020] Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross Girshick.  Momentum contrast for unsupervised visual representation learning.  In _CVPR_ , 2020.
  * Hinton et al. [2012] Geoffrey Hinton, Nitish Srivastava, and Kevin Swersky.  Neural networks for machine learning. lecture 6d - a separate, adaptive learning rate for each connection.  2012\.
  * Jinnai et al. [2020] Yuu Jinnai, Jee Won Park, Marlos C. Machado, and George Konidaris.  Exploration in reinforcement learning with deep covering options.  In _ICLR_ , 2020.
  * Kamienny et al. [2021] Pierre-Alexandre Kamienny, Jean Tarbouriech, Alessandro Lazaric, and Ludovic Denoyer.  Direct then diffuse: Incremental unsupervised skill discovery for state covering and goal reaching.  In _ICML, Workshop on Unsupervised Reinforcement Learning_ , 2021\.
  * Kingma and Ba [2014] Diederik P Kingma and Jimmy Ba.  Adam: A method for stochastic optimization.  _arXiv:1412.6980_ , 2014.
  * Laskin et al. [2020] Michael Laskin, Aravind Srinivas, and Pieter Abbeel.  Curl: Contrastive unsupervised representations for reinforcement learning.  In _ICML_ , 2020.
  * Li et al. [2021] Siyuan Li, Lulu Zheng, Jianhao Wang, and Chongjie Zhang.  Learning subgoal representations with slow dynamics.  In _ICLR_ , 2021.
  * Lillicrap et al. [2015] 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.  _arXiv:1509.02971_ , 2015.
  * Machado [2019] M Cholodovskis Machado.  _Efficient Exploration in Reinforcement Learning through Time-Based Representations_.  PhD thesis, University of Alberta, 2019.
  * Machado et al. [2017] Marlos C. Machado, Marc G. Bellemare, and Michael Bowling.  A laplacian framework for option discovery in reinforcement learning.  _ICML_ , 2017.
  * Machado et al. [2018] Marlos C. Machado, Clemens Rosenbaum, Xiaoxiao Guo, Miao Liu, Gerald Tesauro, and Murray Campbell.  Eigenoption discovery through the deep successor representation.  _ICLR_ , 2018.
  * Machado et al. [2021] Marlos C. Machado, Andre Barreto, and Doina Precup.  Temporal abstraction in reinforcement learning with the successor representation, 2021.  arXiv 2110.05740.
  * Mahadevan [2005] Sridhar Mahadevan.  Proto-value functions: Developmental reinforcement learning.  _ICML_ , 2005.
  * Mahadevan and Maggioni [2007] Sridhar Mahadevan and Mauro Maggioni.  Proto-value functions: A laplacian framework for learning representation and control in markov decision processes.  _Journal of Machine Learning Research_ , 2007.
  * Mazoure et al. [2020] Bogdan Mazoure, Remi Tachet des Combes, Thang Long Doan, Philip Bachman, and R Devon Hjelm.  Deep reinforcement and infomax learning.  In _NeurIPS_ , 2020.
  * Misra et al. [2020] Dipendra Misra, Mikael Henaff, Akshay Krishnamurthy, and John Langford.  Kinematic state abstraction and provably efficient rich-observation reinforcement learning.  In _ICML_ , 2020.
  * Mnih et al. [2015] Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A. Rusu, Joel Veness, Marc G. Bellemare, Alex Graves, Martin A. Riedmiller, Andreas Fidjeland, Georg Ostrovski, Stig Petersen, Charles Beattie, Amir Sadik, Ioannis Antonoglou, Helen King, Dharshan Kumaran, Daan Wierstra, Shane Legg, and Demis Hassabis.  Human-level control through deep reinforcement learning.  _Nature_ , 2015.
  * Mnih et al. [2016] Volodymyr Mnih, Adria Puigdomenech Badia, Mehdi Mirza, Alex Graves, Timothy Lillicrap, Tim Harley, David Silver, and Koray Kavukcuoglu.  Asynchronous methods for deep reinforcement learning.  In _Machine Learning Research_ , 2016.
  * Nasiriany et al. [2019] Soroush Nasiriany, Vitchyr Pong, Steven Lin, and Sergey Levine.  Planning with goal-conditioned policies.  In _NeurIPS_ , 2019.
  * Oudeyer and Kaplan [2009] Pierre-Yves Oudeyer and Frederic Kaplan.  What is intrinsic motivation? a typology of computational approaches.  _Frontiers in neurorobotics_ , 2009.
  * Pong et al. [2020] Vitchyr H Pong, Murtaza Dalal, Steven Lin, Ashvin Nair, Shikhar Bahl, and Sergey Levine.  Skew-fit: State-covering self-supervised reinforcement learning.  In _ICML_ , 2020.
  * Schmidhuber [1991] Jürgen Schmidhuber.  A possibility for implementing curiosity and boredom in model-building neural controllers.  In _International Conference on Simulation of Adaptive Behavior on From Animals to Animats_ , 1991.
  * Schrittwieser et al. [2020] Julian Schrittwieser, Ioannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Simon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, et al.  Mastering atari, go, chess and shogi by planning with a learned model.  _Nature_ , 2020.
  * Stachenfeld et al. [2014] Kimberly L Stachenfeld, Matthew Botvinick, and Samuel J Gershman.  Design principles of the hippocampal cognitive map.  In _NeurIPS_ , 2014.
  * Stooke et al. [2021] Adam Stooke, Kimin Lee, Pieter Abbeel, and Michael Laskin.  Decoupling representation learning from reinforcement learning.  In _ICML_ , 2021.
  * Sutton et al. [1999] Richard S. Sutton, Doina Precup, and Satinder Singh.  Between MDPs and semi-MDPs: A framework for temporal abstraction in reinforcement learning.  _Artificial Intelligence_ , 1999.
  * Todorov et al. [2012] Emanuel Todorov, Tom Erez, and Yuval Tassa.  Mujoco: A physics engine for model-based control.  In _International Conference on Intelligent Robots and Systems_ , 2012\.
  * Trefethen and Bau [1997] Lloyd Nicholas Trefethen and David Bau.  _Numerical Linear Algebra_.  SIAM, Philadelphia, 1997.
  * Vezhnevets et al. [2017] Alexander Sasha Vezhnevets, Simon Osindero, Tom Schaul, Nicolas Heess, Max Jaderberg, David Silver, and Koray Kavukcuoglu.  Feudal networks for hierarchical reinforcement learning.  In _ICML_ , 2017.
  * Watter et al. [2015] Manuel Watter, Jost Springenberg, Joschka Boedecker, and Martin Riedmiller.  Embed to control: A locally linear latent dynamics model for control from raw images.  In _NeurIPS_ , 2015.
  * Wu et al. [2019] Yifan Wu, George Tucker, and Ofir Nachum.  The Laplacian in RL: Learning representations with efficient approximations.  In _ICLR_ , 2019.
  * Yarats et al. [2021] Denis Yarats, Rob Fergus, Alessandro Lazaric, and Lerrel Pinto.  Reinforcement learning with prototypical representations.  In _ICML_ , 2021.
  * Zhang et al. [2021] Tianjun Zhang, Paria Rashidinejad, Jiantao Jiao, Yuandong Tian, Joseph E. Gonzalez, and Stuart Russell.  MADE: Exploration via maximizing deviation from explored regions.  In _NeurIPS_ , 2021.

##  Appendix A Temporal abstractions augmented temporally-contrastive learning
(TATC) in the non-uniform prior setting

The proposed approach consists in simultaneously training the representation
$\phi$ and the hierarchical agent $(\pi_{\text{low}},\pi_{\text{hi}})$. The
idea is to progressively extend the explored area while maintaining the
previously collected knowledge. To do so, in the non-uniform-prior setting,
the agent switches with some probability $p_{rw}$ between following a
uniformly random policy $\pi_{\mu}$ and executing the hierarchical policy
(skills). The latter helps reach further areas more efficiently, where data
collected by $\pi_{\mu}$ is used to train the representation $\phi$. Along
their training, the skills progressively extend to cover newly discovered
areas. Algorithm 1 provides a pseudocode of the proposed approach, in the non-
uniform-prior setting.

Algorithm 1 TATC in the non-uniform prior setting

1: Input: $L,c,p_{rw},N$

2: for $iteration=1,2,\ldots$ do

3: $D_{\pi_{\mu}}=\emptyset$, $D_{s}=\emptyset$

4: for $batch=1,2,\ldots,N$ do

5: Reset to $s_{0}$ with probability $p_{r}$.

6: $p\sim$ _Unif_([0,1])

7: if $p<p_{rw}$ then

8: Run the uniformly random policy $\pi_{\mu}$ to collect $L$ random walk
trajectories $\\{\tau^{\prime}_{i}\\}_{i=1}^{L}$ of $c$ steps each.

9: $D_{\pi_{\mu}}\leftarrow
D_{\pi_{\mu}}\cup\\{\tau^{\prime}_{k}\\}_{k=1}^{L}$

10: else

11: Run $(\pi_{\text{hi}},\pi_{\text{low}})$ to collect $L$ consecutive
skills’ trajectories $\\{(\tau_{k},\bm{\delta}_{k})\\}_{k=1}^{L}$ and their
corresponding directions

12: $D_{s}\leftarrow D_{s}\cup\\{(\tau_{k},\bm{\delta}_{k})\\}_{k=1}^{L}$

13: end if

14: end for

15: Optimize the policies $(\pi_{\text{hi}},\pi_{\text{low}})$ using the
intrinsic objectives 3 and 4

16: Optimize $\phi$ so as to minimize
$\mathcal{L}_{\text{TATC}}(\phi;\mathcal{D}_{s},\mathcal{D}_{\pi_{\mu}})$
(Equation 6).

17: end for

##  Appendix B Representation Objective Augmentation: Ablation Study

###  B.1 Boredom augmentation helps exploration

In order to illustrate the importance of the boredom term $\mathcal{B}$ – in
the final objective (6) we conducted the same representation learning
experiments for the three gridworld domains in the non-uniform prior setting,
but this time with the non-augmented representation learning objective
($\beta^{\prime}=0$).

Figure 9 shows how the agent failed at exploring the whole domain. In T-Maze,
it focuses only on one corridor without getting curious about the other one.
Regarding U-Maze and 4-rooms, the agent stops exploring after discovering the
end of the first corridor and the second room, respectively. This is due to
the lack of incentive to visit the yet unseen states, as they are less
rewarding for $\pi_{\text{hi}}$ (i.e., closer in the representation space,
hence smaller $R^{\text{hi}}$) than the furthest explored state. The effect of
the proposed augmentation compresses the representation of the explored area,
say the first corridor in U-Maze, which makes the rest of the environment more
appealing to explore for $\pi_{\text{hi}}$ (i.e. further in the representation
space, hence larger $R^{\text{hi}}$).

(a) U-Maze

(b) T-Maze

(c) 4-rooms

Figure 9: Learned representations in the gridworld domains with the _non-
augmented_ objective. Without the boredom term, the agent fails to cover the
state space (cf. Figure 3), and may settle for incomplete representations. The
colors reflect the distances in terms of the dynamics. They can be seen as
quantities proportional to the length of the shortest path from the $s_{0}$
(marked in red) to the represented state.

###  B.2 Boredom augmentation enforces dynamics-awareness

To verify the benefit of the boredom term beyond helping exploration, we train
the representation with the non-augmented objective ($\beta^{\prime}=0$) but
this time in the uniform prior setting, so that to marginalize the exploration
problem. Figure 10 illustrates the learned representations in the three
gridworld domains. These representations have failed to capture the dynamics.
For example, in the case of 4-rooms, the distances from the first room to the
fourth and third rooms are comparable in the representation space, which
indicates that the representation does not take into account the relative
order in which the rooms should be visited, when moving from the first room to
the last. Similarly, in U-Maze, the end of the maze is closer to the initial
area than the second corner is. However, in order to reach the former on must
pass by the latter. This proves that the boredom term is not only important
for the desired exploratory behavior (cf. Figure 9), but also enhances the
dynamics-awareness of our representation.

(a) U-Maze

(b) T-Maze

(c) 4-rooms

Figure 10: Learned representations when uniformly sampling over the state
space. Without the boredom term, the representation does not reflect
temporally-extended dynamics. The colors reflect the distances in terms of the
dynamics. They can be seen as quantities proportional to the length of the
shortest path from the $s_{0}$ (marked in red) to the represented state.

##  Appendix C Implementation details

###  C.1 GridWorld

The states are one-hot encoded such that no positional information is provided
to the agent. The domains dimensions are: U-Maze $30\times 30$, T-Maze
$40\times 30$, 4-rooms $21\times 21$.

For all the experiments, we defined the representation network as an MLP of
two hidden layers of size $128$ with tanh activations and a linear output
layer of the size of representation’s dimensionality $d$. The high-level and
the low-level policies are both MLPs of two hidden layers of size $128$ with
tanh activations and a logsoftmax output layer of the size of their respective
action spaces: the environment’s $4$ actions for the low-level policy and $8$
actions for the high-level policy corresponding to the $8$ directions
$\Omega=\\{(\cos(2k\pi/n),\sin(2k\pi/n))\ |\leavevmode\nobreak\
k\in\\{0,...,7\\}\\}$ that define diverse skills.

The policies were trained with vanilla A2C with MC returns from the collected
trajectories (Monte-Carlo estimates), i.e. no bootstrapped values where used.
The skills being of a fixed size they could be trained without any reward
discount ($\gamma=1$). The high-level and low-level policies were entropy-
regularized with coefficients $0.3$ and $0.1$ respectively.

All of these networks were trained with RMSprop [Hinton et al., 2012] and a
step size of $0.001$. Environments specific hyperparameters are provided
below.

####  C.1.1 Representation Learning

#####  U-Maze.

Our representation is learned in the non-uniform prior setting with
$p_{r}{=}0.3$, $p_{rw}{=}0.4$ and $K{=}90$ (around the number of steps between
$s_{0}$ and the furthest state in the maze). We learn a 2-dimensional
representation ($d=2$) using the representation learning objective 6 with
$\beta=0.2$ and $\beta^{\prime}=2$. We fix the skills length to $c=30$ steps
(so $L=K/c=3$), and jointly train the representation $\phi$ and the policies
$(\pi_{\text{hi}},\pi_{\text{low}})$ by collecting, for each update, a batch
of $N=32$ trajectories of length $c$ to fill $D_{s}$ and $D_{\pi_{\mu}}$ as
described in Algorithm 1. We train them for $700$ epochs where each epoch
corresponds to $10$ updates (convergence to the complete representation
required around $500$ epochs).

#####  T-Maze.

Our representation is learned in the non-uniform prior setting with
$p_{r}{=}0.2$, $p_{rw}{=}0.4$ and $K{=}40$ (around the number of steps between
$s_{0}$ and the furthest state in the maze). We learn a 2-dimensional
representation ($d=2$) using the representation learning objective 6 with
$\beta=0.2$ and $\beta^{\prime}=2$. We fix the skills’ length to $c=20$ steps
(so $L=K/c=2$). and jointly train the representation $\phi$ and the policies
$(\pi_{\text{hi}},\pi_{\text{low}})$ by collecting, for each update, a batch
of $N=32$ trajectories of length $c$ to fill $D_{s}$ and $D_{\pi_{\mu}}$ as
described in Algorithm 1. We train them for $700$ epochs where each epoch
corresponds to $10$ updates (convergence to the complete representation
required around $350$ epochs).

#####  4-rooms.

Our representation is learned in the non-uniform prior setting with
$p_{r}{=}0.25$, $p_{rw}{=}0.5$ and $K{=}60$ (around the number of steps
between $s_{0}$ and the furthest state in the maze). We learn a 2-dimensional
representation ($d=2$) using the representation learning objective 6 with
$\beta=0.2$ and $\beta^{\prime}=2$. We fix the skills’ length to $c=20$ steps
(so $L=K/c=3$). and jointly train the representation $\phi$ and the policies
$(\pi_{\text{hi}},\pi_{\text{low}})$ by collecting, for each update, a batch
of $N=32$ trajectories of length $c$ to fill $D_{s}$ and $D_{\pi_{\mu}}$ as
described in Algorithm 1. We train them for $700$ epochs where each epoch
corresponds to $10$ updates (convergence to the complete representation
required around $350$ epochs).

##### The Laplacian representation

(Lap-rep) was trained in the same environments’ settings described above, for
both the uniform and non-uniform prior settings (of course no policy is
trained here so $p_{rw}=1$, and $(s_{0},p_{r})$ are not relevant for the
uniform prior setting). Besides the representation’s dimension $d$, we used
the training configuration and hyperparameters proposed by Wu et al. [2019].
For the uniform prior setting, our online data collection does not cause any
discrepancy compared to the offline scheme used in Wu et al. [2019]. Indeed,
for a minibatch size large enough, the stochastic minibatch based training of
Lap-rep when using a uniform prior is agnostic to the data collection scheme
(offline vs online) since in both cases the minibatches are sampled from the
exact same state distribution.

####  C.1.2 Prediction and Control

In the prediction and control experiments, we evaluate each pretrained
representation by training an actor-critic agent to solve a goal-achieving
task with a sparse reward ($r=1$ upon reaching the goal, and $r=0$ otherwise).
Here are the set goal positions: $(1,30)$ in U-Maze, $(25,30)$ in T-Maze and
$(1,21)$ in 4-rooms. The episode size was set to $100$ steps for all the
gridworld domains.

For the prediction, the critic head is a linear function in the given
representation, while the actor is an MLP with two hidden layers of size $64$
and tanh activations, a logsoftmax output layer of size $4$ (discrete
gridworld actions), and the actor’s input is the state one-hot code. For the
control experiments, the actor-critic agent is defined on top of the
representation as an MLP of two hidden layers of size $64$ with tanh
activations that feed two output heads: a linear critic head, and a logsoftmax
action head for the $4$ actions. The agent is trained with A2C with MC returns
and a discount of $\gamma=0.98$, a batchsize of $80$ episodes, an entropy
regularization with a $0.01$ coefficient, and Adam optimizer [Kingma and Ba,
2014] with a learning rate of $0.001$.

###  C.2 MuJoCo: AntMaze

The Ant agent has a 29 dimensional state space and a 8 dimensional action
space (4 legs with 2 joints each to control). For the sake of simplifying the
RL training algorithm, we mapped each action-dimension interval to a discrete
set of 5 values equally spaced over this interval.

We used 2 mazes similar in shape to those from Wu et al. [2019] (see Figure
11): AntMaze-1 defines a 3D U-shaped corridor and AntMaze-2 is a 3D swirl-
shaped corridor.

(a) AntMaze-1

(b) AntMaze-2

Figure 11: Antmaze domains. Goal positions of the evaluation tasks are shown
in green.

We used the same architectures for the representation and the policies as for
the gridworld, with the only difference that for the low-level policy, the
action head was adapted to the discretization of the action space by having
$8$ logsoftmax output heads of size $5$, one for each action dimension, and
the corresponding $5$ discrete values. This choice makes the training
algorithm simpler as it allows using A2C here as well.

Our representation is learned in the non-uniform prior setting with
$p_{r}=0.2$, $p_{rw}=0.3$ and $K=500$. We learn a 2-dimensional representation
($d=2$) using the representation learning objective 6 with $\beta=0.2$ and
$\beta^{\prime}=5$. We fixed their length to $c=100$ steps (so $L=K/c=5$). and
jointly train the representation $\phi$ and the policies
$(\pi_{\text{hi}},\pi_{\text{low}})$ by collecting, for each update, a batch
of $N=32$ trajectories of length $c$ to fill $D_{s}$ and $D_{\pi_{\mu}}$, as
described in Algorithm 1. We train them for $1000$ epochs where each epoch
corresponds to $10$ updates (convergence to the complete representation
required around $650$ epochs).

The policies were trained with the same A2C used in gridworld domains and the
same RMSprop hyperparameters. The high-level and low-level policies were
entropy-regularized with the coefficients $0.15$ and $0.1$, respectively.

For the reward shaping and skills evaluation experiments, the goal positions
are shown in Figure 11, and success is defined as being within $\epsilon$ from
the goal. Here, $\epsilon=2$ which corresponds to half the size of the
building blocks of the mazes.

####  C.2.1 Reward Shaping

For these experiments, Lap-rep was trained in both the uniform and the non-
uniform prior setting.

For the uniform prior setting, we used $d=2$ and followed the experimental
framework of Wu et al. [2019]. Since our AntMaze environments are larger, we
collected $500,000$ training samples (10 times more than in Wu et al. [2019])
from a uniformly random policy, then we trained the representation on this
large dataset. For all other hyperparameters, we used those provided in Wu et
al. [2019]. With $d=20$, our replication of Lap-rep did not succeed in reward
shaping.

Regarding the non-uniform prior setting, we used the same setting
configuration as for TATC, with the representation objective from Wu et al.
[2019]. We have tested online (similar to TATC) and offline [Wu et al., 2019]
data collection for the representation training, and $d\in\\{2,20\\}$. Both
schemes ended up performed the same way for the reward shaping task.

Now, for the reward shaping, we train a Soft Actor-Critic (SAC) [Haarnoja et
al., 2018] agent to reach a goal area (neighbourhood around the goal state)
with episodes of size $1000$ steps. We use the following hyperparameters:

  * •

Discount $\gamma=0.99$

  * •

Entropy coefficient (temperature) $\alpha=0.1$

  * •

Soft critic updates with smoothing constant $\tau=0.005$

  * •

Replay buffer of size $5\cdot 10^{6}$ (equal to the number of training steps).

  * •

Adam optimizer with step size of $0.0001$

As SAC is sensitive to the reward scale [Haarnoja et al., 2018], we grid-
searched this hyperparameter in $\\{10^{-5},10^{-4},\cdots,1,2,10,20\\}$, and
the best performing one for our representation was $1$ in AntMaze-1 and $0.01$
in AntMaze-2. Regarding Lap-rep, we found that $10$ and $0.01$ worked the best
for these two mazes, respectively. All these coefficients correspond to the
dense reward shaping setting. Their values were doubled for the half-half mix
reward setting, to account for the $0.5$ coefficient.

####  C.2.2 Skills Evaluation

To train DCO, we first collect a dataset to estimate the second eigenvector
and then use the same dataset to train a policy – the option – using DDPG
[Lillicrap et al., 2015]. Each DCO option is tied to its own eigenvector
estimate and its own training set of size $500000$ (10 times the size used in
Jinnai et al. [2020]). As suggested by the authors of DCO [Jinnai et al.,
2020], the remaining hyperparameters to estimate the eigenvectors and train
their corresponding options were taken from Wu et al. [2019]. DIAYN skills
were trained as recommended by Eysenbach et al. [2019]. For fair comparison,
we train $8$ skills for both DCO and DIAYN.

For the skills evaluation stage, we freeze the learned low-level policies and
train a high-level policy to use the $8$ skills as the only available actions
to reach the goal $g$ on the other end of the AntMaze-1 environment using a
sparse reward $r_{t}=\mathds{1}\left[\|s_{t+1}-g\|_{2}\leq\epsilon\right]$
within a finite horizon of $1000$ steps. Note that this tasks is quite
challenging given the type of reward and the length of episode especially in a
continuous state space. As our skills offer some flexibility in their
execution (can be started everywhere and run for arbitrary number of steps),
this episode length was decomposed to $5$ skills of $200$ steps each. The
high-level policy was trained with A2C with MC returns (no discount) a batch
size of $8$ episodes, and RMSprop optimizer with a learning rate of $0.001$.

##  Appendix D The switching utility of the boredom term

Note that $\mathcal{D}_{s}$, in Equation 5, may contain trajectories from
skills that are not yet duly trained; for example early in the training or in
a freshly discovered area. Since at that stage, these skills’ trajectories are
close to random walks, their contribution in the boredom term is similar to
the first attractive term, in Equation 2. This means that a new skill
trajectory initially contributes to the temporal similarity term (attractive
term) in training the representation, thus making the most out of the sampled
skills’ trajectories while these are still early in their training. The more a
skill is trained, the more structured its trajectories become and the more
they contribute to the intended "boredom" effect (Section 3.4), that is
encouraging exploration and dynamics awareness (Appendix B).

##  Appendix E Connection to Behavioral Mutual Information

There are numerous methods in the _intrinsic control_ literature that aim at
maximizing the mutual information between the agent’s behavior and a
conditioning variable that encodes the available skills.

This type of intrinsic control is achieved by training a skill conditioned
policy, $\pi$, to maximize the mutual information between the skill code, $z$,
and some representation of the trajectory, $\tau$, obtained from the
conditioned policy $\pi(\cdot|z)$. This objective can be written as:

$I(z;f(\tau))=\mathcal{H}(z)-\mathcal{H}(z|f(\tau)),$ (7)

where entropy is denoted by $\mathcal{H}$, and $f$ is a function of the
trajectory. It is also common to assume that $z$ is sampled from a fixed prior
[Eysenbach et al., 2019], which simplifies this policy training objective as a
minimization of the conditional entropy term in Equation 7. In practice, the
adopted training loss, can be derived as a lower bound of this quantity, using
an approximate posterior, $q$:

$L_{q}(\pi)=-\mathbb{E}_{z,\pi}[\log q(z|f(\tau))].$ (8)

Traditionally, the integrand of this expectation defines the intrinsic reward
of the skill conditioned policy, while the approximate posterior $q$ is
trained to discriminate the correct $z$ based on the observed behavior $\tau$.

Inspired by the fast inference offered by Successor Features, Hansen et al.
[2020] proposed to use a _log-linear_ discriminator in the successor
representation (SR), $\phi_{\text{SR}}$. In other words, the skill rewards can
be written as:

$r(s)=\phi_{\text{SR}}(s)^{\top}\mathbf{w},$ (9)

with $\mathbf{w}$ playing the role of the skill-identifying variable (denoted
by $z$ in the general case above). Note that in this case, the function $f$
maps, as it is commonly the case, to the final state of the trajectory. Now,
consider the case where the trajectory is instead represented by the
(normalized) latent direction of the final transition $(s,s^{\prime})$. This
reward would be

$r(s,s^{\prime})=\frac{(\phi_{\text{SR}}(s^{\prime})-\phi_{\text{SR}}(s))^{\top}\mathbf{w}}{\|\phi_{\text{SR}}(s^{\prime})-\phi_{\text{SR}}(s)\|_{2}}.$
(10)

Let’s recall that the SR555The SR is encoded by the matrix $(I-\gamma
T)^{-1}$, with $T$ the MDP’s transition matrix. shares the same eigenvectors
as the Laplacian [Stachenfeld et al., 2014, Machado et al., 2018]. This
implies that it can also be approximated with a temporally-contrastive
objective [Wu et al., 2019], and potentially replaced by our alternative TATC
representation $\phi$. Finally, we can rewrite Equation 10 as

$r(s,s^{\prime})=\frac{(\phi(s^{\prime})-\phi(s))^{\top}\mathbf{w}}{\|\phi(s^{\prime})-\phi(s)\|_{2}}$
(11)

This reward corresponds to our skills intrinsic reward from Equation 3, with
$\mathbf{w}\equiv\bm{\delta}$.