#  Goal-Conditioned Reinforcement Learning with Imagined Subgoals

Elliot Chane-Sane  Cordelia Schmid  Ivan Laptev

###### Abstract

Goal-conditioned reinforcement learning endows an agent with a large variety
of skills, but it often struggles to solve tasks that require more temporally
extended reasoning. In this work, we propose to incorporate imagined subgoals
into policy learning to facilitate learning of complex tasks. Imagined
subgoals are predicted by a separate high-level policy, which is trained
simultaneously with the policy and its critic. This high-level policy predicts
intermediate states halfway to the goal using the value function as a
reachability metric. We don’t require the policy to reach these subgoals
explicitly. Instead, we use them to define a prior policy, and incorporate
this prior into a KL-constrained policy iteration scheme to speed up and
regularize learning. Imagined subgoals are used during policy learning, but
not during test time, where we only apply the learned policy. We evaluate our
approach on complex robotic navigation and manipulation tasks and show that it
outperforms existing methods by a large margin.

Reinforcement Learning, Robotics

##  1 Introduction

Figure 1:  Illustration of the KL-regularized policy learning using imagined
subgoals. (Left): The policy fails to reach a distant goal, yet it can reach a
closer subgoal. Our approach automatically generates imagined subgoals for a
task and uses such subgoals to direct the policy search during training.
(Right): At test time, the resulting flat policy can reach arbitrarily distant
goals without relying on subgoals.

An intelligent agent aims at solving tasks of varying horizons in an
environment. It must be able to identify how the different tasks intertwine
with each other and leverage behaviors learned by solving simpler tasks to
efficiently master more complex tasks. For instance, once a legged robot has
learned how to walk in every direction in a simulated maze environment, it
could use these behaviors to efficiently learn to navigate in this
environment. Goal-conditioned reinforcement learning (RL) (Kaelbling, 1993;
Schaul et al., 2015; Andrychowicz et al., 2017) defines each task by the
desired goal and, in principle, could learn a wide range of skills. In
practice, however, reinforcement learning struggles to perform temporally
extended reasoning.

Hierarchical methods have proven effective for learning temporal abstraction
in long-horizon problems (Dayan & Hinton, 1993; Wiering & Schmidhuber, 1997;
Sutton et al., 1999; Dietterich, 2000). In the goal-conditioned setting, a
high-level controller typically finds an appropriate sequence of subgoals that
can be more easily followed by a low-level policy (Nachum et al., 2018; Gupta
et al., 2019). When chosen appropriately, these subgoals effectively decompose
a complex task into easier tasks. However, hierarchical methods are often
unstable to train (Nachum et al., 2018) and rely on appropriate temporal
design choices.

In this work, we propose to use subgoals to improve the goal-conditioned
policy. Instead of reaching the subgoals explicitly, our method builds on the
following intuition. If the current policy can readily reach a subgoal, it
could provide guidance for reaching more distant goals, as illustrated in
Figure 1. We apply this idea to all possible state and goal pairs of the
environment. This self-supervised approach progressively extends the horizon
of the agent throughout training. At the end of training, the resulting flat
policy does not require access to subgoals and can reach distant goals in its
environment.

Policy $\pi$Prior policyHigh-level policy $\pi^{H}$$\pi(.|s,s_{g})$state
$s$goal $g$subgoals $s_{g}$$\pi(.|s,g)$$\pi^{prior}(.|s,g)$$D_{\mathrm{KL}}$
Figure 2:  Overview over our Reinforcement learning with Imagined Subgoals
(RIS) approach. During policy training, the policy $\pi$ is constrained to
stay close to the prior policy $\pi^{prior}$ through KL-regularization. We
define the prior policy $\pi^{prior}$ as the distribution of actions required
to reach intermediate subgoals $s_{g}$ of the task. Given the initial state
$s$ and the goal state $g$, the subgoals are generated by the high-level
policy $\pi^{H}$. Note that the high-level policy and subgoals are only used
during training of the target policy $\pi$. At test time we directly use $\pi$
to generate appropriate actions.

Our method does not require a full sequence of subgoals, but predicts subgoals
that are halfway to the goal, such that reaching them corresponds to a lower
level of temporal abstraction than reaching the goal. To handle this subgoal
prediction problem, we simultaneously train a separate high-level policy
operating in the state space and use the value function of the goal-
conditioned policy as a relative measure of distance between states (Eysenbach
et al., 2019; Nasiriany et al., 2019).

To incorporate subgoals into policy learning, we introduce a prior policy
defined as the distribution over actions required to reach intermediate
subgoals (Figure 2). When using appropriate subgoals that are easier to reach,
this prior policy corresponds to an initial guess to reach the goal.
Accordingly, we leverage a policy iteration with an additional Kullback-
Leibler (KL) constraint scheme towards this prior policy. This adequately
encourages the policy to adapt the simpler behaviors associated with reaching
these subgoals to master more complex goal-reaching tasks in the environment.
Because these subgoals are not actively pursued when interacting with the
actual environment but only used to accelerate policy learning, we refer to
them as imagined subgoals. Imagined subgoals and intermediate subgoals have
equal meaning in this paper.

Our method, Reinforcement learning with Imagined Subgoals (RIS), builds upon
off-policy actor-critic approaches for continuous control and additionally
learns a high-level policy that updates its predictions according to the
current goal-reaching capabilities of the policy. Our approach is general,
solves temporally extended goal-reaching tasks with sparse rewards, and self-
supervises its training by choosing appropriate imagined subgoals to
accelerate its learning.

In summary, our contributions are threefold: (1) we propose a method for
predicting subgoals which decomposes a goal-conditioned task into easier
subtasks; (2) we incorporate these subgoals into policy learning through a KL-
regularized policy iteration scheme with a specific choice of prior policy;
and (3) we show that our approach greatly accelerates policy learning on a set
of simulated robotics environments that involve motor control and temporally
extended reasoning.111Code is available on the project webpage
https://www.di.ens.fr/willow/research/ris/.

##  2 Related Work

Goal-conditioned reinforcement learning has been addressed by a number of
methods (Kaelbling, 1993; Schaul et al., 2015; Andrychowicz et al., 2017;
Veeriah et al., 2018; Pong et al., 2018; Nair et al., 2018; Zhao et al., 2019;
Pitis et al., 2020; Eysenbach et al., 2020). Given the current state and the
goal, the resulting policies predict action sequences that lead towards the
desired goal. Hindsight experience replay (HER) (Kaelbling, 1993; Andrychowicz
et al., 2017) is often used to improve the robustness and sample efficiency of
goal-reaching policies. While in theory such policies can address any goal-
reaching task, they often fail to solve temporally extended problems in
practice (Levy et al., 2019; Nachum et al., 2018).

Long-horizon tasks can be addressed by hierarchical reinforcement learning
(Dayan & Hinton, 1993; Wiering & Schmidhuber, 1997; Dietterich, 2000; Levy et
al., 2019; Vezhnevets et al., 2017). Such methods often design high-level
policies that operate at a coarser time scale and control execution of low-
level policies. To address goal-conditioned settings, high-level policies can
be learned to iteratively predict a sequence of intermediate subgoals. Such
subgoals can then be used as targets for low-level policies (Nachum et al.,
2018; Gupta et al., 2019; Nair & Finn, 2020). As an alternative to the
iterative planning, other methods generate sequences of subgoals with a
divide-and-conquer approach (Jurgenson et al., 2020; Parascandolo et al.,
2020; Pertsch et al., 2020b). While hierarchical RL methods can better address
long-horizon tasks, the joint learning of high-level and low-level policies
may imply instabilities (Nachum et al., 2018). Similar to previous
hierarchical RL methods, we use subgoals to decompose long-horizon tasks into
simpler problems. Our subgoals, however, are only used during policy learning
to guide and accelerate the search of the non-hierarchical policy.

Several recent RL methods use the value function of goal-reaching policies to
measure distances between states and to plan sequences of subgoals (Nasiriany
et al., 2019; Eysenbach et al., 2019; Zhang et al., 2020). In particular, LEAP
(Nasiriany et al., 2019) uses the value function of TDM policies (Pong et al.,
2018) and optimizes sequences of appropriate subgoals at test time. Similar to
previous methods, we use the value function as a distance measure between
states. Our method, however, avoids expensive test-time optimization of
subgolas. We also experimentally compare our method with LEAP and demonstrate
improvements. Moreover, we show that our approach can benefit from recent
advances in learning representations for reinforcement learning from pixels on
a vision-based robotic manipulation task (Kostrikov et al., 2020; Laskin et
al., 2020).

Many approaches cast reinforcement learning as a probabilistic inference
problem where the optimal policy should match a probability distribution in a
graphical model defined by the reward and the environment dynamics (Toussaint,
2009; Kappen et al., 2012; Levine, 2018). Several methods optimize an
objective incorporating the divergence between the target policy and a prior
policy. The prior policy can be fixed (Haarnoja et al., 2017; Abdolmaleki et
al., 2018b; Haarnoja et al., 2018a; Pertsch et al., 2020a) or learned jointly
with the policy (Teh et al., 2017; Galashov et al., 2019; Tirumala et al.,
2019). While previous work imposes explicit priors e.g., in the multi-task and
transfer learning settings (Teh et al., 2017; Galashov et al., 2019; Tirumala
et al., 2019), we constrain our policy search by the prior distribution
implicitly defined by subgoals produced by the high-level policy.

Behavior priors have often been used to avoid value overestimation for out-of-
distribution actions in offline reinforcement learning (Fujimoto et al., 2018;
Wu et al., 2019; Kumar et al., 2019; Siegel et al., 2020; Nair et al., 2020;
Wang et al., 2020). Similar to these methods, we constrain our high-level
policy to avoid predicting subgoals outside of the valid state distribution.

Finally, our work shares similarities with guided policy search methods
(Levine & Koltun, 2013; Levine & Abbeel, 2014; Levine et al., 2016) which
alternate between generating expert trajectories using trajectory optimization
and improving the learned policy. In contrast, our policy search is guided by
subgoals produced by a high-level policy.

##  3 Method

Our method builds on the following key observation. If an action $a$ is well-
suited for approaching an intermediate subgoal $s_{g}$ from state $s$, it
should also be a good choice for approaching the final goal $g$ of the same
task. We assume that reaching subgoals $s_{g}$ is simpler than reaching more
distant goals $g$. Hence, we adopt a self-supervised strategy and use the
subgoal-reaching policy $\pi(\cdot|s,s_{g})$ as a guidance when learning the
goal-reaching policy $\pi(\cdot|s,g)$. We denote our approach as Reinforcement
learning with Imagined Subgoals (RIS) and present its overview in Figures 1
and 2.

To implement the idea of RIS, we first introduce a high-level policy $\pi^{H}$
predicting imagined subgoals $s_{g}$, as described in Section 3.2. Section 3.3
presents the regularized learning of the target policy $\pi$ using subgoals.
The joint training of $\pi$ and $\pi^{H}$ is summarized in Section 3.4. Before
describing our technical contributions, we present the actor-critic paradigm
used by RIS in Section 3.1 below.

###  3.1 Goal-Conditioned Actor-Critic

We consider a discounted, infinite-horizon, goal-conditioned Markov decision
process, with states $s\in\mathcal{S}$, goals $g\in\mathcal{G}$, actions
$a\in\mathcal{A}$, reward function $r(s,a,g)$, dynamics $p(s^{\prime}|s,a)$
and discount factor $\gamma$. The objective of a goal-conditioned RL agent is
to maximize the expected discounted return

$J(\pi)=\mathbb{E}_{g\sim\rho_{g},\tau\sim
d^{\pi}(.|g)}[\sum_{t}\gamma^{t}r(s_{t},a_{t},g)]$

under the distribution

$d^{\pi}(\tau|g)=\rho_{0}(s_{0})\prod_{t}\pi(a_{t}|s_{t},g)p(s_{t+1}|s_{t},a_{t})$

induced by the policy $\pi$ and the initial state and goal distribution. The
policy $\pi(.|s,g)$ in this work generates a distribution over continuous
actions $a$ conditioned on state $s$ and goal $g$. Many algorithms rely on the
appropriate learning of the goal-conditioned action-value function $Q^{\pi}$
and the value function $V^{\pi}$ defined as
$Q^{\pi}(s,a,g)=\mathbb{E}_{s_{0}=s,a_{0}=a,\tau\sim
d^{\pi}(.|g)}[\sum_{t}\gamma^{t}r(s_{t},a_{t},g)]$ and
$V^{\pi}(s,g)=\mathbb{E}_{a\sim\pi(.|s,g)}Q^{\pi}(s,a,g)$.

In this work we assume states and goals to co-exist in the same space, i.e.
$\mathcal{S}=\mathcal{G}$, where each state can be considered as a potential
goal. Moreover, we set the reward $r$ to $-1$ for all actions until the policy
reaches the goal.

We follow the standard off-policy actor-critic paradigm (Silver et al., 2014;
Heess et al., 2015; Mnih et al., 2016; Fujimoto et al., 2018; Haarnoja et al.,
2018a). Experience, consisting of single transition tuples
$(s_{t},a_{t},s_{t+1},g)$, is collected by the policy in a replay buffer $D$.
Actor-critic algorithms maximize return by alternating between policy
evaluation and policy improvement. During the policy evaluation phase, a
critic $Q^{\pi}(s,a,g)$ estimates the action-value function of the current
policy $\pi$ by minimizing the Bellman error with respect to the Q-function
parameters $\phi_{k}$:

$Q_{\phi_{k+1}}=\arg\min_{\phi}\frac{1}{2}\mathbb{E}_{(s_{t},a_{t},s_{t+1},g)\sim
D}[y_{t}-Q_{\phi}(s_{t},a_{t},g)]^{2}$ (1)

with the target value

$y_{t}=r(s_{t},a_{t},g)+\gamma\mathbb{E}_{a_{t+1}\sim\pi(.|s_{t},g)}Q_{\phi_{k}}(s_{t+1},a_{t+1},g).$

During the policy improvement phase, the actor $\pi$ is typically updated such
that the expected value of the current Q-function $Q^{\pi}$, or alternatively
the advantage $A^{\pi}(s,a,g)=Q^{\pi}(s,a,g)-V^{\pi}(s,g)$, under $\pi$ is
maximized:

$\pi_{\theta_{k+1}}=\arg\max_{\theta}\mathbb{E}_{(s,g)\sim
D,a\sim\pi_{\theta}(.|s,g)}[Q^{\pi}(s,a,g)].$ (2)

Reaching distant goals using delayed rewards may require expensive policy
search. To accelerate the training, we propose to direct the policy search
towards intermediate subgoals of a task. Our method learns to predict
appropriate subgoals by the high-level policy $\pi^{H}$. The high-level policy
is trained together with the policy $\pi$ as explained below.

Figure 3:  Given initial states $s$ and goal states $g$, our subgoals are
states on the middle of the path from $s$ to $g$. We measure distances between
states by the value function $|V^{\pi}(s_{1},s_{2})|$ corresponding to the
current policy $\pi$. We obtain a distribution of subgoals from the high-level
policy $s_{g}\sim\pi^{H}(.|s,g)$. We use subgoals only at the training to
regularize and accelerate the policy search.

###  3.2 High-Level Policy

We would like to learn a high-level policy $\pi^{H}(.|s,g)$ that predicts an
appropriate distribution of imagined subgoals $s_{g}$ conditioned on valid
states $s$ and goals $g$. Our high-level policy is defined in terms of the
policy $\pi$ and relies on the goal-reaching capabilities of $\pi$ as
described next.

#### Subgoal search with a value function.

We note that our choice of the reward function $r=-1$ implies that the norm of
the value function $V^{\pi}(s,g)$ corresponds to an estimate of the expected
discounted number of steps required for the policy to reach the goal $g$ from
the current state $s$. We therefore propose to use $|V^{\pi}(s^{i},s^{j})|$ as
a measure of the distance between any valid states $s^{i}$ and $s^{j}$. Note
that this measure depends on the policy $\pi$ and evolves with the
improvements of $\pi$ during training.

Reaching imagined subgoals of a task should be easier than reaching the final
goal. To leverage this assumption for policy learning, we need to find
appropriate subgoals. In this work we propose to define subgoals $s_{g}$ as
midpoints on the path from the current state $s$ to the goal $g$, see Figure
3. More formally, we wish $s_{g}$ (i) to have equal distance from $s$ and $g$,
and (ii) to minimize the length of the paths from $s$ to $s_{g}$ and from
$s_{g}$ to $g$. We can find subgoals that satisfy these constraints by using
our distance measure $|V^{\pi}(s^{i},s^{j})|$ and minimizing the following
cost $C_{\pi}(s_{g}|s,g)$

$C_{\pi}(s_{g}|s,g)=\max\left(|V^{\pi}(s,s_{g})|,|V^{\pi}(s_{g},g)|\right).$
(3)

However, naively minimizing the cost $C_{\pi}$ under the high-level policy
distribution, i.e.

$\pi^{H}_{k+1}=\arg\min_{\pi^{H}}\mathbb{E}_{(s,g)\sim
D,s_{g}\sim\pi^{H}(.|s,g)}[C_{\pi}(s_{g}|s,g)],$ (4)

may lead to undesired solutions where the high-level policy samples subgoals
outside the valid state distribution $p_{s}(.)$. Such predictions may, for
example, correspond to unfeasible robot poses, unrealistic images, or other
adversarial states which have low distance from $s$ and $g$ according to
$|V^{\pi}|$ but are unreachable in practice.

#### Predicting valid subgoals.

To avoid non-valid subgoals, we can additionally encourage the high-level
policy to stay close to the valid state distribution $p_{s}(.)$ using the
following KL-regularized objective:

$\begin{split}\pi^{H}_{k+1}&=\arg\max_{\pi^{H}}\mathbb{E}_{(s,g)\sim
D,s_{g}\sim\pi^{H}(.|s,g)}\left[A^{\pi^{H}_{k}}(s_{g}|s,g)\right]\\\ &\text{
s.t.
}D_{\mathrm{KL}}\left(\pi^{H}(.|s,g)\left|\right|p_{s}(.)\right)\leq\epsilon,\end{split}$
(5)

where the advantage

$A^{\pi^{H}_{k}}(s_{g}|s,g)=\mathbb{E}_{\hat{s_{g}}\sim\pi^{H}_{k}(.|s,g)}\left[C_{\pi}(\hat{s_{g}}|s,g)\right]-C_{\pi}(s_{g}|s,g)$

quantifies the quality of a subgoal $s_{g}$ against the high-level policy
distribution. The Kullback-Leibler divergence term in (5) requires an esimate
of the density $p_{s}(.)$. While estimating the unknown $p_{s}(.)$ might be
challenging, we can obtain samples from this distribution, for example, by
randomly sampling states from the replay buffer $s_{g}\sim D$.

We propose instead to implicitly enforce the KL constraint (5). We first note
that the analytic solution to (5) can be obtained by enforcing the
Karush–Kuhn–Tucker conditions, for which the Lagrangian is:

$\begin{split}\mathcal{L}(\pi^{H},\lambda)&=\mathbb{E}_{s_{g}\sim\pi(.|s,g)}\left[A^{\pi^{H}_{k}}(s_{g}|s,g)\right]\\\
&+\lambda\left(\epsilon-
D_{\mathrm{KL}}\left(\pi(.|s,g)\left|\right|p_{s}(.)\right)\right).\end{split}$

The closed form solution to this problem is then given by:

$\pi_{\star}^{H}(s_{g}|s,g)=\frac{1}{Z(s,g)}p_{s}(s_{g})\exp\left(\frac{1}{\lambda}A^{\pi^{H}_{k}}(s_{g}|s,g)\right)$
(6)

with the normalizing partition function $Z(s,g)=\int
p_{s}(s_{g})\exp\left(\frac{1}{\lambda}A^{\pi^{H}_{k}}(s_{g}|s,g)\right)ds_{g}$
(Peters et al., 2010; Rawlik et al., 2012; Abdolmaleki et al., 2018b, a; Nair
et al., 2020). We project this solution into the policy space by minimizing
the forward KL divergence between our parametric high-level policy
$\pi^{H}_{\psi}$ and the optimal non parametric solution $\pi_{\star}^{H}$:

$\begin{split}\pi^{H}_{\psi_{k+1}}&=\arg\min_{\psi}\mathbb{E}_{(s,g)\sim
D}D_{\mathrm{KL}}\left(\pi_{\star}^{H}(.|s,g)\left|\right|\pi^{H}_{\psi}(.|s,g)\right)\\\
=\arg&\max_{\psi}\mathbb{E}_{(s,g)\sim D,s_{g}\sim D}\\\
&\left[\log\pi^{H}_{\psi}(s_{g}|s,g)\frac{1}{Z(s,g)}\exp\left(\frac{1}{\lambda}A^{\pi^{H}_{k}}(s_{g}|s,g)\right)\right],\end{split}$
(7)

where $\lambda$ is a hyperparameter. Conveniently, this policy improvement
step corresponds to a weighted maximum likelihood with subgoal candidates
obtained from the replay buffer randomly sampled among states visited by the
agent in previous episodes. The samples are re-weighted by their corresponding
advantage, implicitly constraining the high-level policy to stay close to the
valid distribution of states. A more detailed derivation is given in Appendix
A.

###  3.3 Policy Improvement with Imagined Subgoals

Our method builds on the following key insight. If we assume $s_{g}$ to be an
intermediate subgoal on the optimal path from $s$ to $g$, then the optimal
action for reaching $g$ from $s$ should be similar to the optimal action for
reaching $s_{g}$ from $s$ (see Figure 1). We can formalize this using a KL
constraint on the policy distribution, conditioned on goals $g$ and $s_{g}$:

$D_{\mathrm{KL}}\left(\pi(.|s,g)\left|\right|\pi(.|s,s_{g})\right)\leq\epsilon.$

We introduce the non-parametric prior policy $\pi^{prior}(.|s,g)$ as the
distribution over actions that would be chosen by the policy $\pi$ for
reaching subgoals $s_{g}\sim\pi^{H}(.|s,g)$ provided by the high-level policy.
Given a state $s$ and a goal $g$, we bootstrap the behavior of the policy at
subgoals $s_{g}\sim\pi^{H}(.|s,g)$ (see Figure 2):

$\pi^{prior}_{k}(a|s,g):=\mathbb{E}_{s_{g}\sim\pi^{H}(.|s,g)}[\pi_{\theta_{k}^{\prime}}(a|s,s_{g})].$
(8)

As we assume the subgoals to be easier to reach compared to reaching the final
goals, this prior policy provides a good initial guess to constrain the policy
search to the most promising actions.

We then propose to leverage a policy iteration scheme with additional KL
constraint to shape the policy behavior accordingly. During the policy
improvement step, in addition to maximizing the Q-function as in (2), we
encourage the policy to stay close to the prior policy through the KL-
regularization:

$\begin{split}\pi_{\theta_{k+1}}&=\arg\max_{\theta}\mathbb{E}_{(s,g)\sim
D}\mathbb{E}_{a\sim\pi_{\theta}(.|s,g)}\\\ &\left[Q^{\pi}(s,a,g)-\alpha
D_{\mathrm{KL}}\left(\pi_{\theta}(.|s,g)\left|\right|\pi^{prior}_{k}(.|s,g)\right)\right],\end{split}$
(9)

where $\alpha$ is a hyperparameter.

In practice, we found that using an exponential moving average of the online
policy parameters to construct the prior policy is necessary to ensure
convergence:

$\theta_{k+1}^{\prime}=\tau\theta_{k}+(1-\tau)\theta_{k}^{\prime},\text{
}\tau\in]0,1[.$ (10)

This ensures that the prior policy produces a more stable target for
regularizing the online policy.

###  3.4 Algorithm Summary

Algorithm 1 RL with imagined subgoals

Initialize replay buffer $D$

Initialize $Q_{\phi}$, $\pi_{\theta}$, $\pi^{H}_{\psi}$

for k = 1, 2, … do

Collect experience in $D$ using $\pi_{\theta}$ in the environment

Sample batch $(s_{t},a_{t},r_{t},s_{t+1},g)\sim D$ with HER

Sample batch of subgoal candidates $s_{g}\sim D$

Update $Q_{\phi}$ using Eq. 1 (Policy Evaluation)

Update $\pi^{H}_{\psi}$ Using Eq. 7 (High-Level Policy Improvement)

Update $\pi_{\theta}$ using Eq. 9 (Policy Improvement with Imagined Subgoals)

end for

We approximate the policy $\pi_{\theta}$, the Q-function $Q_{\phi}$ and the
high-level policy $\pi_{\psi}^{H}$ with neural networks parametrized by
$\theta$, $\phi$ and $\psi$ respectively, and train them jointly using
stochastic gradient descent. The Q-function is trained to minimize the Bellman
error (1), where we use an exponential moving average of the online Q-function
parameters to compute the target value. The high-level policy is a
probabilistic neural network whose output parametrizes a Laplace distribution
with diagonal variance
$\pi^{H}_{\psi}(.|s,g)=\text{Laplace}(\mu^{H}_{\psi}(s,g),\Sigma^{H}_{\psi}(s,g))$
trained to minimize (7). The policy is also a probabilistic network
parametrizing a squashed Gaussian distribution with diagonal variance
$\pi^{H}_{\psi}(.|s,g)=\tanh\mathcal{N}(\mu_{\theta}(s,g),\Sigma_{\theta}(s,g))$
trained to minimize (9). Finally we can approximate $\pi_{prior}$ using a
Monte-Carlo estimate of (8).

The policy $\pi$ and the high-level policy $\pi^{H}$ are trained jointly. As
the policy learns to reach more and more distant goals, its value function
becomes a better estimate for the distance between states. This allows the
high-level policy to propose more appropriate subgoals for a larger set of
goals. In turn, as the high-level policy improves, imagined subgoals offer a
more relevant supervision to shape the behavior of the policy. This virtuous
cycle progressively extends the policy horizon further and further away,
allowing more complex tasks to be solved by a single policy.

The full algorithm is summarized in Algorithm 1. We use hindsight experience
replay (Andrychowicz et al., 2017) (HER) to improve learning from sparse
rewards. Additional implementation details are given in Appendix B.

Figure 4:  (Left) Heatmap of the subgoal distribution obtained with our high-
level policy and the oracle subgoal for a given state and goal for the ant
U-maze environment. (Right) Distance between oracle subgoals and subgoals
predicted by the high-level policy for RIS and RIS without implicit
regularization. The dimensions of the space are $7.5\times 18$ units and the
ant has a radius of roughly 0.75 units.

Figure 5:  Ablation of our method on the Ant U-Maze environment: simple priors
that do not incorporate subgoals (uniform prior, moving average prior);
ignoring the effect of out-of-distribution subgoal predictions (no implicit
regularization); and using oracle subgoals (with oracle subgoals). Left:
success rate on all configurations throughout training. Right: success rate on
the test configurations.

##  4 Experiments

In this section we first introduce our experimental setup in Section 4.1.
Next, we ablate various design choices of our approach in Section 4.2. We,
then, compare RIS to prior work in goal-conditioned reinforcement learning in
Section 4.3.

###  4.1 Experimental Setup

#### Ant navigation.

We evaluate RIS on a set of ant navigation tasks of increasing difficulty,
each of which requires temporally extended reasoning. In these environments,
the agent observes the joint angles, joint velocity, and center of mass of a
quadruped ant robot navigating in a maze. We consider four different mazes: a
U-shaped maze, a S-shaped maze, a $\Pi$-shaped maze and a $\omega$-shaped maze
illustrated in Figure 6. The obstacles are unknown to the agent.

During training, initial states and goals are uniformly sampled and the agents
are trained to reach any goal in the environment. We evaluate agents in the
most extended temporal settings representing the most difficult configurations
offered by the environment (see Figure 6). We assess the success rate achieved
by these agents, where we define success as the ant being sufficiently close
to the goal position measured by x-y Euclidean distance.

#### Vision-based robotic manipulation.

We follow the experimental setup in (Nasiriany et al., 2019) and also consider
a vision-based robotic manipulation task where an agent controls a 2 DoF
robotic arm from image input and must manipulate a puck positioned on the
table (Figure 9(a)). We define success as the arm and the puck being
sufficiently close to their respective desired positions. During training, the
initial arm and puck positions and their respective desired positions are
uniformly sampled whereas, at test time, we evaluate agents on temporally
extended configurations.

These tasks are challenging because they require temporally extended reasoning
on top of complex motor control. Indeed, a greedy path towards the goal cannot
solve these tasks. We train the agents for 1 million environment steps and
average the results over 4 random seeds. Additional details about the
environments are given in Appendix C.

(a) U-shaped maze

(b) S-shaped maze

(c) $\Pi$-shaped maze

(d) $\omega$-shaped maze

Figure 6:  Comparison of RIS to several state-of-the-art methods (bottom row)
on 4 different ant navigation tasks. We evaluate the success rate of the agent
on the challenging configurations illustrates in the top row, where the ant is
located at the initial state and the desired goal location is represented by a
cyan sphere.

#### Alternative methods.

We compare our approach, RIS, to off-policy reinforcement learning methods for
goal-reaching tasks. We consider Soft Actor-Critic (SAC) (Haarnoja et al.,
2018a) with HER, which trains a policy from scratch by maximizing the entropy
regularized objective using the same sparse reward as ours. We also compare to
Temporal Difference Models (TDM) (Pong et al., 2018) which trains horizon-
aware policies operating under dense rewards in the form of distance to the
goal. We chose to evaluate TDMs with a long policy horizon of 600 steps due to
the complexity of considered tasks. Furthermore, we compare to Latent
Embeddings for Abstracted Planning (LEAP) (Nasiriany et al., 2019), a
competitive approach for these environments, which uses a sequence of subgoals
that a TDM policy must reach one after the other during inference.

We re-implemented SAC, TDM and LEAP and validated our implementations on the
U-shaped ant maze and vision-based robotic manipulation environments.

On the U-shaped Ant maze environment, we additionally report results of HIRO
(Nachum et al., 2018), a hierarchical reinforcement learning method with off-
policy correction, after 1 million environment steps. The results were copied
from Figure 12 in (Nasiriany et al., 2019).

We provide additional details about the hyperparameters used in our
experiments in Appendix D.

###  4.2 Ablative Analysis

We use the Ant U-maze navigation task to conduct ablation studies to validate
the role of our design choices.

(a) S-shaped Maze

(b) $\omega$-shaped Maze

Figure 8: Comparison of different methods for increasingly more difficult
tasks on the S-shaped and $\omega$-shaped ant maze. The goals with increasing
complexity are numbered starting with 1, where 1 is the closest goal.

Evaluating imagined subgoals. To evaluate the quality of the subgoals
predicted by our high-level policy, we introduce an oracle subgoal sampling
procedure. We plan an oracle trajectory, which corresponds to a point-mass
agent navigating in this maze and does not necessarily correspond to the
optimal trajectory of an ant. Oracle subgoals correspond to the midpoint of
this trajectory between state and goal in x-y location (Figure 4 left).

Figure 4 (left) shows the subgoal distribution predicted by a fully-trained
high-level policy for a state and goal pair located at opposite sides of the
U-shaped maze. We can observe that its probability mass is close to the oracle
subgoal.

To quantitatively evaluate the quality of imagined subgoals, we measure the
x-y Euclidean distance between oracle subgoals and subgoals sampled from the
high-level policy throughout training for a set of fixed state and goal tuples
randomly sampled in the environment. Figure 4 (right) shows that RIS
successfully learns to find subgoals that are coherent with the oracle
trajectory during training, despite not having prior knowledge about the
environment.

We also assess the importance of using the implicit regularization scheme
presented in Section 3.2 which discourages high-level policy predictions to
lie outside of the distribution of valid states. We compare our approach
against naively optimizing subgoals without regularization (i.e. directly
optimize (4)). Figure 4 (right) shows that without implicit regularization,
the predicted subgoals significantly diverge from oracle subgoals in x-y
location during training. As result, imagined subgoals are not able to
properly guide policy learning to solve the task, see Figure 5.

Prior policy with imagined subgoals. We evaluate the importance of
incorporating imagined subgoals into the prior policy by comparing to a number
of variants. To disentangle the effects of imagined subgoals from the actor-
critic architecture used by RIS, we first replace our prior policy with
simpler choices of prior distributions that do not incorporate any subgoals:
(i) a uniform prior over the actions
$\pi^{prior}_{k}=\mathcal{U}(\mathcal{A})$, which is equivalent to SAC without
entropy regularization during policy evaluation, and (ii) a parametric prior
policy that uses an exponential moving average of the online policy weights
$\pi^{prior}_{k}=\pi_{\theta^{\prime}_{k}}$. Figure 5 shows that, while they
can learn goal-reaching behaviors on many configurations encountered during
training (left), neither of these variants are able to solve the Ant U-maze
environment in its most difficult setting (right). We also observe that the
agent with a moving average action prior fails to learn (left). This
highlights the benefits of incorporating subgoals into policy learning.

Finally, we propose to incorporate the oracle subgoals into our prior policy.
We replace the subgoal distributions predicted by our high-level policy with
Laplace distributions centered around oracle subgoals. Results in Figure 5
show that RIS with oracle subgoals learns to solve the U-shaped ant maze
environment faster than using a high-level policy simultaneously trained by
the agent. This experiment highlights the efficiency of our approach to guide
policy learning with appropriate subgoals: if we have access to proper
subgoals right from the beginning of the training, our approach could leverage
them to learn even faster. However, such subgoals are generally not readily
available without prior knowledge about the environment. Thus, we introduce a
high-level policy training procedure which determines appropriate subgoals
without any supervision.

###  4.3 Comparison to the State of the Art

Ant navigation. Figure 6 compares RIS to the alternative methods introduced in
section 4.1 for the four ant navigation environments. For all considered mazes
RIS significantly outperforms prior methods in terms of sample efficiency,
often requiring less than 500 thousand environment interactions to solve the
mazes in their most challenging initial state and goal configurations.

LEAP makes progress on these navigation tasks, but requires significantly more
environment interactions. This comparison shows the effectiveness of our
approach, which uses subgoals to guide the policy rather than reaching them
sequentially as done by LEAP. While SAC manages to learn goal-reaching
behaviors, as we will see later in Figure 8, it fails to solve the
environments in their most challenging configurations. The comparison to SAC
highlights the benefits of using our informed prior policy compared to methods
assuming a uniform action prior. On the U-shaped maze environment, HIRO
similarly fails to solve the task within one million environment interactions.
Furthermore, we observe that TDMs fails to learn due to the sparsity of the
reward.

Figure 8 evaluates how SAC, LEAP and RIS perform for varying task horizons in
the S-shaped and $\omega$-shaped mazes. Starting from an initial state located
at the edge of the mazes, we sample goals at locations which require an
increasing number of environment steps to be reached. Figure 8 reports results
for RIS, LEAP and SAC after having been trained for 1 million steps. While the
performances of LEAP and SAC degrades as the planning horizon increases, RIS
consistently solves configurations of increasing complexity.

These results demonstrate that RIS manages to solve complex navigation tasks
despite relying on a flat policy at inference. In contrast, LEAP performs less
well, despite a significantly more expensive planning of subgoals during
inference.

(a) Illustration of the robotic manipulation task

(b) Comparison to prior works

Figure 9:  Robotic manipulation environment: (a) illustration of the task; (b)
results of our method compared to LEAP and SAC.

Vision-based robotic manipulation. While the ant navigation experiments
demonstrate the performances of RIS on environments with low-dimension state
spaces, we also show how our method can be applied to vision-based robotic
manipulation tasks. Our approach takes images of the current and desired
configurations as input. Input images are passed through an image encoder, a
convolutional neural network shared between the policy, high-level policy and
Q-function. The encoder is only updated when training the Q-function during
policy evaluation and is fixed otherwise during policy improvement and high-
level policy improvement. Instead of generating subgoals in the high-
dimensional image space, the high-level policy therefore operates in the
learned compact image representation of the encoder. Following recent works on
reinforcement learning from images, we augment image observations with random
translations (Kostrikov et al., 2020; Laskin et al., 2020). We found that
using such data augmentation was important for training image-based RIS and
SAC policies. Moreover, we found that using a lower learning rate for the
policy was necessary to stabilize training. Additional implementation details
on the image encoding are given in Appendix B.

We compare our approach against LEAP and SAC in Figure 9(b). RIS achieves a
higher success rate than LEAP whereas SAC fails most of the time to solve the
manipulation task consistently enough in the temporally extended configuration
used for evaluation on the vision-based robotic manipulation task. Moreover,
RIS and SAC only requires a single forward pass through their image encoder
and actor network at each time step when interacting with the environment,
whereas LEAP depends in addition upon an expensive planning of image subgoals.

Figure 10 visualizes the imagined subgoals of the high-level policy. Once the
RIS agent is fully trained, we separately train a decoder to reconstruct image
observations from their learned representations. Given observations of the
current state and the desired goal, we then predict the representation of an
imagined subgoal with the high-level policy and generate the corresponding
image using the decoder. Figure 10 shows that subgoals predicted by the high-
level policy are natural intermediate states halfway to the desired goal on
this manipulation task. For example, for the test configuration (Figure 10
top), the high-level policy prediction corresponds to a configuration where
the arm has reached the right side of the puck and is pushing it towards its
desired position.

Additional reconstructions of imagined subgoals for different initial state
and goal configurations in the robotic manipulation environment are given in
Appendix E.

Figure 10: Image reconstruction of an imagined subgoal (middle) given the
current state (left) and the desired goal (right) on a temporally extended
configuration used for evaluation (top) and a random configuration (bottom).

##  5 Conclusion

We introduced RIS, a goal-conditioned reinforcement learning method that
imagines possible subgoals in a self-supervised fashion and uses them to
facilitate training. We propose to use the value function of the goal-reaching
policy to train a high-level policy operating in the state space. We then use
imagined subgoals to define a prior policy and incorporate this prior into
policy learning. Experimental results on challenging simulated navigation and
vision-based manipulation environments show that our proposed method greatly
accelerates learning of temporally extended tasks and outperforms competing
approaches.

While our approach makes use of subgoals to facilitate policy search, future
work could explore how to use them to obtain better Q-value estimates. Future
work could also improve exploration by using imagined subgoals to encourage
the policy to visit all potential states.

##  6 Acknowledgements

This work was funded in part by the French government under management of
Agence Nationale de la Recherche as part of the ”Investissements d’avenir”
program, reference ANR19-P3IA-0001 (PRAIRIE 3IA Institute) and by Louis
Vuitton ENS Chair on Artificial Intelligence.

## References

  * Abdolmaleki et al. (2018a) Abdolmaleki, A., Springenberg, J. T., Degrave, J., Bohez, S., Tassa, Y., Belov, D., Heess, N., and Riedmiller, M.  Relative entropy regularized policy iteration.  _arXiv preprint arXiv:1812.02256_ , 2018a.
  * Abdolmaleki et al. (2018b) Abdolmaleki, A., Springenberg, J. T., Tassa, Y., Munos, R., Heess, N., and Riedmiller, M. A.  Maximum a posteriori policy optimisation.  In _International Conference on Learning Representations_ , 2018b.
  * Andrychowicz et al. (2017) Andrychowicz, M., Crow, D., Ray, A., Schneider, J., Fong, R., Welinder, P., McGrew, B., Tobin, J., Abbeel, P., and Zaremba, W.  Hindsight experience replay.  In _Advances in Neural Information Processing Systems_ , 2017.
  * Dayan & Hinton (1993) Dayan, P. and Hinton, G. E.  Feudal reinforcement learning.  In _Advances in Neural Information Processing Systems_ , 1993.
  * Dietterich (2000) Dietterich, T. G.  Hierarchical reinforcement learning with the maxq value function decomposition.  _Journal of artificial intelligence research_ , 13, 2000.
  * Eysenbach et al. (2019) Eysenbach, B., Salakhutdinov, R., and Levine, S.  Search on the replay buffer: Bridging planning and reinforcement learning.  In _Advances in Neural Information Processing Systems_ , 2019.
  * Eysenbach et al. (2020) Eysenbach, B., Salakhutdinov, R., and Levine, S.  C-learning: Learning to achieve goals via recursive classification.  _arXiv preprint arXiv:2011.08909_ , 2020.
  * Fujimoto et al. (2018) Fujimoto, S., van Hoof, H., and Meger, D.  Addressing function approximation error in actor-critic methods.  In _International Conference on Machine Learning_ , 2018.
  * Galashov et al. (2019) Galashov, A., Jayakumar, S. M., Hasenclever, L., Tirumala, D., Schwarz, J., Desjardins, G., Czarnecki, W. M., Teh, Y. W., Pascanu, R., and Heess, N.  Information asymmetry in kl-regularized RL.  In _International Conference on Learning Representations_ , 2019.
  * Gupta et al. (2019) Gupta, A., Kumar, V., Lynch, C., Levine, S., and Hausman, K.  Relay policy learning: Solving long-horizon tasks via imitation and reinforcement learning.  _arXiv preprint arXiv:1910.11956_ , 2019.
  * Haarnoja et al. (2017) Haarnoja, T., Tang, H., Abbeel, P., and Levine, S.  Reinforcement learning with deep energy-based policies.  In _International Conference on Machine Learning_ , 2017.
  * Haarnoja et al. (2018a) Haarnoja, T., Zhou, A., Abbeel, P., and Levine, S.  Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor.  In _International Conference on Machine Learning_ , 2018a.
  * Haarnoja et al. (2018b) Haarnoja, T., Zhou, A., Hartikainen, K., Tucker, G., Ha, S., Tan, J., Kumar, V., Zhu, H., Gupta, A., Abbeel, P., et al.  Soft actor-critic algorithms and applications.  Technical report, 2018b.
  * Heess et al. (2015) Heess, N., Wayne, G., Silver, D., Lillicrap, T. P., Erez, T., and Tassa, Y.  Learning continuous control policies by stochastic value gradients.  In _Advances in Neural Information Processing Systems_ , 2015.
  * Jurgenson et al. (2020) Jurgenson, T., Avner, O., Groshev, E., and Tamar, A.  Sub-goal trees a framework for goal-based reinforcement learning.  In _International Conference on Machine Learning_ , 2020.
  * Kaelbling (1993) Kaelbling, L. P.  Learning to achieve goals.  In _IJCAI_ , 1993.
  * Kappen et al. (2012) Kappen, H. J., Gómez, V., and Opper, M.  Optimal control as a graphical model inference problem.  _Machine learning_ , 87(2):159–182, 2012.
  * Kostrikov et al. (2020) Kostrikov, I., Yarats, D., and Fergus, R.  Image augmentation is all you need: Regularizing deep reinforcement learning from pixels.  _arXiv preprint arXiv:2004.13649_ , 2020.
  * Kumar et al. (2019) Kumar, A., Fu, J., Soh, M., Tucker, G., and Levine, S.  Stabilizing off-policy q-learning via bootstrapping error reduction.  In _Advances in Neural Information Processing Systems_ , 2019.
  * Laskin et al. (2020) Laskin, M., Lee, K., Stooke, A., Pinto, L., Abbeel, P., and Srinivas, A.  Reinforcement learning with augmented data.  In _Advances in Neural Information Processing Systems_ , 2020.
  * Levine (2018) Levine, S.  Reinforcement learning and control as probabilistic inference: Tutorial and review.  _arXiv preprint arXiv:1805.00909_ , 2018.
  * Levine & Abbeel (2014) Levine, S. and Abbeel, P.  Learning neural network policies with guided policy search under unknown dynamics.  In _Advances in Neural Information Processing Systems_ , 2014.
  * Levine & Koltun (2013) Levine, S. and Koltun, V.  Guided policy search.  In _International Conference on Machine Learning_ , 2013.
  * Levine et al. (2016) Levine, S., Finn, C., Darrell, T., and Abbeel, P.  End-to-end training of deep visuomotor policies.  _The Journal of Machine Learning Research_ , 17(1):1334–1373, 2016.
  * Levy et al. (2019) Levy, A., Konidaris, G. D., Jr., R. P., and Saenko, K.  Learning multi-level hierarchies with hindsight.  In _International Conference on Learning Representations_ , 2019.
  * Mnih et al. (2016) Mnih, V., Badia, A. P., Mirza, M., Graves, A., Lillicrap, T. P., Harley, T., Silver, D., and Kavukcuoglu, K.  Asynchronous methods for deep reinforcement learning.  In _International Conference on Machine Learning_ , 2016.
  * Nachum et al. (2018) Nachum, O., Gu, S., Lee, H., and Levine, S.  Data-efficient hierarchical reinforcement learning.  In _Advances in Neural Information Processing Systems_ , 2018.
  * Nair et al. (2018) Nair, A., Pong, V., Dalal, M., Bahl, S., Lin, S., and Levine, S.  Visual reinforcement learning with imagined goals.  In _Advances in Neural Information Processing Systems_ , 2018.
  * Nair et al. (2020) Nair, A., Dalal, M., Gupta, A., and Levine, S.  Accelerating online reinforcement learning with offline datasets.  _arXiv preprint arXiv:2006.09359_ , 2020.
  * Nair & Finn (2020) Nair, S. and Finn, C.  Hierarchical foresight: Self-supervised learning of long-horizon tasks via visual subgoal generation.  In _International Conference on Learning Representations_ , 2020.
  * Nasiriany et al. (2019) Nasiriany, S., Pong, V., Lin, S., and Levine, S.  Planning with goal-conditioned policies.  In _Advances in Neural Information Processing Systems_ , 2019.
  * Parascandolo et al. (2020) Parascandolo, G., Buesing, L., Merel, J., Hasenclever, L., Aslanides, J., Hamrick, J. B., Heess, N., Neitz, A., and Weber, T.  Divide-and-conquer monte carlo tree search for goal-directed planning.  _arXiv preprint arXiv:2004.11410_ , 2020.
  * Pertsch et al. (2020a) Pertsch, K., Lee, Y., and Lim, J. J.  Accelerating reinforcement learning with learned skill priors.  _arXiv preprint arXiv:2010.11944_ , 2020a.
  * Pertsch et al. (2020b) Pertsch, K., Rybkin, O., Ebert, F., Zhou, S., Jayaraman, D., Finn, C., and Levine, S.  Long-horizon visual planning with goal-conditioned hierarchical predictors.  _Advances in Neural Information Processing Systems_ , 2020b.
  * Peters et al. (2010) Peters, J., Mülling, K., and Altun, Y.  Relative entropy policy search.  In _AAAI Conference on Artificial Intelligence_ , 2010.
  * Pitis et al. (2020) Pitis, S., Chan, H., Zhao, S., Stadie, B. C., and Ba, J.  Maximum entropy gain exploration for long horizon multi-goal reinforcement learning.  In _International Conference on Machine Learning_ , 2020.
  * Pong et al. (2018) Pong, V., Gu, S., Dalal, M., and Levine, S.  Temporal difference models: Model-free deep RL for model-based control.  In _International Conference on Learning Representations_ , 2018.
  * Rawlik et al. (2012) Rawlik, K., Toussaint, M., and Vijayakumar, S.  On stochastic optimal control and reinforcement learning by approximate inference.  _Proceedings of Robotics: Science and Systems VIII_ , 2012.
  * Schaul et al. (2015) Schaul, T., Horgan, D., Gregor, K., and Silver, D.  Universal value function approximators.  In _International Conference on Machine Learning_ , 2015.
  * Siegel et al. (2020) Siegel, N. Y., Springenberg, J. T., Berkenkamp, F., Abdolmaleki, A., Neunert, M., Lampe, T., Hafner, R., and Riedmiller, M.  Keep doing what worked: Behavioral modelling priors for offline reinforcement learning.  _arXiv preprint arXiv:2002.08396_ , 2020.
  * Silver et al. (2014) Silver, D., Lever, G., Heess, N., Degris, T., Wierstra, D., and Riedmiller, M. A.  Deterministic policy gradient algorithms.  In _International Conference on Machine Learning_ , 2014.
  * Sutton et al. (1999) Sutton, R. S., Precup, D., and Singh, S.  Between mdps and semi-mdps: A framework for temporal abstraction in reinforcement learning.  _Artificial intelligence_ , 112(1-2):181–211, 1999.
  * Teh et al. (2017) Teh, Y. W., Bapst, V., Czarnecki, W. M., Quan, J., Kirkpatrick, J., Hadsell, R., Heess, N., and Pascanu, R.  Distral: Robust multitask reinforcement learning.  In _Advances in Neural Information Processing Systems_ , 2017.
  * Tirumala et al. (2019) Tirumala, D., Noh, H., Galashov, A., Hasenclever, L., Ahuja, A., Wayne, G., Pascanu, R., Teh, Y. W., and Heess, N.  Exploiting hierarchy for learning and transfer in kl-regularized rl.  _arXiv preprint arXiv:1903.07438_ , 2019.
  * Toussaint (2009) Toussaint, M.  Robot trajectory optimization using approximate inference.  In _International Conference on Machine Learning_ , 2009.
  * Veeriah et al. (2018) Veeriah, V., Oh, J., and Singh, S.  Many-goals reinforcement learning.  _arXiv preprint arXiv:1806.09605_ , 2018.
  * Vezhnevets et al. (2017) Vezhnevets, A. S., Osindero, S., Schaul, T., Heess, N., Jaderberg, M., Silver, D., and Kavukcuoglu, K.  Feudal networks for hierarchical reinforcement learning.  In _International Conference on Machine Learning_ , 2017.
  * Wang et al. (2020) Wang, Z., Novikov, A., Zolna, K., Merel, J., Springenberg, J. T., Reed, S. E., Shahriari, B., Siegel, N. Y., Gülçehre, Ç., Heess, N., and de Freitas, N.  Critic regularized regression.  In _Advances in Neural Information Processing Systems_ , 2020.
  * Wiering & Schmidhuber (1997) Wiering, M. and Schmidhuber, J.  Hq-learning.  _Adaptive Behavior_ , 6(2):219–246, 1997.
  * Wu et al. (2019) Wu, Y., Tucker, G., and Nachum, O.  Behavior regularized offline reinforcement learning.  _arXiv preprint arXiv:1911.11361_ , 2019.
  * Zhang et al. (2020) Zhang, L., Yang, G., and Stadie, B. C.  World model as a graph: Learning latent landmarks for planning.  _arXiv preprint arXiv:2011.12491_ , 2020.
  * Zhao et al. (2019) Zhao, R., Sun, X., and Tresp, V.  Maximum entropy-regularized multi-goal reinforcement learning.  In _International Conference on Machine Learning_ , 2019.

##  Appendix A High-level Policy Derivation

The following derivation is similar to (Abdolmaleki et al., 2018b) and (Nair
et al., 2020). We consider the subgoal advantage function $A^{\pi^{H}}$, some
distribution over valid states $p_{s}(s)$ and the distribution over states and
goals $\mu(s,g)$ given by previously collected experience and experience
relabelling. During high-level policy improvement, we maximize the following
objective function, where expectations have been replaced by integrals:

$\begin{split}\pi^{H}_{k+1}&=\arg\max_{\pi}\iint\mu(s,g)\int\pi^{H}(s_{g}|s,g)A^{\pi^{H}}(s_{g}|s,g)ds_{g}dsdg\\\
\text{s.t.
}\iint\mu(s,g)&D_{\mathrm{KL}}\left(\pi^{H}(.|s,g)\left|\right|p_{s}(.)\right)dsdg\leq\epsilon,\iint\mu(s,g)\int\pi^{H}(s_{g}|s,g)ds_{g}dsdg=1\end{split}$
(11)

The Lagrangian is:

$\begin{split}\mathcal{L}(\pi^{H},\lambda,\eta)&=\iint\mu(s,g)\int\pi^{H}(s_{g}|s,g)A^{\pi^{H}}(s_{g}|s,g)ds_{g}dsdg\\\
&+\lambda\left(\epsilon-\iint\mu(s,g)\int\pi^{H}(s_{g}|s,g)\log\frac{\pi^{H}(s_{g}|s,g)}{p_{s}(s_{g})}ds_{g}dsdg\right)\\\
&+\eta\left(1-\iint\mu(s,g)\int\pi^{H}(s_{g}|s,g)ds_{g}dsdg\right).\end{split}$
(12)

Differentiating with respect to $\pi^{H}$ yields

$\frac{\partial\mathcal{L}}{\partial\pi^{H}}=A^{\pi^{H}}(s_{g}|s,g)-\lambda\log\pi^{H}(s_{g}|s,g)+\lambda\log
p_{s}(s_{g})+\eta-\lambda.$ (13)

Setting this expression to zero, we get

$\pi^{H}(s_{g}|s,g)=p_{s}(s_{g})\exp\left(\frac{A^{\pi^{H}}(s_{g}|s,g)}{\lambda}\right)\exp\left(-\frac{\lambda-\eta}{\lambda}\right).$
(14)

As the second constraint in (11) must sum to 1, the last exponential is a
normalizing factor. This gives the closed form solution:

$\pi_{\star}^{H}(s_{g}|s,g)=\frac{1}{Z(s,g)}p_{s}(s_{g})\exp\left(\frac{A^{\pi^{H}}(s_{g}|s,g)}{\lambda}\right)$
(15)

with the normalizing partition function

$Z(s,g)=\int
p_{s}(s_{g})\exp\left(\frac{1}{\lambda}A^{\pi^{H}}(s_{g}|s,g)\right)ds_{g}.$
(16)

We next project the closed form solution into the space of parametric policies
by minimizing the reverse KL divergence between our parametric high-level
policy $\pi^{H}_{\psi}$ and the optimal non-parametric solution
$\pi^{H}_{\star}$:

$\begin{split}\pi^{H}_{\psi_{k+1}}&=\arg\min_{\psi}\mathbb{E}_{(s,g)\sim\mu(.)}\left[D_{\mathrm{KL}}\left(\pi_{\star}^{H}(.|s,g)\left|\right|\pi_{\psi}^{H}(.|s,g)\right)\right]\\\
&=\arg\max_{\psi}\mathbb{E}_{(s,g)\sim\mu(.)}\mathbb{E}_{s_{g}\sim\pi_{\star}^{H}(.|s,g)}\left[\log\pi^{H}_{\psi}(s_{g}|s,g)\right]\\\
&=\arg\max_{\psi}\mathbb{E}_{(s,g)\sim\mu(.),s_{g}\sim
p_{s}(.)}\left[\log\pi_{\psi}^{H}(s_{g}|s,g)\frac{1}{Z(s,g)}\exp\left(\frac{1}{\lambda}A^{\pi^{H}}(s_{g}|s,g)\right)\right]\end{split}$
(17)

We choose $p_{s}$ to be the distribution of states given by experience
collected in the replay buffer $D$, such that high-level policy improvement
corresponds to a weighted maximum-likelihood where subgoal candidates $s_{g}$
are randomly sampled from $D$ among states visited by the agent in previous
episodes.

##  Appendix B Implementation Details

#### Actor-Critic.

Our implementation of the actor-critic algorithm is based on Soft Actor-Critic
(Haarnoja et al., 2018b), where we remove the entropy term during policy
evaluation and replace the entropy term by the KL divergence between policy
and our prior policy during policy improvement. The policy is a neural network
that parametrizes the mean and diagonal covariance matrix of a squashed
Gaussian distribution
$\pi_{\theta}(.|s,g)=\tanh\mathcal{N}(\mu_{\theta}(s,g),\Sigma_{\theta}(s,g))$.
We train two seperate Q-networks with target networks and take the minimum
over the two target values to compute the bootstrap value. The target networks
are updated using an exponential moving average of the online Q parameters:
$\phi_{k+1}^{\prime}=\tau\phi_{k}+(1-\tau)\phi_{k}^{\prime}$.

#### High-level policy training.

The high-level policy is a neural network that outputs the mean and diagonal
covariance matrix of a Laplace distribution
$\pi^{H}_{\psi}(.|s,g)=\text{Laplace}(\mu_{\psi}(s,g),\Sigma_{\psi}(s,g))$.
Following (Nair et al., 2020), instead of estimating the normalizing factor
$Z(s,g)$ in (17), we found that computing the weights as softmax of the
advantages over the batch leads to good results in practice. During the high-
level policy improvement, we found that clipping the value function between
$-100$ and $0$, which corresponds to the expected bounds given our choice of
reward function and discount factor, stabilizes the training slightly.

#### KL divergence estimation.

We use an exponentially moving average of the policy weights instead of the
weights of the current policy to construct the prior policy $\pi_{prior}$:
$\theta_{k+1}^{\prime}=\tau\theta_{k}+(1-\tau)\theta_{k}^{\prime}$ with the
same smoothing coefficient $\tau$ as the one used for the Q function. We
estimate the prior density using the following Monte-Carlo estimate:

$\log\pi^{prior}(a|s,g)\approx\log\left[\frac{1}{I}\sum_{i}\pi_{\theta^{\prime}}(a|s,s_{g}^{i})+\epsilon\right],(s_{g}^{i})\sim\pi^{H}_{\psi}(.|s,g),$
(18)

where $\epsilon>0$ is a small constant to avoid large negative values of the
prior log-density. We use $I=10$ samples to estimate $\pi^{prior}(a|s,g)$. We
also use a Monte-Carlo approximation to estimate the KL-divergence term in
Equation (9) of the submission:

$\begin{split}D_{\mathrm{KL}}\left(\pi_{\theta}(.|s,g)\left|\right|\pi^{prior}(.|s,g)\right)&=\mathbb{E}_{a\sim\pi(.|s,g)}[\log\pi_{\theta}(a|s,g)-\log\pi^{prior}(a|s,g)]\\\
&\approx\frac{1}{N}\sum_{n}\left[\log\pi_{\theta}(a_{n}|s,g)-\log\pi^{prior}(a_{n}|s,g)\right]\\\
\text{with }(a_{n})_{n=1,..,N}&\sim\pi_{\theta}(.|s,g).\end{split}$ (19)

Following SAC (Haarnoja et al., 2018a), we use $N=1$, plug the estimate (18)
and use the reparametrization trick to backpropagate the KL divergence term to
the policy weights (Haarnoja et al., 2018a).

#### Experience relabelling.

In all of our experiments we use Hindsight Experience Replay (Andrychowicz et
al., 2017). We use the same relabelling strategy as (Nair et al., 2018) and
(Nasiriany et al., 2019) and relabel the goals in our minibatches as follows:

  * •

$20\%$: original goals from collected trajectories,

  * •

$40\%$: randomly sampled states from the replay buffer, trajectories,

  * •

$40\%$: future states along the same collected trajectory.

#### Vision based environments

On the vision-based robotic manipulation tasks, input images are passed
through an image encoder shared between the policy, high-level policy and
Q-function. Both states, goals and subgoals are encoded using the same encoder
network. The encoder is updated during policy evaluation, where we only update
the representation of the current state images whereas the representations of
desired goal images, next state images and subgoal image candidates are kept
fixed. We augment the observations with random translations by translating the
$84\times 84$ image within a $100\times 100$ empty frame (Laskin et al., 2020;
Kostrikov et al., 2020).

##  Appendix C Environments

We adopt Ant Navigation environments provided by the code repository of
(Nasiriany et al., 2019). In these environments, a 4-legged ant robot must
learn to navigate in various mazes. The state includes the position,
orientation, the joint angles and the velocities of these components. The ant
has a radius of roughly 0.75 units. We consider that the goal is reached if
the x-y position of the ant is within 0.5 units of its desired location in
Euclidean distance. The agent receives a reward $r(s,a,g)=-1$ for all actions
until the goal is reached. The dimensions of the space for the U-shaped maze
are $7.5\times 18$ units. For the S-shaped maze, the dimensions are $12\times
12$. For the $\Pi$-shaped and $\omega$-shaped mazes, the dimensions are
$16\times 16$ units. The walls are $1.5$ units thick. During training, initial
states and goals are uniformly sampled anywhere in the empty space of the
environment. At test time, we evaluate the agent on challenging configurations
that require temporally extended reasoning, as illustrated in Figure 6 of the
submission.

For the robotic manipulation task, we use the same image-based environment as
in (Nasiriany et al., 2019). In this task, the agent operates an arm robot via
2D position control and must manipulate a puck. The agent observes a $84\times
84$ RGB image showing a top-down view of the scene. The dimension of the
workspace are $40\text{cm}\times 20\text{cm}$ and the puck has a radius of
$4\text{cm}$. We consider that the goal is achieved if both the arm and the
puck are within $5\text{cm}$ of their respective target positions. During
training, the initial arm and puck positions and their respective desired
positions are uniformly sampled in the workspace. At test time, we evaluate
the policy on a hard configuration which requires temporally extended
reasoning: the robot must reach across the table to a corner where the puck is
located, move its arm around the puck, pull the puck to a different corner of
the table, and reach again the opposite corner.

##  Appendix D Hyperparameters

Table 1 lists the hyperparameters used for the RIS. We use Adam optimizer and
report results after one million interactions with the environment. For SAC,
following (Haarnoja et al., 2018b), we automatically tune the entropy of the
policy to match the target entropy of $-\text{ dim}(\mathcal{A})$.

Table 1: Hyper-parameters for RIS and SAC.

Hyper-parameter | Ant Navigation | Robotic Manipulation
---|---|---
Q hidden sizes | $[256,256]$ | $[256,256]$
Policy hidden sizes | $[256,256]$ | $[256,256]$
High-level policy hidden sizes | $[256,256]$ | $[256,256]$
Hidden activation functions | ReLU | ReLU
Batch size | $2048$ | $256$
Training batches per environment step | $1$ | $1$
Replay buffer size | $1\times 10^{6}$ | $1\times 10^{5}$
Discount factor $\gamma$ | $0.99$ | $0.99$
polyak for target networks $\tau$ | $5\times 10^{-3}$ | $5\times 10^{-3}$
$\epsilon$ | $1\times 10^{-16}$ | $1\times 10^{-4}$
Critic learning rate | $1\times 10^{-3}$ | $1\times 10^{-3}$
Policy learning rates | $1\times 10^{-3}$ | $1\times 10^{-4}$
High-level policy learning rate | $1\times 10^{-4}$ | $1\times 10^{-4}$
$\alpha$ | $0.1$ | $0.1$
$\lambda$ | $0.1$ | $0.1$

In the vision based environment, our image encoder is a serie of convolutional
layers with kernel sizes $[3,3,3,3]$, strides $[2,2,2,1]$, channel sizes
$[32,32,32,32]$ and $ReLU$ activation functions followed by a fully-connected
layer with output dimension $16$.

Table 2: Environment specific hyper-parameters for LEAP

Hyperparameter | U-shaped maze | S-shaped maze |  $\Pi$-shaped maze |  $\omega$-shaped Maze | Robotic Manipulation
---|---|---|---|---|---
TDM policy horizon | 50 | 50 | 75 | 100 | 25
Number of subgoals | 11 | 11 | 11 | 11 | 3

For LEAP, we re-implemented (Nasiriany et al., 2019) and train TDM (Pong et
al., 2018) policies and Q networks with hidden layers of size $[400,300]$ and
$ReLU$ activation functions. In the ant navigation environments, we pretrain
VAEs with mean squared reconstruction error loss and hidden layers of size
$[64,128,64]$, $ReLU$ activation functions and representation size of $8$ for
the encoders and the decoders. In the vision based robotic manipulation
environment, we pretrain VAEs with mean squared error reconstruction loss and
convolutional layers with encoder kernel of sizes $[5,5,5]$, encoder strides
of sizes $[3,3,3]$, encoder channels of sizes $[16,16,32]$, decoder kernel
sizes of sizes $[5,6,6]$, decoder strides of sizes $[3,3,3]$, and decoder
channels of sizes $[32,32,16]$, representation size of $16$ and $ReLU$
activation functions. Table 2 reports the policy horizon used for each
environment as well as the number of subgoals in the test configuration for
the results in Figure 6 of the submission. For the results presented in Figure
7 of the submission, we adapted the number of subgoals according to the
difficulty of each configuration.

##  Appendix E Additional Results

In Figure 11, we provide image reconstructions of imagined subgoals on
additional configurations of the vision-based robotic manipulation
environment.

Figure 11: Image reconstruction of imagined subgoals (middle column) given
current states (left column) and desired goals (right column) for different
random configurations in the vision-based robotic manipulation environment.