# Flexible Option Learning

Martin Klissarov
Mila, McGill University
<EMAIL_ADDRESS>
&Doina Precup
Mila, McGill University and DeepMind
<EMAIL_ADDRESS>

###### Abstract

Temporal abstraction in reinforcement learning (RL), offers the promise of
improving generalization and knowledge transfer in complex environments, by
propagating information more efficiently over time. Although option learning
was initially formulated in a way that allows updating many options
simultaneously, using off-policy, intra-option learning (Sutton, Precup &
Singh, 1999), many of the recent hierarchical reinforcement learning
approaches only update a single option at a time: the option currently
executing. We revisit and extend intra-option learning in the context of deep
reinforcement learning, in order to enable updating all options consistent
with current primitive action choices, without introducing any additional
estimates. Our method can therefore be naturally adopted in most hierarchical
RL frameworks. When we combine our approach with the option-critic algorithm
for option discovery, we obtain significant improvements in performance and
data-efficiency across a wide variety of domains.

##  1 Introduction

Temporal abstraction is a fundamental component of intelligent agents as it
allows for explicit reasoning at different timescales. The options framework
(Sutton et al., 1999b) provides a clear formalism for such abstractions and
proposes efficient ways to learn directly from the environment’s reward
signal. As the agent needs to learn about the value of a possibly large number
of options, it is crucial to maximize each interaction to ensure sample
efficiency. Sutton et al. (1999b) therefore propose the intra-option value
learning algorithm, which updates all consistent options simultaneously from a
single transition.

Recently there have been important developments on how to learn or discover
options from scratch when using function approximation (Bacon et al., 2016;
Smith et al., 2018; Riemer et al., 2018; Bagaria and Konidaris, 2019; Zhang
and Whiteson, 2019; Harutyunyan et al., 2017). However, most of recent
research only updates a single option at a time, that is, the option selected
by the agent in the sampled state. This is perhaps due to the fact that
consistency between options is less likely to arise naturally when using
function approximation.

We propose a scalable method to better exploit the agent’s experience by
updating all relevant options, where relevance is defined as the likelihood of
the option being selected. We present a decomposition of the state-option
distribution which we leverage to remain consistent in the function
approximation setting while providing performance improvements similarly to
the ones shown in the tabular setting (Sutton et al., 1999b).

When an option set is given to the agent and only the value of each option
remains to be learned, updating all relevant options lets the agent determine
faster when to apply each of them. In the case where all options components
are learned from scratch, our approach can additionally help mitigate the
issue of degenerate solutions (Harb et al., 2017). Indeed, such solutions are
likely due to the "rich-get-richer" phenomenon where an updated option has
more chance of being picked again when compared to a randomly initialized
option. By avoiding degenerate solutions, one can obtain temporally extended
options which leads to meaningful and interpretable behavior.

Unlike recent approaches that leverage inference to update all options (Daniel
et al., 2016; Smith et al., 2018; Wulfmeier et al., 2020), our method is
naturally compatible with many state of the art policy optimization frameworks
(Mnih et al., 2016; Schulman et al., 2017; Fujimoto et al., 2018) and requires
no additional estimators. We empirically verify the merits of our approach on
a wide variety of domains, ranging from gridworlds using tabular
representations (Sutton et al., 1999b; Bacon et al., 2016), control with
linear function approximation (Moore, 1991), continuous control (Todorov et
al., 2012; Brockman et al., 2016) and vision-based navigation (Chevalier-
Boisvert, 2018).

##  2 Related Work

The discovery of meaningful temporally extended behaviour has previously been
explored through many different avenues. One such avenue is by learning to
reach goals in the environment . This can be done in various ways such as
identifying bottleneck states (McGovern and Barto, 2001; Şimşek et al., 2005;
Menache et al., 2002; Machado et al., 2017) or through feudal RL (Dayan and
Hinton, 1993; Vezhnevets et al., 2017; Nachum et al., 2018).

Another avenue of research is through the options framework (Sutton et al.,
1999b) where the discovery of temporal abstractions can be done directly
through the environment’s reward signal. This framework has lead to recent
approaches (Levy and Shimkin, 2011; Bacon et al., 2016; Harb et al., 2017;
Riemer et al., 2018; Tiwari and Thomas, 2018) that rely on policy gradient
methods to learn the option policies. As policy gradient methods are derived
for the on-policy case, these approaches have only updated the sampled option
within a state. Our work complements them by providing a way to transform
single-option updates into updates that consider all relevant options without
introducing any additional estimators.

Within the options framework, Harutyunyan et al. (2017) have recently
decoupled the way options terminate from the target termination functions.
This allowed for learning longer options as if their duration was less
extended in time. This can be beneficial as shorter options are shown to
produce better solutions and longer options are preferred for planning (Mann
et al., 2015). By contrast, our approach looks at updating all relevant
options through the experience generated by any such option, irrespective of
the duration.

Another line of work is proposed by Daniel et al. (2016) in which options are
considered latent unobserved variables which allows for updating all options.
The authors adopt an expectation-maximization approach which assumes a linear
structure of the option policies. Smith et al. (2018) alleviate this
assumption and derive a policy gradient objective that improves the data-
efficiency and interpretability of the learned options. A conceptually related
work is proposed by Wulfmeier et al. (2020) which leverages dynamic
programming to infer option and action probabilities in hindsight. As these
approaches rely on inference methods to perform updates, it is not directly
compatible with other option learning methods nor with various policy
optimization algorithms, which is another important feature of our approach.

##  3 Background and Notation

We assume a Markov Decision Process $\mathcal{M}$, defined as a tuple
$\langle\mathcal{S},\mathcal{A},r,P\rangle$ with a finite state space
$\mathcal{S}$, a finite action space $\mathcal{A}$, a scalar reward function
$r:\mathcal{S}\times\mathcal{A}\to Dist(\mathbb{R})$ and a transition
probability distribution $P:\mathcal{S}\times\mathcal{A}\to
Dist(\mathcal{S})$. A policy $\pi:\mathcal{S}\to Dist(\mathcal{A})$ specifies
the agent’s behaviour, and the value function is the expected discounted
return obtained by following $\pi$ starting from any state:
$V_{\pi}(s)\dot{=}\mathbb{E}_{\pi}\left[\sum_{i=t}^{\infty}\gamma^{t-i}r(S_{i},A_{i})\middle|S_{t}=s\right]$,
where $\gamma\in[0,1)$ is the discount factor. The policy gradient theorem
(Sutton et al., 1999a) provides the gradient of the expected discounted return
from an initial state distribution $d(s_{0})$ with respect to a parameterized
stochastic policy $\pi_{\zeta}(\cdot|s)$: $\frac{\partial
J(\zeta)}{\partial\zeta}=\sum_{s}d^{\gamma}(s)\sum_{a}\frac{\partial\pi_{\zeta}\left(a\middle|s\right)}{\partial\zeta}Q_{\pi}(s,a)$
where
$d^{\gamma}(s)=\sum_{s_{0}}d(s_{0})\sum_{t=0}^{\infty}\gamma^{t}P^{\pi}(S_{t}=s|S_{0}=s_{0})$
is the discounted state occupancy measure.

Options. The options framework (Sutton et al., 1999b) provides a formalism for
abstractions over time. An option within the option set $\mathcal{O}$ is
composed of an intra-option policy $\pi(a|s,o)$ for selecting actions, a
termination function $\beta(s,o)$ to determine how long an option executes and
an initiation set $I(s,o)$ to restrict the states in which an option may be
selected. It is then possible to choose which option to execute through the
policy over options $\mu:\mathcal{S}\to Dist(\mathcal{O})$. The set of options
along with the policy over options defines an option-value function which can
be written recursively as,

$\displaystyle Q_{\mathcal{O}}(s,o)$
$\displaystyle\dot{=}\sum_{a}\pi(a|s,o)r(s,a)+\gamma\sum_{s^{\prime},o^{\prime}}P_{\pi,\mu,\beta}(s^{\prime},o^{\prime}|s,o)Q_{\mathcal{O}}(s^{\prime},o^{\prime})$

where $P_{\pi,\mu,\beta}(s^{\prime},o^{\prime}|s,o)$ is the probability of
transitioning from $(s,o)$ to any other state-option pair in one step in the
MDP. This probability is defined as,

$\displaystyle
P_{\pi,\mu,\beta}(s^{\prime},o^{\prime}|s,o)\dot{=}p_{\mu,\beta}(o^{\prime}|s^{\prime},o)\sum_{a}P(s^{\prime}|s,a)\pi(a|s,o)$

where
$p_{\mu,\beta}(o^{\prime}|o,s^{\prime})=\beta(s^{\prime},o)\mu(o^{\prime}|s^{\prime})+(1-\beta(s^{\prime},o))\delta_{o,o^{\prime}}$
is the upon-arrival option probability and $\delta$ is the Kronecker delta.
Intuitively, this distribution considers whether or not the previous option
terminates in $s^{\prime}$, and accordingly either chooses an option or
continues with the previous one.

##  4 Learning Multiple Options Simultaneously

In this work we build upon the call-and-return option execution model.
Starting from a state $s$, the agent picks an option according to the policy
over options $o\sim\mu(\cdot|s)$. This option then dictates the action
selection process through the intra-option policy $a\sim\pi(\cdot|s,o)$. Upon
arrival at the next state $s^{\prime}$, the agent queries the option’s
termination function $\beta(s^{\prime},o)$ to decide whether the option
continues or a new option is to be picked. This process then continues until
the episode terminates.

This execution model, coupled with the assumption that the options are Markov,
will let us leverage a notion of how likely an option is to be selected,
irrespective of the actual sampled option. This measure will then be used to
weight the updates of each possible option.

###  4.1 Option Evaluation

In this section we present an algorithm that updates simultaneously the option
value functions of all the options, written as $Q_{\mathcal{O}}(s,o;\theta)$
and parameterized by $\theta$. Let’s recall the on-policy one-step TD update
for state-option pairs,

$\displaystyle\theta_{t+1}$
$\displaystyle=\theta_{t}+\alpha\big{(}R_{t+1}+\gamma\hat{Q}(S_{t+1},O_{t+1})-\hat{Q}(S_{t},O_{t})\big{)}\nabla_{\theta_{t}}\hat{Q}(S_{t},O_{t})$
$\displaystyle\theta_{t+1}$
$\displaystyle=\theta_{t}+\alpha\big{(}R_{t+1}+\gamma\theta_{t}^{\top}\phi(S_{t+1},O_{t+1})-\theta_{t}^{\top}\phi(S_{t},O_{t})\big{)}\phi(S_{t},O_{t})$
$\displaystyle=\theta_{t}+\alpha\big{(}\underbrace{R_{t+1}\phi(S_{t},O_{t})}_{\textbf{b}_{t}}+\underbrace{\phi(S_{t},O_{t})(\phi(S_{t},O_{t})-\gamma\phi(S_{t+1},O_{t+1}))^{\top}}_{\textbf{A}_{t}}\theta_{t}\big{)}$
$\displaystyle\theta_{t+1}$
$\displaystyle=\theta_{t}+\alpha(\textbf{b}_{t}-\textbf{A}_{t}\theta_{t})$ (1)

where the option value functions are linear functions, i.e.
$Q_{\mathcal{O}}(s,o;\theta)\dot{=}\theta_{t}^{\top}\phi(s,o)$ of the
parameters $\theta_{t}$ at time $t$ and the features $\phi(s,o)$. Since the
random variables within the update rule, summarized through the vector
$\textbf{b}_{t}$ and the matrix $\textbf{A}_{t}$, are sampled from interacting
with the environment, the stability and convergence properties can be verified
by inspecting its expected behaviour in the limit, written as,

$\displaystyle\theta_{t+1}$
$\displaystyle=\theta_{t}+\alpha(\textbf{b}-\textbf{A}\theta_{t})$

where $\textbf{A}=\lim_{{t}\rightarrow\infty}\mathbb{E}[\textbf{A}_{t}]$ and
$\textbf{b}=\lim_{{t}\rightarrow\infty}\mathbb{E}[\textbf{b}_{t}]$. As the
matrix A has a multiplicative effect on the parameters, it will be crucial for
this matrix to be positive-definite in order to assure the stability of the
algorithm, and subsequently its convergence. In the on-policy case, this
matrix is written as:

A
$\displaystyle=\sum_{s,o}\lim_{t\rightarrow\infty}P_{\pi,\mu,\beta}(S_{t}=s,O_{t}=o)\phi(s,o)\bigg{(}\phi(s,o)-\gamma\sum_{s^{\prime},o^{\prime}}P_{\pi,\mu,\beta}(s^{\prime},o^{\prime}|s,o)\phi(s^{\prime},o^{\prime})\bigg{)}^{\top}$

where $P_{\pi,\mu,\beta}(S_{t}=s,O_{t}=o)$ is the probability of being in a
particular state-option pair at time $t$. In order to update all options, we
will exploit the form of this state-option distribution to preserve the
positive definite property of the matrix A. In the following lemma we first
show that this limiting distribution exists under some assumptions and we
derive a useful decomposition. To do so, we introduce the following two
assumptions. Regarding Assumption 1, it is common in standard RL to assume
that the policy induces ergodicity on the state space, which we extend to all
option policies. Assumption 2 is specific to the hierarchical setting and is a
way to ensure irreducibility, that is, that each state-option pair is
reachable. It would be possible to find less strict assumptions by having
access to more information about how the MDP is constructed.

######  Assumption 1.

The option policies $\\{\pi(a|s,o):\forall o\in\mathcal{O}\\}$ each induce an
ergodic Markov chain on the state space $\mathcal{S}$.

######  Assumption 2.

The termination function $\beta(s,o)$ and policy over options $\mu(o|s)$ have
strictly positive probabilities in all state-option pairs.

######  Lemma 1.

Under assumptions 1-2 the following limit exists,

$\displaystyle\lim_{t\rightarrow\infty}P_{\pi,\mu,\beta}(S_{t}=s,O_{t}=o|S_{0},O_{0})=d_{\pi,\mu,\beta}(s,o)$
(2)

and the augmented chain process is an ergodic Markov chain in the augmented
state-option space. Moreover the stationary distribution can be decomposed in
the following form,

$\displaystyle
d_{\pi,\mu,\beta}(s,o)=\sum_{\bar{o}}\bar{d}_{\pi,\mu,\beta}(\bar{o},s)p_{\mu,\beta}(o|s,\bar{o})$
(3)

where $p_{\mu,\beta}(o|s,\bar{o})$ is the probability over options upon
arriving in a state.

Remarks. In the derivation we have avoided mentioning the concept of
initiation sets such that we would not encumber the notation and presentation.
The derivation would be identical if we had assumed that the initation set for
each option is equal to the state space, which has been done for most of
recent work on options (Bacon et al., 2016). Another way to reconcile our
method with the original options framework would be to consider interest
functions (Khetarpal et al., 2020), which are a generalization of the concept
of initiation sets. We could then add an assumption similar to Assumption 2 on
the interest function by only allowing positive values. This would let us
heavily favor a subset of options for a certain region of the state space.
However, this assumption would still preclude us from using hard constraints
on the option set, such as option affordances (Khetarpal et al., 2021). More
empirical evidence would be necessary to investigate whether this is a
limitation in practice.

Ideally, if we seek to update all options in a state, we could weight each of
these updates by the stationary distribution $d_{\pi,\mu,\beta}(s,o)$.
However, in most case, it is impractical to assume access to this
distribution. A possible solution would be to estimate this quantity by making
use of recent approaches (Hallak and Mannor, 2017; Xie et al., 2019; Zhang et
al., 2020). Instead, we propose to leverage the decomposition presented in
Equation 3.

Indeed, this decomposition highlights the dependency on the upon-arrival
option distribution, $p_{\mu,\beta}(o|s,\bar{o})$, which is a known quantity
as it is only dependent of the agent’s policy over options $\mu(o|s)$ and
termination function $\beta(s,\bar{o})$. Although
$\bar{d}_{\pi,\mu,\beta}(s,\bar{o})$ is itself unknown, we can collect its
samples by adopting the call-and-return option execution model. More
concretely, under this model, an option which is previously selected leads us
to a state of interest at time $t$. These random variables, being sampled on-
policy, follow the distribution $\bar{d}_{\pi,\mu,\beta}(s,\bar{o})$. Upon
arriving in any state of interest $S_{t}$ with option $O_{t-1}$, although an
option $O_{t}$ will be sampled according to call-and-return, we can use the
distribution $p_{\mu,\beta}(O_{t}=o|S_{t},O_{t-1})$ to estimate how likely it
would have been for each option $o$ to be selected. The update rule would
become,

$\displaystyle\theta_{t+1}$
$\displaystyle=\theta_{t}+\alpha\sum_{o}p_{\mu,\beta}(O_{t}=o|S_{t},O_{t-1})\big{(}\mathbb{E}[U_{t}|S_{t},O_{t}=o]-\theta_{t}^{\top}\phi(S_{t},O_{t}=o)\big{)}\phi(S_{t},O_{t}=o)$
(4)

where $\mathbb{E}[U_{t}|S_{t}=s,O_{t}=o]$ is the target for a given state and
option pair. In Section 4.3 we show how to obtain this quantity without using
any additional estimation or simulation-resetting. Our update rule would then
preserve the positive-definite property of the matrix A.

###  4.2 Option Control

In this section we extend the approach from evaluating the option value
functions to the case where we seek to learn all option policies
simultaneously. To do so we build upon the option-critic (OC) update rules
derived by Bacon et al. (2016). Let us first recall the intra-option policy
gradient,

$\displaystyle\sum_{s,o}d^{\gamma}_{\pi,\mu,\beta}(s,o)\sum_{a}\frac{\partial\pi_{\zeta}(a|s,o)}{\partial\zeta}Q_{\mathcal{O}}(s,o,a)$
(5)

where,
$d^{\gamma}_{\pi,\mu,\beta}(s,o)=\sum_{t}\gamma^{t}P_{\pi,\mu,\beta}(S_{t}=s,O_{t}=o)$,
is the $\gamma$-discounted occupancy measure over state-option pairs. As in
the previous section, we wish to adopt the same call-and-return option
execution model, in particular as it is also used by the option-critic
algorithm in order to perform stochastic gradient updates following Equation
5. When updating all options within a state $s$, we notice that we can
leverage the structure of distribution state-option occupancy measure,

$\displaystyle
d^{\gamma}_{\pi,\mu,\beta}(s,o)=\sum_{\bar{o}}p_{\mu,\beta}(o|s,\bar{o})\sum_{\bar{s}}p_{\pi}(s|\bar{s})d^{\gamma}_{\pi,\mu,\beta}(\bar{s},\bar{o})$
(6)

By expanding the state-occupancy measure we can separate unknown
probabilities, that is the next state distribution $p_{\pi}(s|\bar{s})$ and
the previous state-option pair $d^{\gamma}_{\pi,\mu,\beta}(\bar{s},\bar{o})$,
from the known upon-arrival option distribution, that is
$p_{\mu,\beta}(o|s,\bar{o})$. Executing Markov options in a call-and-return
fashion will then provide us samples from the unknown distributions, and for
each sampled $(S_{t},O_{t-1})$-pair we can weight the likelihood of choosing
option $o$ at time $t$. Concretely, this translates into the following update
rule,

$\displaystyle\zeta_{t+1}=\zeta_{t}+\alpha_{\zeta}\sum_{o}p_{\mu,\beta}(O_{t}=o|S_{t},O_{t-1})\frac{\log\partial\pi_{\zeta}(A_{t}|S_{t},O_{t}=o)}{\partial\zeta_{t}}Q_{\mathcal{O}}(S_{t},O_{t}=o,A_{t})$
(7)

We highlight that the update rules for option control do not rely on
Assumption 1 and 2, as these were specific to the evaluation setting. By
updating all options within a state, we can expect better sample efficiency
and final performance. Moreover, this may also contribute to reducing
variance. Indeed, the update rule in Equation 7 can be seen as a variation on
the all-action policy gradient (Sutton et al., 1983; Petit et al., 2019),
which updates all actions applied to the distribution of options. Such all-
action or all-option updates can be derived as applications of the extended
(Bratley et al., 1987) form of conditional Monte-Carlo estimators (Hammersley,
1956; Liu et al., 2019), which provide a conditioning effect on the updates
that tend to reduce the variance. As we will see in our experiments, such
variance reduction can be observed, especially with tabular and linear
function estimators.

It is possible to combine evaluation and control into a single method. We do
so by instantiating these update rules within the option-critic (OC) algorithm
which we present in Algorithm 1 in appendix A.1 and name it Multi-updates
Option Critic (MOC) algorithm.

###  4.3 Suitable targets

We have proposed a way in which all relevant options can be updated
simultaneously. In the case of option evaluation we relied on having access to
appropriate targets $\mathbb{E}[U_{t}|S_{t}=s,O_{t}=o]$ for each option within
a state. However, as we execute options in the call-and-return model and we do
not assume access to the simulator for arbitrary resetting, we only have one
sampled action for each state: the one taken by the executing option. To
leverage this experience for any arbitrary option, we will leverage standard
importance sampling (Precup et al., 2000). Using importance sampling naturally
excludes the possibility that options be defined through deterministic
policies, which can be seen as a limitation of our approach. That being said,
importance sampling is not necessary for the option control algorithm.

The importance sampling target for any arbitrary option $\tilde{o}$ given the
action $a$ in state $s$ generated by the executing option $o$ is the
following,

$\displaystyle U^{\rho}_{t}(s,\tilde{o})$
$\displaystyle=\dfrac{\pi(A_{t}=a|S_{t}=s,O_{t}=\tilde{o})}{\pi(A_{t}=a|S_{t}=s,O_{t}=o)}R_{t+1}+\dfrac{p_{\mu,\beta}(O_{t+1}=o^{\prime}|S_{t+1}=s^{\prime},O_{t}=\tilde{o})}{p_{\mu,\beta}(O_{t+1}=o^{\prime}|S_{t+1}=s^{\prime},O_{t}=o)}\gamma\theta_{t}^{\top}\phi(s^{\prime},o^{\prime})$

It is straightforward to show that this target is unbiased and corrects for
the off-policyness introduced by updating any other option than the sampled
one, leading to consistency (Sutton et al., 1999b). It would be possible to
consider different approaches to correct for difference between intra-option
policies (Mahmood et al., 2014; Li et al., 2015; Harutyunyan et al., 2016;
Jiang and Li, 2016; Jain and Precup, 2018) which we leave for future work.

###  4.4 Controlling the number of updates

A possible drawback of updating all options is the prospect of reducing the
degree of diversity in the option set. A possible remedy would be to favor
diversity, or similar measures, through reward shaping (Gregor et al., 2016;
Eysenbach et al., 2018; Pathak et al., 2019; Klissarov and Precup, 2020). In
this work, we instead introduce a hyperparameter $\eta$ which represents the
probability of updating all options, as opposed to updating only the sampled
option. In our experiments we will notice that high values of $\eta$ are not
necessarily detrimental to the expressivity of the option set.

Another interesting property of this hyperparameter, which we do not leverage
in this work, is that it could be use in a state dependent way (i.e.
$\eta(s)$) without changing the stability of the algorithm. This
generalization adds flexibility to how we want to update options. In
particular, it could be a way to limit the variance of the algorithm by only
performing such off-policy updates for state-option pairs where the importance
sampling is below a threshold.

##  5 Experiments

We now validate our method on a wide range of tasks and try to assess the
following questions: (1) By updating all relevant options are we able to
improve the sample-efficiency and final performance? (2) Does our approach
induce temporally extended options without the need of regularization? (3) Is
the learned option set diverse and useful?

Most of our experiments are done in a transfer learning setting. Our method is
compatible with any modern policy optimization framework and we validate this
by building our hierarchichal implementations on top of the synchronous
variant of the asynchronous advantage actor-critic algorithm (Mnih et al.,
2016), that is the A2C algorithm, as well as the proximal policy optimization
algorithm (PPO) (Schulman et al., 2017). All hyperparameters are available in
the appendix. Importantly, for each set of experiments we use the same value
of $\eta$. For all experiments, the shaded area represents the $80\%$
confidence interval.

We would also like to emphasize that although our implementation is based upon
the option-critic algorithm (Bacon et al., 2016), our approach can be
straightforwardly adapted to most hierarchical RL algorithms (Riemer et al.,
2018; Jain et al., 2018; Jinnai et al., 2019; Zhang and Whiteson, 2019;
Bagaria and Konidaris, 2019).

###  5.1 Tabular Domain

Experimental Setup: We first evaluate our approach in the classic FourRooms
domain (Sutton et al., 1999b) where the agent starts in a random location and
has to navigate to the goal in the right hallway. As we are interested in the
algorithms’ sample efficiency, we plot the number of steps taken per episodes.

In Figure 1(a), we work under the assumption that the options have been given
and remain fixed throughout learning. The only element being updated is the
option value functions, $Q_{\mathcal{O}}(s,o)$, which is used by the policy
over options to select among options. We compare our approach, leveraging the
multi-options update rule in Equation 4, to the case of only updating the
sampled option within a state. The options used are the 12 hallway options
designed in Sutton et al. (1999b) (see their Figure 2 for a brief summary). In
this setting, our approach achieves significantly better sample efficiency,
illustrating how our method can accelerate learning the utility of a set of
well-suited options.

We now move to the case where all the options’ parameters are learned and
leverage the learning rule for control in Equation 7. In this setting we
compare our approach, multi-updates option critic (MOC), to the option-critic
(OC) algorithm as well as a flat agent, denoted as actor-critic (AC). Both
hierarchical agents outperform the flat baselines, however, our agent is able
to reach the same performance as OC in about half the number of episodes.
Moreover, we notice that MOC shows less variance across runs, which is in
accordance with previous work investigating the variance reduction effect of
expected updates (Sutton et al., 1983; Petit et al., 2019).

(a)

(b)

(c)

Figure 1: FourRooms domain. In a) we used the fixed set of hallway options as
defined in Sutton et al. (1999b) and only learned the option-value functions.
In b) all the parameters are learned from scratch before the goal location is
change mid-training. In c) we plot the information radius between options as a
measure of diversity. Results are averaged over 50 independent runs.

In Figure 6 in the appendix, we verify the kind of options that are being
learned under both the option-critic and our approach. For each option we plot
the best action in each state and we highlight in yellow whether the option is
likely to be selected by the policy over options. We first notice that for the
OC agent, there is an option that is selected much more often, whereas the MOC
agent learns a good separation of the state space. Moreover, we notice in
purple states (where the option is less likely to be chosen), the OC agent’s
actions do not seem to lead to a meaningful behaviour in the environment,
which is in contrast to our approach. As a potential benefit, by sharing
information between options, the MOC agent will tend to be more robust to the
environment’s stochasticity as well as to the learning process of the policy
over options.

Finally, we also would like to point out a possible challenge in our approach
that we have mentioned in Section 4. It is possible that by learning all
options simultaneously we may reduce the diversity between options. We verify
this in Figure 1(c) where we plot the information radius (Sibson, 1969)
between options as a measure of diversity. We plot the OC agent as well as our
approach, MOC, with different values of the $\eta$ hyperparameter which
controls the degree to which we apply updates to all options. We notice indeed
a clear different between the OC agent and our agent when $\eta=1.0$ (only
updates to all options). However, intermediate values of $\eta$ provide an
effective to increase the information radius. As in this experiment we use
tabular representations, the difference between updated options and non-
updated options is accentuated. In the function approximation case we will see
that the concern about the diversity of options is mitigated.

###  5.2 Linear Function Approximation

Experimental Setup: We now perform experiments in the case where the value
function is estimated through linear function approximation. The environment
we consider is a sparse variant of the classic Mountain Car domain (Moore,
1991; Sutton and Barto, 2018) where a positive reward is given upon reaching
the goal. We consider the transfer setting where after half the total number
of timesteps the gravity is doubled.

(a)

(b)

(c)

Figure 2: MountainCar domain. In a) we plot the performance of our agent, MOC,
the option-critic and a flat baseline. The options are learned from scratch.
At mid-training, we modify the gravity to double its original value. In b) we
perform a hyperparamter study on the value of $\eta$ which controls the amount
of updates to all options. In c) we plot the information radius between
options as a measure of expressivity and notice that higher values of $\eta$
do not affect it negatively. The performance is averaged across 20 independent
runs.

In Figure 2(a), we observe similar results to the tabular case, that is,
hierarchical agents perform better than standard RL in the transfer task. MOC
also improves upon OC’s performance and sample efficiency. We verify how this
improvement in performance translates with respect to the value of $\eta$, the
hyperparameter controlling for the probability of updating all options. In
Figure 2(b) we can observe an almost monotonic relationship between $\eta$ and
the algorithm’s final performance, which confirms the effectiveness of our
approach when using function approximation.

Finally, in Figure 2(c) we verify the information radius between options for
different values of $\eta$. We notice that unlike the tabular case where
higher values of $\eta$ lead to a smaller information radius, there is no
clear correlation in this case. Perhaps surprisingly we notice that in the
first task high values of $\eta$ lead to an increased information radius. This
might indicate that at the beginning of training, the agent can learn to
specialize faster by updating all options.

###  5.3 Experiments at Scale

####  5.3.1 Continuous Control

Experimental Setup: We turn our attention to the deep RL setting where all
option components are estimated through neural networks. We first perform
experiments in the MuJoco domain (Todorov et al., 2012) where both states and
actions are continuous variables. We build on classic tasks (Brockman et al.,
2016) and modify them to make the learning process more challenging. The
description of each task as well as a visual representation is available in
the appendix.

We leverage the PPO algorithm (Schulman et al., 2017) to build our
hierarchical agents. We also compare our approach, MOC, to the interest
option-critic (IOC) (Khetarpal et al., 2020) which has been shown to
outperform the option-critic algorithm derived for the continuous control
setting (Harb et al., 2017; Klissarov et al., 2017). Additionally, we compare
to a flat agent using PPO and to the inferred option policy gradient (IOPG)
algorithm (Smith et al., 2018), an HRL algorithm also designed for updating
all options by leveraging inference.

(a)

(b)

(c)

(d)

(e)

(f)

(g)

Figure 3: Transfer on MuJoco. In a-b-c) we compare the performance of our
approach, MOC, to various hierarchical baselines and to a flat agent. At about
a quarter into the training process, the original task is modified to a more
challenging, but related, setting where the agents must transfer their
knowledge. In d-e-f-g), we plot the option duration for both the MOC agent and
the OC baseline. We notice that our approach tends to produce options that are
extended in time, which can lead to meaningful behavior. All experiments are
plotted using 10 independent runs.

Across environments in Figure 3abc we witness that MOC shows better sample
efficient and performance compared to the other hierarchical agents as well as
the PPO baseline. In particular, although IOPG updates all options similarly
to our approach, this doesn’t translate in a practical advantage. This
highlights the adaptability of the update rules we propose to any state-of-
the-art policy optimization algorithm. We also emphasize that in all these
experiments, the same value of $\eta$ is used, making it a general approach.

(a)

(b)

Figure 4: a) Option visualization. In this experiment, AntWalls-v1, the agent
has first learned how to walk in the forward direction before we add obstacles
(walls) which it needs to circumvent. We highlight in purple the states in
which our agent uses the first option and in yellow the second option. b)
Performance of MOC and IOC algorithms for different neural network
architectures. Independence indicates that the second last hidden layer does
is not shared across optiosn.

In Figure 3defg we investigate the options’ duration for both the IOC agent
and our approach. We investigate both methods using two options, so as to
provide more interpretable results We argue that such a metric can be an
important indicator of the usefulness of options, which are meant to be
temporally extended. We notice that both in HalfCheetahDir-v1 and AntWalls-v1,
both options learned by MOC are used by the agent over multiple timesteps.
This is in contrast to the IOC agent where options at the end of training
exhibit an imbalance. We also highlight that although MOC improves this aspect
of option learning, the shaded are is still large which indicates that it does
not fully solve the problem.

For the MOC agent, it is interesting to notice temporally extended options are
truly present once the agent faces a new task. This perhaps suggests that
future inquiry should look into whether challenging the agent with multiple
tasks (such as in the continual learning setting) is actually a catalyst for
the discovery of options.

In Figure 4a, we provide a visualization of meaningful behaviour exhibited by
the MOC agent. We notice that in the AntWalls-v1 environment, one agent learns
an option to move forward (highlighted in purple), while the other one is used
to move around the walls (highlighted in yellow). We provide additional
visualization (including from the IOC agent) in the appendix. We mention that
although such meaningful behaviour was witnessed more often in our approach,
we did not completely avoid degenerate solutions for all random seeds or all
values of the learning rate.

Table 1: Walker2dStand-v1 Algorithm | 2 options | 4 options | 8 options
---|---|---|---
MOC | 617 (98) | 574 (76) | 535 (81)
IOC | 326 (92) | 238 (55) | 212 (64)

In our experiments we compared HRL algorithms when using two options. We now
verify how MOC and IOC scale when we grow the option set to four and eight
options. In Table LABEL:options, we notice that most of the gains from
updating all options are persistent for each cardinality. However, we notice
eventually a significant drop in performance when using 8 options. This could
attributed to the fact that the domains we experiment with are simple enough
that such a high number of options is unnecessary and leads to reduced sample
efficiency. One way to verify this would be to apply the same algorithm in a
continual learning setting. It could also be the case that the OC baseline on
which we implement our flexible update rules could be itself improved for
better performance. More experiments are needed to investigate these
questions, which we leave for future work.

Finally, a promising possibility of our method is to re-organize the way
parameters are shared between options. The neural network architecture used by
OC-based algorithms (including ours) shares most of the parameters between
options: only each head (last layer of the network) is specific to each option
and independent of others. In Figure 4b, we investigate the performance of MOC
and IOC when the parameters of the second to last hidden layer are also
independent for each option. As we notice through the dashed lines, MOC
remains more robust to this configuration. This would indicate that by using
our proposed update rules, more sophisticated network architectures could be
devised without sacrificing sample efficiency.

####  5.3.2 MiniWorld Vision Navigation

Experimental Setup: In this section we verify how our approach extends to
tasks with high dimensional inputs such as the first-person visual navigation
domain of MiniWorld (Chevalier-Boisvert, 2018). We provide a description for
each environment in the appendix.

In MiniWorld-YMazeTransfer-v0, the transfer is not drastic and the flat agent
is able to recover well in the transfer task. Moreover, MOC seemed to be
slower to learn at first, although it was able to transfer with a smaller drop
in performance. For the more challenging MiniWorld-WallGapTransfer-v0, we
notice that HRL algorithms are more relevant and adapt better to the change,
with MOC being more sample efficient.

(a)

(b)

(c)

Figure 5: Transfer in MiniWorld. In a-b) we plot the performance of our
approach and compare it to the option-critic (OC) and a flat baseline (PPO).
In particular in b) we can appreciate the need for hierarchy when transferring
to a similar, but more challenging task. In c) we present qualitative results
on the kind of options that are being learned by our agent. In these
experiments we average the performance across six independent runs.

Finally, in 5c, we plot the options learned by our MOC agent in the MiniWorld-
YMazeTransfer-v0 environment. The first option learns to get to the middle of
the maze and locate the red goal. When the goal has been seen, the second
option moves towards it. Learning such a structure can potentially help an
agent adapt to a new and related task.

##  6 Conclusion

In this work we have proposed a way to update all relevant options for a
visited state. Our method does not introduce any additional estimators and is
therefore readily applicable to many state-of-the-art RL, and HRL, algorithms
in order to improve their performance and sample-efficiency. Updating all
options simultaneously has been an integral part of the intra-option learning
proposed in the options framework (Sutton et al., 1999b). This work therefore
aims to extend this flexibility to function approximation and deep RL.

Our proposed update rules can be seen as updates done in expectation with
respect to options. As such, they are reminescent of generalizations of
standard RL algorithms, such Expected SARSA (van Seijen et al., 2009) and
Expected Policy Gradients (Ciosek and Whiteson, 2018). Amongst other
qualities, such approaches have theoretically been shown to reduce variance. A
similar analysis could be performed in our framework to verify what
theoretical guarantees can be obtained from updating all relevant options.

We also proposed a flexible way to control between updating all relevant
options and updating only the sampled one through the hyperparameter $\eta$.
Potential future work could investigate whether a state-dependent $\eta$
function, similarly to interest functions (Mahmood et al., 2015), could be
used to highlight important transitions, such as bottlenecks, and update all
options only in those states. Finally, learning all options simultaneously
could be further leveraged in the case where we learn a very large number of
them simultaneously (Precup, 2018; Schaul et al., 2015), before deciding which
are useful. Such an investigation would most likely be most promising in a
continual learning setting. Updating all relevant options also means that we
could implement less parameter sharing between options and discover neural
network architectures better suited for HRL.

## Acknowledgments and Disclosure of Funding

The authors would like to thank the National Science and Engineering Research
Council of Canada (NSERC) for funding this research; Khimya Khetarpal for
early discussions on the project and the anonymous reviewers for providing
critical and constructive feedback.

## References

  * Bacon et al. (2016) Pierre-Luc Bacon, Jean Harb, and Doina Precup.  The option-critic architecture.  _CoRR_ , abs/1609.05140, 2016.  URL http://arxiv.org/abs/1609.05140.
  * Bagaria and Konidaris (2019) Akhil Bagaria and George Konidaris.  Option discovery using deep skill chaining.  In _International Conference on Learning Representations_ , 2019.
  * Bratley et al. (1987) Paul Bratley, Bennett L. Fox, and Linus E. Schrage.  _A Guide to Simulation (2nd Ed.)_.  Springer-Verlag, Berlin, Heidelberg, 1987.  ISBN 0387964673.
  * Brockman et al. (2016) Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojciech Zaremba.  Openai gym.  _CoRR_ , abs/1606.01540, 2016.  URL http://arxiv.org/abs/1606.01540.
  * Chevalier-Boisvert (2018) Maxime Chevalier-Boisvert.  gym-miniworld environment for openai gym.  https://github.com/maximecb/gym-miniworld, 2018.
  * Ciosek and Whiteson (2018) Kamil Ciosek and Shimon Whiteson.  Expected policy gradients, 2018.
  * Şimşek et al. (2005) Özgür Şimşek, Alicia P. Wolfe, and Andrew G. Barto.  Identifying useful subgoals in reinforcement learning by local graph partitioning.  In _Proceedings of the 22nd International Conference on Machine Learning_ , ICML ’05, page 816–823, New York, NY, USA, 2005. Association for Computing Machinery.  ISBN 1595931805.  doi: 10.1145/1102351.1102454.  URL https://doi.org/10.1145/1102351.1102454.
  * Daniel et al. (2016) Christian Daniel, Herke Van Hoof, Jan Peters, and Gerhard Neumann.  Probabilistic inference for determining options in reinforcement learning.  _Machine Learning_ , 104(2):337–357, 2016.
  * Dayan and Hinton (1993) Peter Dayan and Geoffrey E Hinton.  Feudal reinforcement learning.  In S. Hanson, J. Cowan, and C. Giles, editors, _Advances in Neural Information Processing Systems_ , volume 5. Morgan-Kaufmann, 1993.  URL https://proceedings.neurips.cc/paper/1992/file/d14220ee66aeec73c49038385428ec4c-Paper.pdf.
  * Dhariwal et al. (2017) Prafulla Dhariwal, Christopher Hesse, Oleg Klimov, Alex Nichol, Matthias Plappert, Alec Radford, John Schulman, Szymon Sidor, Yuhuai Wu, and Peter Zhokhov.  Openai baselines.  https://github.com/openai/baselines, 2017.
  * Eysenbach et al. (2018) Benjamin Eysenbach, Abhishek Gupta, Julian Ibarz, and Sergey Levine.  Diversity is all you need: Learning skills without a reward function.  _CoRR_ , abs/1802.06070, 2018.  URL http://arxiv.org/abs/1802.06070.
  * Fujimoto et al. (2018) Scott Fujimoto, Herke van Hoof, and David Meger.  Addressing function approximation error in actor-critic methods.  _CoRR_ , abs/1802.09477, 2018.  URL http://arxiv.org/abs/1802.09477.
  * Gregor et al. (2016) Karol Gregor, Danilo Jimenez Rezende, and Daan Wierstra.  Variational intrinsic control.  _CoRR_ , abs/1611.07507, 2016.  URL http://arxiv.org/abs/1611.07507.
  * Hallak and Mannor (2017) Assaf Hallak and Shie Mannor.  Consistent on-line off-policy evaluation, 2017.
  * Hammersley (1956) J. M. Hammersley.  Conditional monte carlo.  _J. ACM_ , 3(2):73–76, April 1956.  ISSN 0004-5411.  doi: 10.1145/320825.320827.  URL https://doi.org/10.1145/320825.320827.
  * Harb et al. (2017) Jean Harb, Pierre-Luc Bacon, Martin Klissarov, and Doina Precup.  When waiting is not an option : Learning options with a deliberation cost.  _CoRR_ , abs/1709.04571, 2017.  URL http://arxiv.org/abs/1709.04571.
  * Harutyunyan et al. (2016) Anna Harutyunyan, Marc G. Bellemare, Tom Stepleton, and Remi Munos.  Q($\lambda$) with off-policy corrections, 2016.
  * Harutyunyan et al. (2017) Anna Harutyunyan, Peter Vrancx, Pierre-Luc Bacon, Doina Precup, and Ann Nowe.  Learning with options that terminate off-policy, 2017.
  * Jain et al. (2018) Arushi Jain, Khimya Khetarpal, and Doina Precup.  Safe option-critic: Learning safety in the option-critic architecture.  _CoRR_ , abs/1807.08060, 2018.  URL http://arxiv.org/abs/1807.08060.
  * Jain and Precup (2018) Ayush Jain and Doina Precup.  Eligibility traces for options.  In _Proceedings of the 17th International Conference on Autonomous Agents and MultiAgent Systems_ , AAMAS ’18, page 1008–1016, Richland, SC, 2018. International Foundation for Autonomous Agents and Multiagent Systems.
  * Jiang and Li (2016) Nan Jiang and Lihong Li.  Doubly robust off-policy value evaluation for reinforcement learning.  In Maria Florina Balcan and Kilian Q. Weinberger, editors, _Proceedings of The 33rd International Conference on Machine Learning_ , volume 48 of _Proceedings of Machine Learning Research_ , pages 652–661, New York, New York, USA, 20–22 Jun 2016. PMLR.  URL http://proceedings.mlr.press/v48/jiang16.html.
  * Jinnai et al. (2019) Yuu Jinnai, Jee Won Park, David Abel, and George Konidaris.  Discovering options for exploration by minimizing cover time.  In Kamalika Chaudhuri and Ruslan Salakhutdinov, editors, _Proceedings of the 36th International Conference on Machine Learning_ , volume 97 of _Proceedings of Machine Learning Research_ , pages 3130–3139. PMLR, 09–15 Jun 2019.  URL http://proceedings.mlr.press/v97/jinnai19b.html.
  * Khetarpal et al. (2020) Khimya Khetarpal, Martin Klissarov, Maxime Chevalier-Boisvert, Pierre-Luc Bacon, and Doina Precup.  Options of interest: Temporal abstraction with interest functions.  _CoRR_ , abs/2001.00271, 2020.  URL http://arxiv.org/abs/2001.00271.
  * Khetarpal et al. (2021) Khimya Khetarpal, Zafarali Ahmed, Gheorghe Comanici, and Doina Precup.  Temporally abstract partial models, 2021.
  * Klissarov and Precup (2020) Martin Klissarov and Doina Precup.  Reward propagation using graph convolutional networks.  In H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin, editors, _Advances in Neural Information Processing Systems_ , volume 33, pages 12895–12908. Curran Associates, Inc., 2020.  URL https://proceedings.neurips.cc/paper/2020/file/970627414218ccff3497cb7a784288f5-Paper.pdf.
  * Klissarov et al. (2017) Martin Klissarov, Pierre-Luc Bacon, Jean Harb, and Doina Precup.  Learnings options end-to-end for continuous action tasks.  _CoRR_ , abs/1712.00004, 2017.  URL http://arxiv.org/abs/1712.00004.
  * Kostrikov (2018) Ilya Kostrikov.  Pytorch implementations of reinforcement learning algorithms.  https://github.com/ikostrikov/pytorch-a2c-ppo-acktr-gail, 2018.
  * Levy and Shimkin (2011) Kfir Y Levy and Nahum Shimkin.  Unified inter and intra options learning using policy gradient methods.  In _European Workshop on Reinforcement Learning_ , pages 153–164. Springer, 2011.
  * Li et al. (2015) Lihong Li, Remi Munos, and Csaba Szepesvari.  Toward minimax off-policy value estimation.  In _Proceedings of the 18th International Conference on Artificial Intelligence and Statistics (AISTATS)_. JMLR: Workshop and Conference Proceedings, May 2015.  URL https://www.microsoft.com/en-us/research/publication/toward-minimax-off-policy-value-estimation/.
  * Liu et al. (2019) Yao Liu, Pierre-Luc Bacon, and Emma Brunskill.  Understanding the curse of horizon in off-policy evaluation via conditional importance sampling.  _CoRR_ , abs/1910.06508, 2019.  URL http://arxiv.org/abs/1910.06508.
  * Machado et al. (2017) Marlos C Machado, Marc G Bellemare, and Michael Bowling.  A laplacian framework for option discovery in reinforcement learning.  _arXiv preprint arXiv:1703.00956_ , 2017.
  * Mahmood et al. (2014) A. Rupam Mahmood, Hado P van Hasselt, and Richard S Sutton.  Weighted importance sampling for off-policy learning with linear function approximation.  In Z. Ghahramani, M. Welling, C. Cortes, N. Lawrence, and K. Q. Weinberger, editors, _Advances in Neural Information Processing Systems_ , volume 27. Curran Associates, Inc., 2014.  URL https://proceedings.neurips.cc/paper/2014/file/be53ee61104935234b174e62a07e53cf-Paper.pdf.
  * Mahmood et al. (2015) Ashique Rupam Mahmood, Huizhen Yu, Martha White, and Richard S. Sutton.  Emphatic temporal-difference learning.  _CoRR_ , abs/1507.01569, 2015.  URL http://arxiv.org/abs/1507.01569.
  * Mann et al. (2015) Timothy A. Mann, Shie Mannor, and Doina Precup.  Approximate value iteration with temporally extended actions.  _J. Artif. Intell. Res._ , 53:375–438, 2015.
  * McGovern and Barto (2001) Amy McGovern and Andrew G Barto.  Automatic discovery of subgoals in reinforcement learning using diverse density.  In _ICML_ , volume 1, pages 361–368, 2001.
  * Menache et al. (2002) Ishai Menache, Shie Mannor, and Nahum Shimkin.  Q-cut-dynamic discovery of sub-goals in reinforcement learning.  In _European Conference on Machine Learning_ , pages 295–306. Springer, 2002.
  * Mnih et al. (2016) Volodymyr Mnih, Adrià Puigdomènech Badia, Mehdi Mirza, Alex Graves, Timothy P. Lillicrap, Tim Harley, David Silver, and Koray Kavukcuoglu.  Asynchronous methods for deep reinforcement learning.  _CoRR_ , abs/1602.01783, 2016.  URL http://arxiv.org/abs/1602.01783.
  * Moore (1991) Andrew Moore.  _Efficient Memory-based Learning for Robot Control_.  PhD thesis, Carnegie Mellon University, Pittsburgh, PA, March 1991.
  * Nachum et al. (2018) Ofir Nachum, Shixiang Gu, Honglak Lee, and Sergey Levine.  Data-efficient hierarchical reinforcement learning.  _CoRR_ , abs/1805.08296, 2018.  URL http://arxiv.org/abs/1805.08296.
  * Pathak et al. (2019) Deepak Pathak, Dhiraj Gandhi, and Abhinav Gupta.  Self-supervised exploration via disagreement.  _CoRR_ , abs/1906.04161, 2019.  URL http://arxiv.org/abs/1906.04161.
  * Petit et al. (2019) Benjamin Petit, Loren Amdahl-Culleton, Yao Liu, Jimmy Smith, and Pierre-Luc Bacon.  All-action policy gradient methods: A numerical integration approach.  _CoRR_ , abs/1910.09093, 2019.  URL http://arxiv.org/abs/1910.09093.
  * Precup (2018) Doina Precup.  Temporal abstraction.  2018\.  URL http://videolectures.net/site/normal_dl/tag=1199094/DLRLsummerschool2018_precup_temporal_abstraction_01.pdf.
  * Precup et al. (2000) Doina Precup, Richard S. Sutton, and Satinder P. Singh.  Eligibility traces for off-policy policy evaluation.  In _Proceedings of the Seventeenth International Conference on Machine Learning_ , ICML ’00, page 759–766, San Francisco, CA, USA, 2000. Morgan Kaufmann Publishers Inc.  ISBN 1558607072.
  * Riemer et al. (2018) Matthew Riemer, Miao Liu, and Gerald Tesauro.  Learning abstract options.  _CoRR_ , abs/1810.11583, 2018.  URL http://arxiv.org/abs/1810.11583.
  * Schaul et al. (2015) Tom Schaul, Dan Horgan, Karol Gregor, and David Silver.  Universal value function approximators.  In _Proceedings of the 32nd International Conference on International Conference on Machine Learning - Volume 37_ , ICML’15, page 1312–1320. JMLR.org, 2015.
  * Schulman et al. (2017) John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov.  Proximal policy optimization algorithms.  _CoRR_ , abs/1707.06347, 2017.  URL http://arxiv.org/abs/1707.06347.
  * Sibson (1969) Robin Sibson.  Information radius.  _Zeitschrift für Wahrscheinlichkeitstheorie und verwandte Gebiete_ , 14(2):149–160, 1969.
  * Smith et al. (2018) Matthew Smith, Herke van Hoof, and Joelle Pineau.  An inference-based policy gradient method for learning options.  In Jennifer Dy and Andreas Krause, editors, _Proceedings of the 35th International Conference on Machine Learning_ , volume 80 of _Proceedings of Machine Learning Research_ , pages 4703–4712. PMLR, 10–15 Jul 2018.  URL http://proceedings.mlr.press/v80/smith18a.html.
  * Sutton and Barto (2018) Richard S. Sutton and Andrew G. Barto.  _Reinforcement Learning: An Introduction_.  The MIT Press, second edition, 2018.  URL http://incompleteideas.net/book/the-book-2nd.html.
  * Sutton et al. (1983) Richard S. Sutton, Satinder Singh, and David Mcallester.  Comparing policy-gradient algorithms.  _IEEE Trans. on Systems, Man, and Cybernetics_ , 1983.
  * Sutton et al. (1999a) Richard S. Sutton, David A. McAllester, Satinder P. Singh, and Yishay Mansour.  Policy gradient methods for reinforcement learning with function approximation.  In _NIPS_ , pages 1057–1063, 1999a.
  * Sutton et al. (1999b) Richard S Sutton, Doina Precup, and Satinder Singh.  Between mdps and semi-mdps: A framework for temporal abstraction in reinforcement learning.  _Artificial intelligence_ , 112(1-2):181–211, 1999b.
  * Tiwari and Thomas (2018) Saket Tiwari and Philip S. Thomas.  Natural option critic.  _CoRR_ , abs/1812.01488, 2018.  URL http://arxiv.org/abs/1812.01488.
  * Todorov et al. (2012) Emanuel Todorov, Tom Erez, and Yuval Tassa.  Mujoco: A physics engine for model-based control.  In _IROS_ , pages 5026–5033. IEEE, 2012.  ISBN 978-1-4673-1737-5.  URL http://dblp.uni-trier.de/db/conf/iros/iros2012.html#TodorovET12.
  * van Seijen et al. (2009) Harm van Seijen, Hado van Hasselt, Shimon Whiteson, and Marco Wiering.  A theoretical and empirical analysis of expected sarsa.  In _2009 IEEE Symposium on Adaptive Dynamic Programming and Reinforcement Learning_ , pages 177–184, 2009.  doi: 10.1109/ADPRL.2009.4927542.
  * 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 _International Conference on Machine Learning_ , pages 3540–3549. PMLR, 2017.
  * Wulfmeier et al. (2020) Markus Wulfmeier, Dushyant Rao, Roland Hafner, Thomas Lampe, Abbas Abdolmaleki, Tim Hertweck, Michael Neunert, Dhruva Tirumala, Noah Y. Siegel, Nicolas Heess, and Martin A. Riedmiller.  Data-efficient hindsight off-policy option learning.  _CoRR_ , abs/2007.15588, 2020.  URL https://arxiv.org/abs/2007.15588.
  * Xie et al. (2019) Tengyang Xie, Yifei Ma, and Yu-Xiang Wang.  Optimal off-policy evaluation for reinforcement learning with marginalized importance sampling.  _CoRR_ , abs/1906.03393, 2019.  URL http://arxiv.org/abs/1906.03393.
  * Zhang et al. (2020) Ruiyi Zhang, Bo Dai, Lihong Li, and Dale Schuurmans.  Gendice: Generalized offline estimation of stationary values, 2020.
  * Zhang and Whiteson (2019) Shangtong Zhang and Shimon Whiteson.  DAC: the double actor-critic architecture for learning options.  _CoRR_ , abs/1904.12691, 2019.  URL http://arxiv.org/abs/1904.12691.

## Checklist

  1. 1.

For all authors…

    1. (a)

Do the main claims made in the abstract and introduction accurately reflect
the paper’s contributions and scope? [Yes]

    2. (b)

Did you describe the limitations of your work? [Yes]

    3. (c)

Did you discuss any potential negative societal impacts of your work? [Yes]

    4. (d)

Have you read the ethics review guidelines and ensured that your paper
conforms to them? [Yes]

  2. 2.

If you are including theoretical results…

    1. (a)

Did you state the full set of assumptions of all theoretical results? [Yes]

    2. (b)

Did you include complete proofs of all theoretical results? [Yes]

  3. 3.

If you ran experiments…

    1. (a)

Did you include the code, data, and instructions needed to reproduce the main
experimental results (either in the supplemental material or as a URL)? [Yes]

    2. (b)

Did you specify all the training details (e.g., data splits, hyperparameters,
how they were chosen)? [Yes]

    3. (c)

Did you report error bars (e.g., with respect to the random seed after running
experiments multiple times)? [Yes]

    4. (d)

Did you include the total amount of compute and the type of resources used
(e.g., type of GPUs, internal cluster, or cloud provider)? [Yes]

  4. 4.

If you are using existing assets (e.g., code, data, models) or
curating/releasing new assets…

    1. (a)

If your work uses existing assets, did you cite the creators? [No]

    2. (b)

Did you mention the license of the assets? [N/A]

    3. (c)

Did you include any new assets either in the supplemental material or as a
URL? [No]

    4. (d)

Did you discuss whether and how consent was obtained from people whose data
you’re using/curating? [N/A]

    5. (e)

Did you discuss whether the data you are using/curating contains personally
identifiable information or offensive content? [N/A]

  5. 5.

If you used crowdsourcing or conducted research with human subjects…

    1. (a)

Did you include the full text of instructions given to participants and
screenshots, if applicable? [N/A]

    2. (b)

Did you describe any potential participant risks, with links to Institutional
Review Board (IRB) approvals, if applicable? [N/A]

    3. (c)

Did you include the estimated hourly wage paid to participants and the total
amount spent on participant compensation? [N/A]

##  Appendix A Appendix

###  A.1 Algorithm

The code is available at https://github.com/mklissa/MOC.

Algorithm 1 Multi-updates Option Critic

Set $s\longleftarrow s_{0}$

Choose $o$ at $s$ according to $\mu_{z}(.|s)$

repeat

Choose $a$ according to $\pi_{\zeta}(a|s,o)$

Take action $a$ in $s$, observe $s^{\prime},r$

Sample termination from $\beta_{\nu}(s^{\prime},o)$

if $o$ terminates in $s^{\prime}$ then

Sample $o^{\prime}$ according to $\mu_{z}(.|s)$

else

$o^{\prime}$ = $o$

end if

Define previous option $\bar{o}=o$ 1\. Evaluation step:for each option
$\tilde{o}$ in the option set $\mathcal{O}$ do

$\delta\leftarrow\,\mathbb{E}[U^{\rho}|s,\tilde{o}]-Q_{\theta}(s,\tilde{o})$
where  $U^{\rho}$ is an importance sampling weighted target

$\theta\leftarrow\,\theta+p_{\mu,\beta}(\tilde{o}|s,\bar{o})\alpha\delta\phi(s,\tilde{o})$end
for 2\. Improvement step for each option $\tilde{o}$ in the option set
$\mathcal{O}$ do

$\zeta\leftarrow\zeta+p_{\mu,\beta}(\tilde{o}|s,\bar{o})\alpha_{\zeta}\frac{\partial\log\pi_{\zeta}(a|s,o)}{\partial\zeta}Q_{\theta}(s,o,a)$end
for

$\nu\leftarrow\nu-\textbf{}\alpha_{\nu}\frac{\partial\beta_{\nu}(s^{\prime},o)}{\partial\nu}(Q_{\theta}(s^{\prime},o)-V_{\theta}(s^{\prime}))$

$z\leftarrow
z+\alpha_{z}\beta_{\nu}(s^{\prime},o)\frac{\partial\mu_{z}(o^{\prime}|s^{\prime})}{\partial
z}Q_{\theta}(s^{\prime},o^{\prime})$ $s\leftarrow s^{\prime}$

until $s^{\prime}$ is a terminal state

###  A.2 Tabular experiments

####  A.2.1 Implementation Details

For our experiments of the FourRooms domain we based our implementation on
[Bacon et al., 2016] and ran the experiments for 500 episodes that last a
maximum of 1000 steps with goal located in the right hallway.

In the first experiment we verify whether learning a fixed set of options can
be accelerated by our method. We define this fixed set as the hallway options
from Sutton et al. [1999b]. As the policies of these options were
deterministic and we use importance sampling, we relax them to stochastic
policies where the most likely action is the one leading to a hallway.

In the second experiment, after 500 episodes where the goal is located in the
right hallway, we move the goal a random location in the bottom left room and
augment the stochasticity of the environment. Both of these modifications are
done to verify the agent’s ability to adapt to change.

For all experiments and all algorithms we performed a learning rate sweep in
the rage of $\\{2\cdot 10^{i},4\cdot 10^{i},6\cdot 10^{i},8\cdot
10^{i}:i=-1,-2,-3\\}$. We used the same learning rate for both intra-option
updates as well as the udpates on the value function. We did so to avoid
introducing a larger hyperparameters search for the hierarchical approaches.
For the MOC hyperparameter $\eta$, we looked at values in
$\\{0.2,0.3,0.5,0.7,1.0\\}$. All values are available in Table 2

Hyperparameter | Value
---|---
AC lr | 2e-1
OC lr | 8e-1
MOC lr | 8e-1
$\eta$ | 0.3
Table 2: Best hyperparameters for the FourRooms domain.

####  A.2.2 Option visualization

In the following figure we present a qualitative representation of the options
being learned by our method MOC and compare it to the OC method.

(a) Option-Critic

(b) Multi-updates Option-Critic

Figure 6: Visualization of the learned options. For each option, the
brightness of the state represent the probability of selecting the said
option. For each state and option, we plot by a arrow the most likely action.
In a) we present two options learned by the OC algorithm. We notice a certain
imbalance between options in the sense that option 1 is almost always taken by
the policy over options. In b) we present two options learned by our approach,
MOC, and notice a better balance between options. Moreover, even if an option
is not likely to be selected, it generally chooses relevant actions, which is
not the case for the OC agent.

###  A.3 Linear Function Approximation

####  A.3.1 Implementation Details

The environment we consider is a sparse variant of the classic Mountain Car
domain [Moore, 1991, Sutton and Barto, 2018] where a positive reward is given
upon reaching the goal. We consider the transfer setting where after half the
total number of timesteps the gravity is doubled.

In this experiment, the critic is a linear function of the features, while the
actor is a two layer network with a tanh non-linearity in the hidden layer.
The features are defined through four radial basis functions with radii
$\\{5,2,1,0.5\\}$. These RBF kernels are fit by sampling random states from
the environment for 100000 steps before the training starts.

For all algorithms we performed a learning rate sweep in the set
$\\{0.0008,0.001,0.003,0.004,0.005,0.006,0.008,0.01\\}$. As in the tabular
case, we used the same learning rate for both intra-option policy updates, the
option value function udpates, the termination functions updates as well as
the udpates on to the policy over options. We did so to avoid introducing a
larger hyperparameters search for the hierarchical approaches. Other
hyperparameters values are standard values for the A2C algorithm. For the MOC
hyperparameter $\eta$, we did a hyperparameter study for values in
$\\{0.1,0.2,0.3,0.4,0.5,0.7,0.8,0.9,1.0\\}$ in order to verify how the
algorithm behaves with respect to the number of updates done to all options.
The results are plotted in Figure 2(b) and we see that higher values of $\eta$
lead to generally better final performance. All values are available in Table
3

Hyperparameter | Value
---|---
AC lr | 0.005
OC lr | 0.004
MOC lr | 0.004
$\eta$ | 0.9
# steps before update | 5
actor hidden features size | 128
value loss coefficient | 0.5
# RBF kernels | 5
Table 3: Best hyperparameters for the MountainCar-v0 domain.

###  A.4 Continuous Control

####  A.4.1 Tasks Description

We build on classic continuous control domains [Brockman et al., 2016] and
modify them to make the learning process more challenging. The first
experiment is HalfCheetahDir-v1, visualized in Figure 7(a) , in which the
cheetah first learns to walk in the forward direction and then needs to adapt
to a change in reward signal which now encourages walking backward. The second
experiment is AntWalls-v1 illustrated in Figure 7(b) where the agent first
learns how to walk forward as in the classic Ant-v0 domain [Brockman et al.,
2016]. In the transfer task, we add walls as obstacles and provide a terrain-
read which the agent can use to circumvent them. The final experiment is
Walker2dStand-v1 shown in 7(c) in which the walker agent first learns how to
stand up and in the transfer task needs to learn how to walk.

(a)

(b)

(c)

Figure 7: MuJoCo tasks visual representation.  Table 4: MuJoCo experiments hyperparameters Hyperparameter | AntWalls-v1 | Walker2dStand-v1 | HalfCheetahDir-v1
---|---|---|---
PPO Learning rate | 2e-4 | 2e-4 | 3e-4
IOC LR (policy, termination, master) | 1e-4 | 1e-4 | 1e-4,5e-4,5e-4
MOC LR (policy, termination, master) | 8e-5 | 9e-5 | 3e-4,5e-4,5e-4
IOPG LR (policy, termination, master) | 9e-5 | 8e-5 | 1e-4,3e-4,5e-4
$\gamma$ | 0.99 | 0.99 | 0.99
$\lambda$ | 0.95 | 0.95 | 0.95
Entropy coefficient | 0.0 | 0.0 | 0.0
LR schedule | constant | constant | constant
PPO steps | 2048 | 2048 | 2048
PPO cliping value | 0.1 | 0.1 | 0.1
# of minibatches | 32 | 32 | 32
# of processes | 1 | 1 | 1
$\eta$ | 0.9 | 0.9 | 0.9

####  A.4.2 Implementation Details

We based our implementation on [Dhariwal et al., 2017, Klissarov et al., 2017,
Khetarpal et al., 2020] and we ran the experiments for 2M steps for
AntWalls-v1 and Walker2dStand-v1, and for 1M steps for HalfCheetahDir-v1.
After about one quarter of the total timesteps we change the original task to
a transfer task that is meant to be more challenging.

For the Walker2dStand-v1 and AntWalls-v1 experiments, we performed a
hyperparameter search for all algorithms on the learning rate in the set
$\\{8\cdot 10^{-5},9\cdot 10^{-5},2\cdot 10^{-4},5\cdot 10^{-4},7\cdot
10^{-4}\\}$. As in our previous experiments, we used the same learning for all
options components. For the HalfCheetahDir-v1 experiment, we performed a
hyperparameter search for all algorithms on the learning rate in the set
$\\{1\cdot 10^{-4},3\cdot 10^{-4},5\cdot 10^{-4},7\cdot 10^{-4}\\}$. For this
experiment we tried different values of the learning rates for each option
component in order to stay consistent with the hyperparameter search done in
previous work [Khetarpal et al., 2020]. For the MOC $\eta$ hyperparameter, we
simply used $0.9$ as we have previously shown that higher values are usually
better.

The architecture for the algorithms was kept to be the same with respect to
the original code base. All other PPO hyperparameters were kept as is. All the
values are available in Table 4.

####  A.4.3 Option visualization

We share qualitative results of the kind of options our method, MOC, learns as
well as the ones learned by the IOC approach. These results are available in
the form of videos in the following URL.

###  A.5 MiniWorld Navigation

We build on the first-person visual navigation domain of MiniWorld [Chevalier-
Boisvert, 2018] to propose a transfer learning setting. In the transfer task
we aim to increase the difficulty of the source. Moreover, we look at a
setting where the agents could benefit from reusing the options learned in the
source task. The tasks are defined as in the following.

We experiment with the MiniWorld-YMazeTransfer-v0 shown in Figure 8(a) domain
which contains of a maze in a Y shape. The agent spawns in a random position
in the bottom part and has to reach the goal represented by a red box. At
first this goal is in a random location on the left side and in the transfer
task the goal is moved to the right side. We also explore the MiniWorld-
WallGapTransfer-v0 domain shown in 8(b), where a gap in the wall separates two
large rooms. At first, the agent has to navigate to the adjacent room, scan
the environment and collect the goal represented by a red box. In the transfer
task, we introduce a blue box, which now is the new goal and is located in the
room adjacent to where the agent spawns. We also include a red box in the same
room as the agent, however this box no longer gives the agent any reward.

For both environment, the agent and the goal locations are randomized at each
iteration in order to avoid deterministic solutions.

(a)

(b)

Figure 8: MiniWorld tasks visual representation.  Hyperparameter | Value
---|---
Learing rate | 2.5e-4
$\gamma$ | 0.99
$\lambda$ | 0.95
Entropy coefficient | 0.01
LR schedule | constant
PPO steps | 128
PPO cliping value | 0.1
# of minibatches | 4
# of processes | 32
$\eta$ | 0.7
Table 5: Hyperparameters for the MiniWorld experiments

####  A.5.1 Implementation Details

We based our implementation on [Kostrikov, 2018] and we ran the experiments
for 2M steps for the MiniWorld-YMazeTransfer-v0 experiment and 5M steps for
the MiniWorld-WallGapTransfer-v0 experiment. After 1M steps and 2M
respectively, we change the original task to a transfer task that is meant to
be more challenging.

The architecture for the algorithms was kept to be the same with respect to
the original code base. Similarly, all PPO hyperparameters, including the
learning rate, were kept as is. For the MOC algorithm we used $0.7$ as a value
for the hyperparameter $\eta$ after looking for values in $\\{0.7,1.0\\}$.
Best hyperparameters values are found in Table 5.

####  A.5.2 Option visualization

Once again we share results on the options learned by our MOC algorithm and
the OC algorithm. These results are available in the form of videos in the
following URL.

###  A.6 Proof of Lemma 1

See 1

###### Proof.

The limit in Equation 2 exists if and only if the Markov chain on the
augmented state-option space is ergodic. To ensure ergodicity, we first need
to verify irreducibility, that is, whether it is possible to get from any
state-option pair $(s,o)$ to any other pair $(s^{\prime},o^{\prime})$ in a
finite amount of time. Concretely, there exist a positive integer $n$ such
that,

$[(P_{\pi,\mu,\beta})^{n}]_{(s,o),(s^{\prime},o^{\prime})}>0$

As we consider a finite augmented space, positive recurrrence is implied by
irreducibility. Secondly, we need to verify that the Markov chain on the
augmented space also satisfies aperiodicity.

To verify both criteria, we will use the following expansion on the state-
option transition function,

$\displaystyle P_{\pi,\mu,\beta}(s^{\prime},o^{\prime}|s,o)$
$\displaystyle=p_{\mu,\beta}(o^{\prime}|s^{\prime},o)\sum_{a}P(s^{\prime}|s,a)\pi(a|s,o)$
(8)
$\displaystyle=\big{(}\beta(s^{\prime},o)\mu(o^{\prime}|s^{\prime})+(1-\beta(s^{\prime},o))\delta_{o,o^{\prime}}\big{)}P^{\pi_{o}}(s^{\prime}|s)$
(9)

For each option, Assumption 1 verifies that $P^{\pi_{o}}(s^{\prime}|s)$ is
ergodic. Moreover, Assumption 2 implies that both the termination functions
$\\{\beta\\}$ and the policy over options $\mu$ have strictly positive values.
As per the definition in Equation 9, both assumptions lead to satisfying the
irreducibility and aperiodicity criteria and therefore the chain on the
augmented space is ergodic. We can then write,

$\displaystyle\lim_{t\rightarrow\infty}P_{\pi,\mu,\beta}(S_{t}=s,O_{t}=o)$
$\displaystyle=\sum_{\bar{o}}\lim_{t\rightarrow\infty}P_{\pi,\mu,\beta}(S_{t}=s,O_{t-1}=\bar{o})$
$\displaystyle\big{(}(1-\beta(s,\bar{o}))\delta_{o=\bar{o}}+\beta(s,\bar{o})\mu(o|s)\big{)}$
$\displaystyle=\sum_{\bar{o}}\bar{d}_{\pi,\mu,\beta}(\bar{o},s)\big{(}(1-\beta(s,\bar{o})\delta_{o=\bar{o}}+\beta(s,\bar{o})\mu(o|s)\big{)}$

∎


# Benchmarking deep generative models for diverse antibody sequence design

Igor Melnyk
IBM Research AI
<EMAIL_ADDRESS>
Payel Das
IBM Research AI
<EMAIL_ADDRESS>
Vijil Chenthamarakshan
IBM Research AI
<EMAIL_ADDRESS>
Aurelie Lozano
IBM Research AI
<EMAIL_ADDRESS>

###### Abstract

Computational protein design, i.e. inferring novel and diverse protein
sequences consistent with a given structure, remains a major unsolved
challenge. Recently, deep generative models that learn from sequences alone or
from sequences and structures jointly have shown impressive performance on
this task. However, those models appear limited in terms of modeling
structural constraints, capturing enough sequence diversity, or both. Here we
consider three recently proposed deep generative frameworks for protein
design: (AR) the sequence-based autoregressive generative model, (GVP) the
precise structure-based graph neural network, and Fold2Seq that leverages a
fuzzy and scale-free representation of a three-dimensional fold, while
enforcing structure-to-sequence (and vice versa) consistency. We benchmark
these models on the task of computational design of antibody sequences, which
demand designing sequences with high diversity for functional implication. The
Fold2Seq framework outperforms the two other baselines in terms of diversity
of the designed sequences, while maintaining the typical fold.

##  1 Introduction

Antibodies and their functional domains play key roles in research,
diagnostics, and therapeutics. Among them, an attractive class is comprised of
nanobodies, which are functional antibody domains with small size ($\sim$15
kDa) and high stability (Tm up to 90 C), and therefore are of increasing
therapeutic interest [3]. Designing functional sequences typically requires a
combinatorial exploration of sequence space. To address this issue, one can
impose sequence and/or structural constraints to narrow the search space.

Functionally diverse antibodies correspond to the highly stable immunoglobulin
fold, a universal protein scaffold structure. Antigen binding specificity is
largely determined by the sequence and structural diversity of the
complementarity-determining regions (CDRs) that are displayed on canonical
frameworks. Among the CDRs, CDR3 contributes most of the sequence and length
diversity. As a result, a big focus in computational antibody design has been
on sampling diverse CDR3s.

Along this direction, recently deep generative models have been used for
sampling virtual sequences that were then prioritized for experimental
validation. Most of those models are sequence-based that leveraged
autoregressive neural nets such as LSTM. Such models have been successfully
used for the entire functional antibody sequence design [13], as well as for
designing only the CDRs given rest of the sequence as a context [14].
Nevertheless, such model does not explicitly account for the 3D structural
constraints associated with the characteristic immunoglobulin fold and may be
limited in terms of capturing the broad landscape of CDRs.

In this paper, we benchmark three recently proposed state-of-the-art deep
generative models – sequence-based, 3D structure-based, and 3D fold-based – by
comparing the generated sequences of llama nanobodies. The structure-based and
fold-based models generate full protein sequences by using a single
representative 3D structure of the llama nanobody as an input, whereas the
autoregressive sequence-based model designs CDR3 by conditioning on the
preceding germline framework-CDR1-CDR2 nanobody sequence. The structure-based
model leverages the precise backbone coordinates, whereas the fold-based model
uses a fuzzy description of the secondary structures elements in 3D as the
input. Machine learning, physicochemical, bioinformatics, and structural
metrics of the generated sequences from a single llama nanobody
sequence/structure and their comparison with the natural llama nanobody
repertoire suggest that generated sequences strongly vary in term of novelty,
diversity, consistency, and coverage.

Figure 1: Overview of our study for benchmarking models for antibody sequence
design. Starting with the prefix sequence (CDR1, CDR2 and the framework),
structure or the fold of the PDB ID 3K3Q, we use three models i.e., AR [14],
GVP [9] and Fold2Seq [4], respectively to generate the protein’s CDR regions.
For GVP and Fold2Seq, whose output is the entire protein sequence, we also run
ANARCI to extract the CDRs. The output from all the models is then processed
and filtered for the final property analysis.

##  2 Related Work

The protein/antibody design problem has been recently addressed using multiple
machine learning approaches. For example, the sequence-based methods of [2,
13, 1] formulated the problem as an autoregressive process, leveraging deep
recurrent networks such as an LSTM, to design the entire antibody sequence. On
the other hand, the work of [14] focused on the variable complementarity
determining regions (CDRs), which are complex and highly diverse parts of the
sequence, determining the specificity of the antibody. Instead of recurrent
networks, they applied causal dilated convolutions to model the autoregressive
likelihood.

Alternatively, ML-physics hybrid methods have also been explored. For
instance, [15], proposed a system, that starts from a random protein sequence,
and the search process is guided by the physics-based constraints to satisfy
the desired structural motifs. EVOdesign [11], on the other hand, is a
structure-based approach that uses the evolutionary profiles to guide the
sequence search. It first identifies the structural analogs from PDB and then
constructs a structural sequence profile using MSA. The physics-based force
fields are then used to search for the low free-energy sequence states.

Recently, a class of deep generative models that account for the 3D structural
constraints, have been proposed. For example, [7] used a generative model for
protein sequences design given a target structure, represented as a graph over
the residues. The key challenge was to account for long-range dependencies in
the protein sequence that are usually short-range in the 3D structure space.

Based on [7] the work of [9] then proposed Geometric Vector Perceptrons (GVPs)
to allow for the embedding of geometric information at nodes and edges without
reducing such information to scalars that may not fully capture complex
geometry. The proposed graph network ensures that the vector and scalar
outputs are equivariant and invariant with respect to rotations and
reflections.

Designing proteins based on a rigid backbone structure is known to restrict
the diversity and novelty of the sequences. For this reason, [4] proposed
Fold2Seq, a transformer-based generative framework for designing protein
sequences conditioned on a generic target fold, rather than the specific high
resolution 3D structure. The fold here was defined as the spatial arrangement
of the local secondary structure elements. The method also uses joint
sequence–fold embedding to better capture the relationship between the two
modalities.

Finally, a recent work of [8] proposed a generative model to specifically
design only the CDRs of antibodies rather than the whole protein sequence. The
approach is based on co-designing the CDR sequence and 3D structure of CDR as
graphs. The output is built autoregressively while iteratively refining its
predicted global structure, which in turn guides the next residue choice.

Figure 2: Kernel density plot as a function of pairwise sequence similarity
within the same ensemble (f2s= Fold2Seq, gvp=GVP, AR=AR, nat=Natural) and with
the natural ensemble (*-nat). Higher score indicates better similarity.

##  3 Method

Figure 1 shows the overview of our benchmarking study.

Models We examine three different models. AR [14]: the autoregressive approach
that uses the causal dilated convolutions for the input prefix sequence to
generate the CDR3 protein subsequence. GVP [9]: the encode-decoder GNN model
that uses Geometric Vector Perceptrons to encode the scalar and vector
structure information which is then decoded in the autoregressive manner to
generate the entire protein sequence. Fold2seq [4]: another encode-decoder
model based on transformer architecture that embeds the fuzzy 3D protein
structure information (fold) in the joint sequence-fold embedding space and
then decodes it autoregressively into the corresponding protein sequence. For
the GVP and Fold2Seq outputs we run ANARCI [6] to extract the CDR regions.

Sequence Design The sequence and structure corresponding to the Chain A of pdb
id 3K3Q were used as inputs for this study. It is worth noting that this
structure is included in the training of both GVP and Fold2Seq model, whereas
a maximum of 58.94% sequence identity was found to be present between the
input sequence and the AR training set.

We compare the full sequences as well as the CDRs across the generated
ensembles. For this purpose, we extracted the CDRs from the generated
sequences using the IMGT numbering scheme as returned by ANARCI software [6].
For the extracted CDRs, we estimate the percentage of unique sequences
(uniqueness). Sequences that contain glycosylation sites, asparagine
deamination motifs, or sulfur-containing amino acids (cysteine and methionine)
were removed. For the AR model, we optionally considered an extra filter to
exclude sequences that do not end with the final beta-strand of the nanobody
template as in [14]. We denote the approach with final beta-strand filtering
by _AR filtered_ , while we call _AR unfiltered_ the version without final
beta-strand filtering.

Evaluation Metrics We define the set of the generated sequences (structures)
conditioned on sequence/structure $j$ as $\mathcal{G}_{j}$. In structure-based
design, Sequence Recovery rate is defined as ($SRR$) for $\bm{y}_{j}$ as
$SRR_{structure}(j)=\frac{1}{|\mathcal{G}_{j}|}\sum_{g\in\mathcal{G}_{j}}SIM(\bm{x}_{g},\bm{x}_{j})$.
A global alignment scheme and BLAST62 matrix, with a gap opening penalty of
-10 and gap extending penalty of -1, were used for estimating pairwise
sequence identity (SIM) and alignment score. Negative log likelihood (NLL) was
estimated by using the autoregressive (AR) generative model trained on 1.2
million natural llama nanobody sequences [14], as following:
$NLL=-\sum_{k=1}^{K}{\log(p(x_{k}|X_{<k}))}$, which is sum of the cross-
entropy between the true residue at each position and the predicted
distribution over possible residues, conditioned on the preceding characters.
Structural recovery of the three sequence design models by predicting the 3D
structure of the top 100 generated sequences using pretrained models from
Alphafold2111https://github.com/kalininalab/alphafold_non_docker.

##  4 Results

Table 1: Sequence recovery rate (SRR) and NLL from the autoregressive sequence model [14] trained on natural llama nanobody repertoire for the sequences generated by the comparison methods and the sequences from the natural llama library, synthetic library, and next-generation sequencing library. Model |  Seq Recovery Rate (%) |  NLL
---|---|---
Fold2Seq |  30.711 |  2.572
GVP |  40.131 |  2.987
AR |  48.865 |  0.375
Natural |  – |  0.371
Synthetic |  – |  4.912
NGS |  – |  5.102
Table 2: Uniqueness and novelty of the CDR3, CDR2, and CDR1 regions of the sequences generated by Fold2Seq, GVP, AR without final beta-stand filtering (AR unfiltered), AR with final beta-strand filtering (AR filtered), and the sequences from the natural llama library. Note that the AR approach has been trained to generate CDR3 only, given the preceding portion of a ground truth sequence. |  | Fold2Seq | GVP | AR unfiltered | AR filtered | Natural Llama
---|---|---|---|---|---|---
CDR3 | Uniqueness | 100 | 88.33 | 87.57 | 13.85 | 100
Novelty | 43.36 | 52.71 | 11.92 | 8.97 | 52.64
CDR2 | Uniqueness | 100 | 9.15 | – | – | 100
Novelty | 58.70 | 9.15 | – | – | 83.83
CDR1 | Uniqueness | 92.49 | 56.20 | – | – | 100
Novelty | 60.75 | 51.99 | – | – | 83.37
Figure 3: 2D kernel density plot as a function of isoelectric point and length
of generated/natural CDR3 sequences. Black dot indicates the groundtruth CDR3.

Figure 4: (a) Scatter plot of sequence identity and structure similarity (TM-
score) for Fold2Seq, GVP and AR generated sequences. The structure of the
generated sequences are estimated using AlphaFold and the TM-Score is
calculated with the ground truth structure using TM-Align. Note that since the
sequence identity is computed over the entire protein chain, the AR method,
generating only the CDR region (as opposed to others, that generate the whole
sequence), by construction has higher similarity to the ground truth. (b)
Alignment of the structure corresponding to a Fold2seq generated sequence
(green) with the ground truth structure for PDB Id 3K3Q chain A using TM-
align. Arrow indicates CDR3. TM-score is 0.75 and backbone RMSD is 1.86 Å.

Table 1 reports the average sequence recovery rate and average negative log
likelihood from the trained autoregressive model in [14], estimated using 10k
generated sequences. The autoregressive sequence model provides highest
sequence recovery rate, followed by GVP and Fold2Seq. Both methods yield
sequences that share $\geq$ 30% identity on average with the sequence of the
input structure, implying fold-consistency of the generated sequences, as a
30% sequence identity threshold typically suggests fold homology [12]. SRR of
GVP is higher than Fold2Seq, consistent with earlier obeservation that a
structure-based model yields higher sequence recovery than a fold-based model
[4]. The negative log likelihood (NLL) from the trained autoregressive model
was found to be consistent with the experimentally reported thermostability of
unseen llama nanobody sequences, which is an important aspect of nanobody
fitness [14]. Therefore, we also estimate the NLL of the generated sequences
using Fold2Seq, GVP, AR, as well as a state-of-art synthetic library
(Synthetic, constructed combinatorically using the position-specific amino
acid frequencies of nanobody sequences with crystal structures in the PDB
database) [10] and a large collection of nanobody sequences from next-
generation sequencing repositories (ngs) [5]. As expected, AR-generated
sequences show a NLL close to the natural nanobodies that were used for
training of the model. Both synthetic and ngs libraries show very high NLL,
indicating the bias of the trained AR toward natural llama nanobody
repertoire. Among the two 3D-conditional generative models, interestingly, F2S
sequences show lower NLL than GVP, implying F2S sequences are closer to the
broader llama nanobody sequences.

As reported in Table 2, we estimate uniqueness and novelty of all three
generative models with respect to the training set used for the AR model
training. Fold2Seq outperforms both GVP and AR by a significant margin in term
of uniqueness for all CDRs. Trend in novelty is not as clear though between
GVP and Fold2Seq: while GVP produces more novel CDR3s, Fold2Seq performs
better for CDR1s and CDR2s in term of novelty. For reference we also include
in Table 2 uniqueness and novelty for 10,000 sequences selected at random from
the natural llama library of [14]. This result implies that a fuzzy
representation of a 3D fold allows generating sequences with novel and unique
CDRs, when compared to a sequence-based and a backbone structure-based model.

This result is confirmed as we estimate the pairwise sequence similarity
within the generated ensemble from a particular model (Figure 2). When
compared to the database comprised of natural llama nanobody sequences,
Fold2Seq sequences are more diverse (low similarity score within the ensemble)
than GVP and AR sequences. GVP sequences show an interesting characteristic –
majority of sequences are highly similar - indicating a mode collapse. In term
of similarity with the natural CDR3s, Fold2Seq sequences are more distant,
followed by GVP and AR. The similarity distribution corresponding to GVP seem
narrower, consistent with lack of diversity in generated samples.

Further, we analyze the physicochemical properties, such as isoelectric point
(pI) and length of the CDR3 sequences. Figure 3 shows that the natural
nanobody CDR3s populate broad ranges of isoelectric point and sequence length,
with a preference toward pI around 6 and length around 14 and 18. Fold2Seq
produces a significant coverage of the natural sequences, with a clustering
around pI = 6 and length = 14, while still exhibiting non-zero density around
the input. GVP produces sequences that are very close to the one corresponding
to the input pdb only, while AR tends to generate short (around 6 amino acid)
sequences with pI around 4.

Finally, we compare the structural recovery of the three generative frameworks
by predicting the 3D structure of 100 random generated sequences using
Alphafold2. Figure 4(a) shows that, almost all of the sequences generated by
the Fold2Seq and GVP share a sequence identity of $\geq$ 30% and TM-score of
$\geq$ 0.5 with the groundtruth, indicating both methods perform similarly in
term of recovering correct fold topology [16]. However, GVP sequences exhibit
higher TM-score than Fold2Seq ones, consistent with high consistency of the
generated sequences with the input while lacking coverage and diversity (see
Figure 4(a)). The AR sequences show high sequences identity as well as high
TM-score, which is not surprising given that the model “grafts” a short
(around 6 residue) sequences within the input sequence. In contrast, though
Fold2Seq generates diverse and dissimilar sequences, corresponding structures
turn out to be high-consistency, with a 1.8-3 Å backbone RMSD from the input
structure (see Figure 4(b)).

In conclusion, this benchmarking study highlights key performance differences
of three protein sequence design models on the nanobody sequence generation
task. The backbone based GVP model provides high consistency with the input,
while lacking the diversity present within the llama nanobody sequences. A
sequence-based model, on the other hand, captures a limited spectrum of the
llama nanobody repertoire. The sequence design model that employs a fuzzy
representation of the corresponding 3D fold as the input and generates
sequences by leveraging representation learned by join training on 3D folds
and 1D sequences provides better diversity, while still providing sufficient
consistency with the input sequence/structure and coverage of the llama
nanobody repertoire.

Ultimately, the choice of method should be guided by the task at hand. If one
wishes to generate sequences very close to an input sequence/structure, GVP
might be more appropriate. If one wishes to graft specific regions such as a
short CDR3 within an input sequence, then the AR approach would be better
suited. If one wishes to generate diverse sequences with broader coverage that
“extrapolate” to a greater extent from a given input sequence, while still
maintaining consistency with the input, then Fold2Seq would be a method of
choice.

## References

  * [1] Rahmad Akbar, Philippe A Robert, Cédric R Weber, Michael Widrich, Robert Frank, Milena Pavlović, Lonneke Scheffer, Maria Chernigovskaya, Igor Snapkov, and Andrei Slabodkin.  In silico proof of principle of machine learning-based antibody design at unconstrained scale.  BioRXiV, 2021.
  * [2] Ethan C Alley, Grigory Khimulya, Surojit Biswas, Mohammed AlQuraishi, and George M Church.  Unified rational protein engineering with sequence-based deep representation learning.  Nature methods, 16(12):1315–1322, 2019.
  * [3] Peter Bannas, Julia Hambach, and Friedrich Koch-Nolte.  Nanobodies and nanobody-based human heavy chain antibodies as antitumor therapeutics.  Frontiers in immunology, 8:1603, 2017.
  * [4] Yue Cao, Payel Das, Vijil Chenthamarakshan, Pin-Yu Chen, Igor Melnyk, and Yang Shen.  Fold2seq: A joint sequence (1d)-fold (3d) embedding-based generative model for protein design.  In International Conference on Machine Learning, pages 1261–1271. PMLR, 2021.
  * [5] Piotr Deszynski, Jakub Mlokosiewicz, Adam Volanakis, Igor Jaszczyszyn, Natalie Castellana, Stefano Bonissone, Rajkumar Ganesan, and Konrad Krawczyk.  Indi-integrated nanobody database for immunoinformatics.  medRxiv, 2021.
  * [6] James Dunbar and Charlotte M Deane.  Anarci: antigen receptor numbering and receptor classification.  Bioinformatics, 32(2):298–300, 2016.
  * [7] John Ingraham, Vikas Garg, Regina Barzilay, and Tommi Jaakkola.  Generative models for graph-based protein design.  In Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
  * [8] Wengong Jin, Jeremy Wohlwend, Regina Barzilay, and Tommi Jaakkola.  Iterative refinement graph neural network for antibody sequence-structure co-design, 2021.
  * [9] Bowen Jing, Stephan Eismann, Patricia Suriana, Raphael JL Townshend, and Ron Dror.  Learning from protein structure with geometric vector perceptrons.  arXiv preprint arXiv:2009.01411, 2020.
  * [10] Conor McMahon, Alexander S Baier, Roberta Pascolutti, Marcin Wegrecki, Sanduo Zheng, Janice X Ong, Sarah C Erlandson, Daniel Hilger, Søren GF Rasmussen, and Aaron M Ring.  Yeast surface display platform for rapid discovery of conformationally selective nanobodies.  Nature structural & molecular biology, 25(3):289–296, 2018.
  * [11] Robin Pearce, Xiaoqiang Huang, Dani Setiawan, and Yang Zhang.  Evodesign: Designing protein-protein binding interactions using evolutionary interface profiles in conjunction with an optimized physical energy function.  Journal of molecular biology, 431 13:2467–2476, 2019.
  * [12] William R Pearson.  An introduction to sequence similarity (“homology”) searching.  Current protocols in bioinformatics, 42(1):3–1, 2013.
  * [13] Koichiro Saka, Taro Kakuzaki, Shoichi Metsugi, Daiki Kashiwagi, Kenji Yoshida, Manabu Wada, Hiroyuki Tsunoda, and Reiji Teramoto.  Antibody design using lstm based deep generative model from phage display library for affinity maturation.  Scientific reports, 11(1):1–13, 2021.
  * [14] Jung-Eun Shin, Adam J Riesselman, Aaron W Kollasch, Conor McMahon, Elana Simon, Chris Sander, Aashish Manglik, Andrew C Kruse, and Debora S Marks.  Protein design and variant prediction using autoregressive generative models.  Nature Communications, 12, 2021.
  * [15] Doug Tischer, Sidney Lisanza, Jue Wang, Runze Dong, Ivan Anishchenko, Lukas F Milles, Sergey Ovchinnikov, and David Baker.  Design of proteins presenting discontinuous functional sites using deep learning.  bioRxiv, 2020.
  * [16] Jinrui Xu and Yang Zhang.  How significant is a protein structure similarity with tm-score= 0.5?  Bioinformatics, 26(7):889–895, 2010.