# Composing Task Knowledge with
Modular Successor Feature Approximators

Wilka Carvalho ,1 Angelos Filos2
Richard L. Lewis1 Honglak lee1,3 Satinder Singh1
1University of Michigan 2University of Oxford 3LG AI Research  Contact author:
<EMAIL_ADDRESS>
###### Abstract

Recently, the Successor Features and Generalized Policy Improvement (SF&GPI)
framework has been proposed as a method for learning, composing, and
transferring predictive knowledge and behavior. SF&GPI works by having an
agent learn predictive representations (SFs) that can be combined for transfer
to new tasks with GPI. However, to be effective this approach requires state
features that are useful to predict, and these state-features are typically
hand-designed. In this work, we present a novel neural network architecture,
“Modular Successor Feature Approximators” (MSFA), where modules both discover
what is useful to predict, and learn their own predictive representations. We
show that MSFA is able to better generalize compared to baseline architectures
for learning SFs and modular architectures for learning state representations.

##  1 Introduction

Consider a household robot that needs to learn tasks including picking up
dirty dishes and cleaning up spills. Now consider that the robot is deployed
and encounters a table with both a spill and a set of dirty dishes. Ideally
this robot can combine its training behaviors to both clean up the spill and
pickup the dirty dishes. We study this aspect of generalization: combining
knowledge from multiple tasks.

Combining knowledge from multiple tasks is challenging because it is not clear
how to synthesize either the behavioral policies or the value functions
learned during training. This challenge is exacerbated when an agent also
needs to generalize to novel appearances and environment configurations.
Returning to our example, our robot might need to additionally generalize to
both novel dirty dishes and to novel arrangements of chairs.

Successor features (SFs) and Generalized Policy Improvement (GPI) provide a
mechanism to combine knowledge from multiple training tasks (Barreto et al.,
2017; 2020). SFs are predictive representations that estimate how much state-
features (known as “cumulants”) will be experienced given a behavior. By
assuming that reward has a linear relationship between cumulants and a task
vector, an agent can efficiently compute how much reward it can expect to
obtain from a given behavior. If the agent knows multiple behaviors, it can
leverage GPI to compute which behavior would provide the most reward (see
Figure 2 for an example). However, SF&GPI commonly assume hand-designed
cumulants and don’t have a mechanism for generalizing to novel environment
configurations.

Modular architectures are a promising method for generalizing to distributions
outside of the training distribution (Goyal et al., 2019; Madan et al., 2021).
Recently, Carvalho et al. (2021a) presented “FARM” and showed that learning
multiple state modules enabled generalization to environments with unseen
environment parameters (e.g. to larger maps with more objects). In this work,
we hypothesize that modules can further be leveraged to discover state-
features that are useful to predict.

Figure 1:  (1) FARM learns multiple state modules. This promotes
generalization to novel environments. However, it has no mechanism for
combining task solutions. (2) USFA learns a single monolithic architecture for
predicing SFs and can combine task solutions. However, it relies on hand-
designed state features and has no mechanism for generalization to novel
environments. (3) We combine the benefits of both. We leverage modules for
reward-driven discovery of state features that are useful to predict. These
form the basis of their own predictive representations (SFs) and enables
combining task solutions in novel environments.

We present “Modular Successor Feature Approximators” (MSFA), a novel neural
network for discovering, composing, and transferring predictive knowledge and
behavior via SF&GPI. MSFA is composed of a set of modules, which each learn
their own state-features and corresponding predictive representations (SFs).
Our core contribution is showing that an inductive bias for modularity can
enable reward-driven discovery of state-features that are useful for zero-shot
transfer with SF&GPI. We exemplify this with a simple state-feature discovery
method presented in Barreto et al. (2018) where the dot-product between state-
features and a task vector is regressed to environment reward. This method
enabled transfer with SF&GPI in a continual learning setting but had limited
success in the zero-shot transfer settings we study. While there are other
methods for state-feature discovery, they add training complexity with mutual
information objectives (Hansen et al., 2019) or meta-gradients (Veeriah et
al., 2019). With MSFA, by adding only an architectural bias for modularity, we
discover state-features that (1) support zero-shot transfer competitive with
hand-designed features, and (2) enable zero-shot transfer in visually diverse,
procedurally generated environments. We are hopeful that our architectural
bias can be leveraged with other discovery methods in future work.

##  2 Related Work on Generalization in RL

Hierarchical RL (HRL) is one dominant approach for combining task knowledge.
The basic idea is that one can sequentially combine policies in time by having
a “meta-policy” that sequentially activates “low-level” policies for
protracted periods of time. By leveraging hand-designed or pre-trained low-
level policies, one can generalize to longer instructions (Oh et al., 2017;
Corona et al., 2020), to new instruction orders (Brooks et al., 2021), and to
novel subtask graphs (Sohn et al., 2020; 2022). We differ in that we focus on
combining policies concurrently in time as opposed to sequentially in time. To
do so, we develop a modular neural network for the SF&GPI framework.

SFs are predictive representations that represent the current state as a
summary of the successive features to follow (see §3 for a formal definition).
By combining them with Generalized Policy Improvement, researchers have shown
that they can transfer behaviors across object navigation tasks (Borsa et al.,
2019; Zhang et al., 2017; Zhu et al., 2017), across continuous control tasks
(Hunt et al., 2019), and within an HRL framework (Barreto et al., 2019).
However, these works tend to require hand-designed cumulants which are
cumbersome to design for every new environment. In our work, we integrate SFs
with Modular RL to facilitate reward-driven discovery of cumulants and improve
successor feature learning.

Modular RL (MRL) (Russell & Zimdars, 2003) is a framework for generalization
by combining value functions. Early work dates back to (Singh, 1992), who had
a mixture-of-experts system select between separately trained value functions.
Since then, MRL has been applied to generalize across robotic morphologies
(Huang et al., 2020), to novel task-robot combinations (Devin et al., 2017;
Haarnoja et al., 2018), and to novel language instructions (Logeswaran et al.,
2021). MSFA, is the first to integrate MRL with SF&GPI. This integration
enables combining task solutions in novel environment configurations.

Generalizing to novel environment configurations with modules. Goyal et al.
(2019) showed that leveraging modules to learn a state function improved out-
of-distribution generalization. Carvalho et al. (2021a) showed that a modified
attention mechanism led to strong generalization improvements with RL. MSFA
differs from both in that it employ modules for learning value functions in
the form of SFs. This enables a principled way to compose task knowledge while
additionally generalizing to novel environment configurations.

##  3 Problem Setting and Background

We study a reinforcement learning agent’s ability to transfer knowledge
between tasks in an environment. During training, the experiences $n_{\tt
train}$ tasks $\mathbb{M}_{\tt train}=\\{\mathcal{M}_{i}\\}_{i=1}^{n_{\tt
train}}$, sampled from a training distribution $p_{\tt train}(\mathcal{M})$.
During testing, the agent is evaluated on $n_{\tt test}$ tasks,
$\\{\mathcal{M}_{i}\\}_{i=1}^{n_{\tt test}}$, sampled from a testing
distribution $p_{\tt test}(\mathcal{M})$. Each task $\mathcal{M}_{i}$ is
specified as Partially Observable Markov Decision Process (POMDP),
$\mathcal{M}_{i}=\langle\mathcal{S}^{e},\mathcal{A},\mathcal{X},R,p,f_{x}\rangle$.
Here, $\mathcal{S}^{e}$, $\mathcal{A}$ and $\mathcal{X}$ are the environment
state, action, and observation spaces. $p(\cdot|s^{e}_{t},a_{t})$ specifies
the next-state distribution based on taking action $a_{t}$ in state
$s^{e}_{t}$, and $f_{x}(s^{e}_{t})$ maps the underlying environment state to
an observation $x_{t}$. We focus on tasks where rewards are parameterized by a
task vector $w$, i.e. $r^{w}_{t}=R(s^{e}_{t},a_{t},s^{e}_{t+1},w)$ is the
reward obtained for transition $(s^{e}_{t},a_{t},s^{e}_{t+1})$ given task
vector $w$. Since this is a POMDP, we need to learn a state function that maps
histories to agent state representations. We do so with a recurrent function:
$s_{t}=s_{\theta}(x_{t},s_{t-1},a_{t-1})$. Given this learned state, we want
to obtain a behavioral policy $\pi(s_{t})$ that best maximises the expected
reward it will obtain when taking an action $a_{t}$ at a state $s_{t}$:
$Q^{\pi,w}_{t}=Q^{\pi,w}(s_{t},a_{t})=\mathbb{E}_{\pi}\left[\sum^{\infty}_{t=0}\gamma^{t}r^{w}_{t}\right]$.

Figure 2: High-level diagram of how MSFA can be leveraged for transfer with
SF&GPI. During training, we can have the agent learn policies for tasks—e.g.
“open drawer” and “open fridge”. Each task leads the agent to experience
different aspects of the environment—e.g. a “fork” during “open drawer” or an
“apple” during “open fridge”. We can leverage MSFA to have different modules
learn different “cumulants”, $\phi$, and SFs, $\psi$. For example, module 1
($\theta_{1}$) can estimate SFs for apple cumulants. Module SFs are combined
to form the SF for a policy. When the agent wants to transfer its knowledge to
a test task—e.g., “get milk”—it can compute Q-values for that task as a dot-
product with the SFs of each training task. The highest Q-value is then used
to select actions.

Transfer with SF&GPI. In order to leverage SFs (Barreto et al., 2017), one
assumes an agent has access to state features known as “cumulants”,
$\phi_{t}=\phi(s_{t},a_{t},s_{t+1})$. Given a behavioral policy $\pi(a|s)$,
SFs are a type of value function that use $\phi_{t}$ as pseudo-rewards:

$\psi^{\pi}_{t}=\psi^{\pi}(s_{t},a_{t})=\mathbb{E}_{\pi}\left[\sum_{i=0}^{\infty}\gamma^{i}\phi_{t+i}\right]$
(1)

If reward is (approximately) $r^{w}_{t}=\phi_{t}^{\top}w$, then action-values
can be decomposed as $Q^{\pi,w}_{t}={\psi^{\pi}_{t}}^{\top}w$. This is
interesting because it provides an easy way to reuse task-agnostic features
$\psi^{\pi}_{t}$ for new tasks.

We can re-use the SFs we’ve learned from training tasks $\mathbb{M}_{\tt
train}$ for transfer with GPI. Assume we have learned (potentially optimal)
policies $\\{\pi_{i}\\}^{n_{\tt train}}_{i=1}$ and their corresponding SFs
$\\{\psi^{\pi_{i}}(s,a)\\}^{n_{\tt train}}_{i=1}$. Given a test task $w_{\tt
test}$, we can obtain a new policy with GPI in two steps: (1) compute Q-values
using the training task SFs (2) select actions using the highest Q-value. This
operation is summarized as follows:

$\pi(s_{t};w_{\tt
test})\in\operatorname*{arg\,max}_{a\in\mathcal{A}}\max_{i\in\\{1,\ldots,n_{\tt
train}\\}}\\{{Q^{\pi_{i},w_{\tt
test}}_{t}}\\}=\operatorname*{arg\,max}_{a\in\mathcal{A}}\max_{i\in\\{1,\ldots,n_{\tt
train}\\}}\\{{\psi^{\pi_{i}}_{t}}^{\top}w_{\tt test}\\}$ (2)

This is useful because the GPI theorem states that $\pi$ will perform as well
as all of the training policies, i.e. that $Q^{\pi,w_{\tt
test}}(s,a)\geq\max_{i}Q^{\pi_{i},w_{i}}(s,a)\forall(s,a)\in(\mathcal{S}\times\mathcal{A})$
(Barreto et al., 2017).

SF&GPI enable transfer by exploiting structure in the RL problem: a policy
that maximizes a value function is guaranteed to perform at least as well as
the policy that defined that value function. However, SF&GPI relies on
combining a fixed set of SFs. Another form of transfer comes from “Universal
Value Function Approximators” (UVFAs) (Schaul et al., 2015), which add the
task-vector $w$ as a parameter to a Q-approximator parameterized by $\theta$,
$Q_{\theta}(s,a,w)$. If $Q_{\theta}$ is smooth with respect to $w$, then
$Q_{\theta}$ should generalize to test tasks nearby to train tasks in task
space. Borsa et al. (2019) showed that one could combine the benefits of both
with “Universal Successor Feature Approximators”. Since rewards $r^{w}$, and
therefore task vectors $w$, reference deterministic task policies $\pi_{w}$,
one can parameterize successor feature approximators with task-vectors
$\tilde{\psi}^{\pi_{w}}=\tilde{\psi}^{w}\approx\psi_{\theta}(s,a,w)$. However,
USFA assumed hand-designed cumulants. We introduce an architecture for reward-
driven discovery of cumulants and improved function approximation of universal
successor features.

Figure 3:  Left: Universal Successor Feature Approximator (USFA) learns a
single, monolithic successor feature estimator that uses hand-designed
cumulants. Right: Modular Successor Feature Approximator (MSFA) learns a set
of successor feature modules, each with their own functions for (a) updating
module-state, (b) computing cumulants, and (c) estimating successor features.
Modules then share information with an attention mechanism. We hypothesize
that isolated module computations facilitate learning cumulants that suppot
generalization with GPI.

##  4 Modular Successor Feature Approximators

We propose a new architecture Modular Successor Feature Approximators (MSFA)
for approximating SFs, shown in Figure 3. Our hypothesis is that learning
cumulants and SFs with modules improves zero-shot composition of task
knowledge with SF&GPI. MSFA accomplishes this by learning $n$ state modules
$\\{s_{\theta_{k}}\\}_{k=1}^{n}$ that evolve with independent parameters
$\theta_{k}$ and have sparse inter-module information flow. MSFA then produces
modular cumulants $\\{\tilde{\phi}^{(k)}_{t}\\}^{n}_{k=1}$ and SFs
$\\{\tilde{\psi}^{\pi,k}_{t}\\}^{n}_{k=1}$ by having their computations depend
only on individual module-states. For example, a cumulant may correspond to
information about apples, and would be a function only of the module
representing state information related to apples. This is in contrast to prior
work, which learns a single monolithic prediction module for computing
cumulants and SFs (see Figure 3).

The rest of section is structured as follows. In section §4.1, we derive
Modular Successor Feature Learning within the Modular RL framework. We then
describe our architecture, MSFA, for learning modular successor features in
§4.2. In §4.3, we describe how to generate behavior with MSFA. Finally, we
describe the learning algorithm for MSFA in section §4.4.

###  4.1 Modular Successor Feature Learning

Following the Modular RL framework (Russell & Zimdars, 2003), we assume that
reward has an additive structure
$R(s_{t},a_{t},s_{t+1})=\sum_{k}R^{k}(s_{t},a_{t},s_{t+1})=\sum_{k}R^{(k)}_{t}$,
where $R^{(k)}_{t}$ is the reward of the $k$-th module. We enforce that every
module decomposes reward into an inner-product between its own task-
description $w^{(k)}$ and task-agnostic cumulants $\phi^{(k)}\in\mathbb{R}$:
$R^{(k)}_{t}={\phi^{(k)}_{t}}\cdot w^{(k)}$. Here, we simply break up the task
vector into $n$ pieces so individual modules are responsible for subsets of
the task vector. This allows us to decompose the action-value function as

$\displaystyle
Q^{\pi}(s_{t},a_{t},w)=\sum^{n}_{k=1}Q^{\pi,k}(s_{t},a_{t},w^{(k)})=\sum^{n}_{k=1}\psi^{\pi,k}(s_{t},a_{t})\cdot
w^{(k)}$ (3)

where we now have modular SFs $\\{\psi^{\pi,k}(s,a)\\}^{n}_{k=1}$ (see the
Appendix for a derivation). Rather than hand-designing modules or cumulants,
we aim to discover them from the environment reward signal.

###  4.2 Architecture

We learn a set of modules with states
$\mathbb{S}_{t}=\\{s^{(k)}_{t}\\}^{n}_{k=1}$. They update at each time-step
$t$ with the observation $x_{t}$, the previous module-state ${s}^{(k)}_{t-1}$,
and information from other modules
$A_{\theta}(s^{(k)}_{t-1},\mathbb{S}_{t-1})$. Following prior work (Santoro et
al., 2018; Goyal et al., 2019; Carvalho et al., 2021a), we have $A_{\theta}$
combine transformer-style attention (Vaswani et al., 2017) with a gating
mechanism (Parisotto et al., 2020) to enforce that inter-module interactions
are sparse. Since $A_{\theta}$ is not the main contribution of this paper, we
describe these computations in more detail with our notation in the Appendix.
We summarize the high-level update below.

$\displaystyle s^{(k)}_{t}$
$\displaystyle=s_{\theta_{k}}(x_{t},{s}^{(k)}_{t-1},A_{\theta}(s^{(k)}_{t-1},\mathbb{S}_{t-1}))$
(4)

We learn modular cumulants and SFs by having sets of cumulants and SFs depend
on individual module-states. Module cumulants depend on the module-state from
the current and next time-step. Module SFs depend on the current module-state
and on their subset of the task description. We summarize this below:

$\tilde{\phi}^{(k)}_{t}=\phi_{\theta}(s^{(k)}_{t},a_{t},s^{(k)}_{t+1})\quad\quad\tilde{\psi}^{w,k}_{t}=\psi_{\theta}(s^{(k)}_{t},a_{t},w^{(k)})$
(5)

We highlight that cumulants share parameters but differ in their input. This
suggests that the key is not having cumulants and SFs with separate parameters
but that they are functions of sparse subsets of state (rather than all state
information). We show evidence for this hypothesis in Figure 6.

We concatenate module-specific cumulants and SFs to form the final outputs:
$\tilde{\phi}_{t}=\left[\tilde{\phi}^{(1)}_{t},\ldots,\tilde{\phi}^{(n)}_{t}\right]$
and
$\tilde{\psi}^{w}_{t}=\psi_{\theta}(s_{t},a_{t},w)=\left[\tilde{\psi}^{{w},1}_{t},\ldots,\tilde{\psi}^{{w},n}_{t}\right]$.
Note that cumulants, are only used during learning, update with the module-
state from the next time-step.

###  4.3 Behavior

During training, actions are selected in proportion to Q-values computed using
task SFs as
$\pi(s_{t},w)\propto{\tilde{Q}(s_{t},a,w)}={{\psi}_{\theta}(s_{t},a,w)}^{\top}w$.
In practice we use an epsilon-greedy policy, though one can use other choices.
During testing, we compute policies with GPI as $\pi(s_{t},w_{\tt
test})\in\operatorname*{arg\,max}_{a}\max_{z\in\mathbb{M}_{\tt
train}}\\{{{\psi}_{\theta}(s_{t},a,z)}^{\top}w_{\tt test}\\}$, where
$\mathbb{M}_{\tt train}$ are train task vectors.

###  4.4 Learning Algorithm

MSFA relies on three losses. The first loss, $\mathcal{L}_{Q}$, is a standard
Q-learning loss, which MSFA uses to learn optimal policies for the training
tasks. The main difference here is that MSFA uses a particular
parameterization of the Q-function
$Q^{\pi_{w},w}(s,a)=\psi^{\pi_{w}}(s,a)^{\top}w$. The second loss,
$\mathcal{L}_{\psi}$, is an SF learning loss, which we use as a regularizer to
enforce that the Q-values follow the structure in the reward function
$r^{w}_{t}=\phi_{t}^{\top}w$. For this, we again apply standard Q-learning but
using SFs as value functions and cumulants as pseudo-rewards. The final loss,
$\mathcal{L}_{\phi}$ is a loss for learning cumulants that grounds them in the
environment reward signal. The losses are summarised as follows

$\displaystyle\mathcal{L}_{Q}$
$\displaystyle=||r_{t}+\gamma\psi_{\theta}(s_{t+1},a^{\prime},w)^{\top}w-\psi_{\theta}(s_{t},a_{t},w)^{\top}w||^{2}$
(6) $\displaystyle\mathcal{L}_{\psi}$
$\displaystyle=||{\tilde{\phi_{t}}}+\gamma\psi_{\theta}(s_{t+1},a^{\prime},w)-\psi_{\theta}(s_{t},a_{t},w)||^{2}$
(7) $\displaystyle\mathcal{L}_{\phi}$
$\displaystyle=||r^{w}_{t}-\tilde{\phi_{t}}^{\top}w||^{2}$ (8)

where
$a^{\prime}=\operatorname*{arg\,max}_{a}\psi_{\theta}(s_{t+1},a,w)^{\top}w$.
Selecting the next action via the combination of all modules ensures they
individually convergence to optimal values (Russell & Zimdars, 2003). The
final loss is
$\mathcal{L}=\alpha_{Q}\mathcal{L}_{Q}+\alpha_{\psi}\mathcal{L}_{\psi}+\alpha_{\phi}\mathcal{L}_{\phi}$.

##  5 Experiments

Figure 4: We study an agent’s ability to combine task knowledge in three
environments. (a) In BabyAI, an agent learns to pick up one object type at a
time during training. During testing, the agent must pickup combinations of
object types while avoiding other object types. This is the setting used by
USFA which assumed hand-designed cumulants. (b) In Procgen, we study extending
this form of generalization to a visually diverse, procedurally generated
environment. (c) In Minihack, we go beyond combining object navigation skills.
Here, an agents needs to combine (1) avoiding teleportation traps, (2)
avoiding monsters, and (3) partial visibility around the agent.

We study generalization when training behaviors must be combined concurrently
in time in the presence of novel object appearances and layouts. The need to
combine training behaviors tests how well MSFA can leverage SF&GPI.
Generalization to novel object appearances and layouts tests how well MSFA’s
modular construction supports generalization to novel environment
configurations.

Baselines. (1) Universal Value Function Approximator (UVFA) (Schaul et al.,
2015), which takes the task as input: $Q_{\theta}(s,w)$. This comparison shows
shows the transfer benefits of SF&GPI. (2) UVFA with Feature-Attending
Recurrent Modules (UVFA+FARM) instead takes state-factors as input
$Q_{\theta}(s^{(1)},\ldots,s^{(n)},w)$. Each state-factor $s^{(k)}$ is the
output of a FARM module. (3) Modular Value Function Approximator(MVFA) is an
adaptation of (Haarnoja et al., 2018) where modules learn individual Q-values
$Q^{(i)}_{\theta}(s^{(i)},w^{(i)})$. Comparing to UVFA+FARM and MVFA enables
us to study the benefits of leveraging modules for learning value functions in
the form of SFs. (4) Universal Successor Function Approximator (USFA) (Borsa
et al., 2019) leverages a single monolithic function for successor features
with hand-designed cumulants. USFA is an upper-bound baseline that allows us
to test the quality of cumulants and successor features that MSFA learns. We
also test a variant of USFA with learned cumulants, USFA-Learned-$\phi$, which
shows how the architecture degrades without oracle cumulants.

Implementation. We implement the state modules of MSFA with FARM modules
(Carvalho et al., 2021b). For UVFA and USFA, we learn a state representation
with an LSTM (Hochreiter & Schmidhuber, 1997). We implement all $\phi$,
$\psi$, and $Q$ functions with Mutli-layer Perceptrons. We train UVFA and
UVFA+FARM with n-step Q-learning (Watkins & Dayan, 1992). When learning
cumulants, USFA and MSFA have the exact same losses and learning alogirthm.
They both learn Q-values and SFs with n-step Q-learning. We use $n=5$. When
not learning cumulants, following (Borsa et al., 2019), USFA only learns SFs
with n-step Q-learning. All agents are built with JAX (Bradbury et al., 2018)
using the open-source ACME codebase (Hoffman et al., 2020) for reinforcement
learning.

###  5.1 Combining object navigation task knowledge

Figure 5:  MSFA matches USFA, which has hand-designed cumulants. We show mean
and standard error generalization episode return across 10 runs. We put a
task’s L2 distance to the closest train task in parenthesis. USFA best
generalizes to novel combinations of picking up and avoiding objects. Once
USFA learns cumulants, its performance degrades significantly. UVFA-based
methods struggle as more objects should be avoided or tasks are further in
distance to train tasks.

MSFA learns modular functions for computing $\phi$ and estimating $\psi$. We
hypothesize that this facilitates learning cumulants that respond to different
aspects of the environment (e.g. to different object categories). This leads
to the following research questions. R1: Can we recover prior generalization
results that relied on hand-designed cumulants for different object
categories? R2: How important is it to learn modular functions for $\phi$ and
$\psi$? R3: Without GPI, do learning modular functions for $\phi$ and $\psi$
still aid in generalization?

Setup. We implement a simplified version of the object navigation task of
(Borsa et al., 2019) in the (Chevalier-Boisvert et al., 2019) BabyAI
environment. The environment consists of $3$ instances of $4$ object
categories. Observations are partial and egocentric. Actions: the agent can
rotate left or right, move forward, or pickup an object. When it picks up an
object, following (Borsa et al., 2019), the object is respawned somewhere on
the grid. Task vectors lie in $w\in\mathbb{R}^{4}$ with training tasks being
the standard unit vectors. For example, $[0,1,0,0]$ specifies the agent must
obtain objects of category $2$. Generalization tasks are linear combinations
of training tasks. For example, $[-1,1,-1,1]$ specifies the agent must obtain
categories $2$ and $4$ while avoiding categories $1$ and $3$. Borsa et al.
(2019) showed that USFA could generalize with hand-designed cumulants that
described whether an object was picked up. We describe challenges for this
task in detail in the Appendix.

R1: MSFA is competitive with USFA, which uses oracle $\phi$. Figure 5 shows
USFA with a similar generalization trend to (Borsa et al., 2019). Tasks get
more challenging as they are further from train tasks or involve avoiding more
objects. For simply going to combinations of objects, USFA-Learned-$\phi$ does
slightly worse than MSFA. However, with more objects to avoid, all methods
except MSFA (including USFA-Learned-$\phi$) degrades significantly. For
comparison, we show performance by an oracle bread-first-search policy with
access to ground-truth state (BFS State Oracle). All methods have room for
improvement when objects must be avoided. In the appendix, we present heat-
maps for how often object categories were picked up during different tasks. We
find that MSFA most matches USFA, while USFA-Learned-$\phi$ commonly picks up
all objects regardless of task.

Figure 6:  Learning modular $\phi$ and $\psi$ is key to generalization and
improves generalization even without GPI. We show mean and standard error
generalization episode return across 10 runs. (a) We ablate having modular
functions for $\phi_{\theta}$ and $\psi_{\theta}$. Generalization results
degrade significantly. (b) We ablate leveraging GPI for generalization from
all SF-based methods. MSFA without GPI can outperform both USFA-Learned-$\phi$
with GPI and USFA without GPI. This shows the utility of modularity for
generalization.

R2: Learning modular $\phi$ and $\psi$ functions is critical for
generalization. Learning an entangled function corresponds to learning a
monolithic function for $\psi$ or $\phi$ where we concatenate module-states,
e.g. $\tilde{\psi}^{w}_{t}=\psi_{\theta}(s^{(1)}_{t},\ldots,s^{(n)}_{t},a,w)$.
Modular functions correspond to equation 5. Figure 6 (a) shows that without
modular functions for $\phi$ and $\psi$, performance severely degrades. This
also highlights that a naive combination of USFA+FARM—with entangled functions
for $\phi$ and $\psi$—does not recover our generalization performance.

R3: Modularity alone improves generalization. For all SF-based methods, we
remove GPI and select actions with a greedy policy:
$\pi(s)=\operatorname*{arg\,max}_{a}\psi_{\theta}{(s_{t},a,w)}^{\top}w$.
Figure 6 (b) shows that GPI is critical for generalization with USFA as
expected. USFA-Learned-$\phi$ benefits less from GPI (presumably because of
challenges in learning $\phi$). Interestingly, MSFA can generalize relatively
well without GPI, sometimes doing better than USFA without GPI.

###  5.2 Combining object navigation task knowledge with novel appearances and
environment configurations

Beyond generalizing to combinations of tasks, agents will need to generalize
to different layouts and appearances of objects. R4: Can MSFA enable combining
task knowledge in a visually diverse, procedurally generated environment?

Setup. We leverage the “Fruitbot” environment within ProcGen (Cobbe et al.,
2020). Here, an agent controls a paddle that tries to obtain certain
categories of objects while avoiding others. When the agent hits a wall or
fence, it dies and the episode terminates. If the agent collects a non-task
object, nothing happens. Observations are partial. Actions: At each time-step
the agent moves one step forward and can move left or right or shoot pellets
to open fences. Training and generalization tasks follow the same setup as
§5.1.

Figure 7: MSFA is able to combine task knowledge in a visually diverse,
procedurally generated ProcGen environment. We find that no method is able to
do well when there are objects to avoid ($w=[1,0,0,-1]$) in this setting (see
text for more). However, as more objects need to be collected MSFA best
generalizes (10 runs).

R4: MSFA enables combining task knowledge of object navigation tasks in a
visually diverse, procedurally generated environment. Figure 7 shows that when
an agent has to generalize to collecting more objects, modular architectures
generalize best, with MSFA doing best. When objects have to be avoided, we see
that no architecture does well, though MSFA tends to do better. We observe
that avoiding objects leads agents to hit walls. Since the agent always moves
forward at each time-step in Fruibot, this makes avoiding objects a
particularly challenging type of generalization.

##  6 Discussion and Conclusion

We have presented “Modular Successor Feature Approximators”, a modular neural
network for learning cumulants and SFs produced by their own modules. We first
showed that MSFA is competitive with prior object navigation generalization
results that relied on hand-designed cumulants (§5.1). Afterwards, we showed
that MSFA improves an agent’s ability to combine task knowledge in a visually
diverse, procedurally generated environment (§5.2). We also show that MSFA can
combine solutions of heterogeneous tasks (§B). Our ablations show that
learning modular cumulants and SFs is critical for generalization with GPI.

We compared MSFA to (1) USFA, a monolithic architecture for learning cumulants
and SFs; (2) FARM, an architecture which learns multiple state-modules but
combines them with a monolithic Q-value function. Our results show that when
learning cumulants, MSFA improves generalization with SF&GPI compared to USFA.
Additionally, without GPI, MSFA as an architecture improves generalization as
compared to both FARM and USFA.

Limitations. While we demonstrated reward-driven discovery of cumulants for
transfer with SF&GPI, we focused on relatively simple task encodings. Future
work can extend this to more expressive encodings such as language embeddings.
Another limitation is that we did not explore more sophisticated state-feature
discovery methods such as meta-gradients (Veeriah et al., 2019). Nonetheless,
we think that MSFA provides an important insight for future work: that
modularity is a simple but powerful inductive bias for discovering state-
features that enable zero-shot transfer with SF&GPI.

Future directions. SFs are useful for exploration (Janz et al., 2019; Machado
et al., 2020); for discovering and combining options (Barreto et al., 2019;
Hansen et al., 2019); for transferring policies across environments (Zhang et
al., 2017); for improving importance sampling (Fujimoto et al., 2021); and for
learning policies from other agents (Filos et al., 2021). We hope that future
work can leverage MSFA for improved state-feature discovery and SF-learning in
all of these settings.

## Acknowledgments

This work was supported in part by a University of Michigan Rackham Merit
Fellowship and grant from LG AI Research. The authors would also like to thank
Andrew Lampinen, Cameron Allen, and the anonymous reviewers for their helpful
comments.

## References

  * Barreto et al. (2017) André Barreto, Will Dabney, Rémi Munos, Jonathan J Hunt, Tom Schaul, Hado P van Hasselt, and David Silver.  Successor features for transfer in reinforcement learning.  _NIPS_ , 30, 2017.
  * Barreto et al. (2018) André Barreto, Diana Borsa, John Quan, Tom Schaul, David Silver, Matteo Hessel, Daniel Mankowitz, Augustin Zidek, and Remi Munos.  Transfer in deep reinforcement learning using successor features and generalised policy improvement.  In _ICML_ , pp. 501–510. PMLR, 2018.
  * Barreto et al. (2019) André Barreto, Diana Borsa, Shaobo Hou, Gheorghe Comanici, Eser Aygün, Philippe Hamel, Daniel Toyama, Shibl Mourad, David Silver, Doina Precup, et al.  The option keyboard: Combining skills in reinforcement learning.  _Advances in Neural Information Processing Systems_ , 32, 2019.
  * Barreto et al. (2020) André Barreto, Shaobo Hou, Diana Borsa, David Silver, and Doina Precup.  Fast reinforcement learning with generalized policy updates.  _PNAS_ , 117(48):30079–30087, 2020.
  * Borsa et al. (2019) Diana Borsa, André Barreto, John Quan, Daniel Mankowitz, Rémi Munos, Hado Van Hasselt, David Silver, and Tom Schaul.  Universal successor features approximators.  _ICLR_ , 2019.
  * Bradbury et al. (2018) James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milne, and Qiao Zhang.  JAX: composable transformations of Python+NumPy programs, 2018\.  URL http://github.com/google/jax.
  * Brooks et al. (2021) Ethan Brooks, Janarthanan Rajendran, Richard L Lewis, and Satinder Singh.  Reinforcement learning of implicit and explicit control flow instructions.  In _International Conference on Machine Learning_ , pp. 1082–1091. PMLR, 2021.
  * Carvalho et al. (2021a) Wilka Carvalho, Andrew Lampinen, Kyriacos Nikiforou, Felix Hill, and Murray Shanahan.  Feature-attending recurrent modules for generalization in reinforcement learning.  _arXiv_ , 2021a.
  * Carvalho et al. (2021b) Wilka Carvalho, Anthony Liang, Kimin Lee, Sungryull Sohn, Honglak Lee, Richard L Lewis, and Satinder Singh.  Reinforcement learning for sparse-reward object-interaction tasks in a first-person simulated 3d environment.  _IJCAI_ , 2021b.
  * Chevalier-Boisvert et al. (2018) Maxime Chevalier-Boisvert, Dzmitry Bahdanau, Salem Lahlou, Lucas Willems, Chitwan Saharia, Thien Huu Nguyen, and Yoshua Bengio.  Babyai: A platform to study the sample efficiency of grounded language learning.  _arXiv preprint arXiv:1810.08272_ , 2018.
  * Chevalier-Boisvert et al. (2019) Maxime Chevalier-Boisvert, Dzmitry Bahdanau, Salem Lahlou, Lucas Willems, Chitwan Saharia, Thien Huu Nguyen, and Yoshua Bengio.  BabyAI: First steps towards grounded language learning with a human in the loop.  In _International Conference on Learning Representations_ , 2019.  URL https://openreview.net/forum?id=rJeXCo0cYX.
  * Cobbe et al. (2020) Karl Cobbe, Chris Hesse, Jacob Hilton, and John Schulman.  Leveraging procedural generation to benchmark reinforcement learning.  In _International conference on machine learning_ , pp. 2048–2056. PMLR, 2020.
  * Corona et al. (2020) Rodolfo Corona, Daniel Fried, Coline Devin, Dan Klein, and Trevor Darrell.  Modular networks for compositional instruction following.  _arXiv preprint arXiv:2010.12764_ , 2020.
  * Devin et al. (2017) Coline Devin, Abhishek Gupta, Trevor Darrell, Pieter Abbeel, and Sergey Levine.  Learning modular neural network policies for multi-task and multi-robot transfer.  In _2017 IEEE international conference on robotics and automation (ICRA)_ , pp. 2169–2176. IEEE, 2017.
  * Espeholt et al. (2018) Lasse Espeholt, Hubert Soyer, Remi Munos, Karen Simonyan, Vlad Mnih, Tom Ward, Yotam Doron, Vlad Firoiu, Tim Harley, Iain Dunning, et al.  Impala: Scalable distributed deep-rl with importance weighted actor-learner architectures.  In _International conference on machine learning_ , pp. 1407–1416. PMLR, 2018.
  * Filos et al. (2021) Angelos Filos, Clare Lyle, Yarin Gal, Sergey Levine, Natasha Jaques, and Gregory Farquhar.  Psiphi-learning: Reinforcement learning with demonstrations using successor features and inverse temporal difference learning.  In _International Conference on Machine Learning_ , pp. 3305–3317. PMLR, 2021.
  * Fujimoto et al. (2021) Scott Fujimoto, David Meger, and Doina Precup.  A deep reinforcement learning approach to marginalized importance sampling with the successor representation.  In _International Conference on Machine Learning_ , pp. 3518–3529. PMLR, 2021.
  * Goyal et al. (2019) Anirudh Goyal, Alex Lamb, Jordan Hoffmann, Shagun Sodhani, Sergey Levine, Yoshua Bengio, and Bernhard Schölkopf.  Recurrent independent mechanisms.  _arXiv preprint arXiv:1909.10893_ , 2019.
  * Haarnoja et al. (2018) Tuomas Haarnoja, Vitchyr Pong, Aurick Zhou, Murtaza Dalal, Pieter Abbeel, and Sergey Levine.  Composable deep reinforcement learning for robotic manipulation.  In _2018 IEEE international conference on robotics and automation (ICRA)_ , pp. 6244–6251. IEEE, 2018.
  * Hansen et al. (2019) Steven Hansen, Will Dabney, Andre Barreto, Tom Van de Wiele, David Warde-Farley, and Volodymyr Mnih.  Fast task inference with variational intrinsic successor features.  _arXiv preprint arXiv:1906.05030_ , 2019.
  * Hochreiter & Schmidhuber (1997) Sepp Hochreiter and Jürgen Schmidhuber.  Long short-term memory.  _Neural computation_ , 9(8):1735–1780, 1997.
  * Hoffman et al. (2020) Matt Hoffman, Bobak Shahriari, John Aslanides, Gabriel Barth-Maron, Feryal Behbahani, Tamara Norman, Abbas Abdolmaleki, Albin Cassirer, Fan Yang, Kate Baumli, et al.  Acme: A research framework for distributed reinforcement learning.  _arXiv preprint arXiv:2006.00979_ , 2020.
  * Huang et al. (2020) Wenlong Huang, Igor Mordatch, and Deepak Pathak.  One policy to control them all: Shared modular policies for agent-agnostic control.  In _International Conference on Machine Learning_ , pp. 4455–4464. PMLR, 2020.
  * Hunt et al. (2019) Jonathan Hunt, Andre Barreto, Timothy Lillicrap, and Nicolas Heess.  Composing entropic policies using divergence correction.  In _International Conference on Machine Learning_ , pp. 2911–2920. PMLR, 2019.
  * Janz et al. (2019) David Janz, Jiri Hron, Przemysław Mazur, Katja Hofmann, José Miguel Hernández-Lobato, and Sebastian Tschiatschek.  Successor uncertainties: exploration and uncertainty in temporal difference learning.  _Advances in Neural Information Processing Systems_ , 32, 2019.
  * Kapturowski et al. (2018) Steven Kapturowski, Georg Ostrovski, John Quan, Remi Munos, and Will Dabney.  Recurrent experience replay in distributed reinforcement learning.  In _International conference on learning representations_ , 2018.
  * Logeswaran et al. (2021) Lajanugen Logeswaran, Wilka Torrico Carvalho, and Honglak Lee.  Learning compositional tasks from language instructions.  In _Deep RL Workshop NeurIPS 2021_ , 2021.
  * Machado et al. (2020) Marlos C Machado, Marc G Bellemare, and Michael Bowling.  Count-based exploration with the successor representation.  In _Proceedings of the AAAI Conference on Artificial Intelligence_ , number 04, pp. 5125–5133, 2020.
  * Madan et al. (2021) Kanika Madan, Nan Rosemary Ke, Anirudh Goyal, Bernhard Schölkopf, and Yoshua Bengio.  Fast and slow learning of recurrent independent mechanisms.  _arXiv preprint arXiv:2105.08710_ , 2021.
  * Mnih et al. (2015) Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Marc G Bellemare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, et al.  Human-level control through deep reinforcement learning.  _nature_ , 518(7540):529–533, 2015.
  * Oh et al. (2017) Junhyuk Oh, Satinder Singh, Honglak Lee, and Pushmeet Kohli.  Zero-shot task generalization with multi-task deep reinforcement learning.  In _International Conference on Machine Learning_ , pp. 2661–2670. PMLR, 2017.
  * Parisotto et al. (2020) Emilio Parisotto, Francis Song, Jack Rae, Razvan Pascanu, Caglar Gulcehre, Siddhant Jayakumar, Max Jaderberg, Raphael Lopez Kaufman, Aidan Clark, Seb Noury, et al.  Stabilizing transformers for reinforcement learning.  In _International conference on machine learning_ , pp. 7487–7498. PMLR, 2020.
  * Russell & Zimdars (2003) Stuart J Russell and Andrew Zimdars.  Q-decomposition for reinforcement learning agents.  In _Proceedings of the 20th International Conference on Machine Learning (ICML-03)_ , pp. 656–663, 2003.
  * Samvelyan et al. (2021) Mikayel Samvelyan, Robert Kirk, Vitaly Kurin, Jack Parker-Holder, Minqi Jiang, Eric Hambro, Fabio Petroni, Heinrich Küttler, Edward Grefenstette, and Tim Rocktäschel.  Minihack the planet: A sandbox for open-ended reinforcement learning research.  _arXiv preprint arXiv:2109.13202_ , 2021.
  * Santoro et al. (2018) Adam Santoro, Ryan Faulkner, David Raposo, Jack Rae, Mike Chrzanowski, Theophane Weber, Daan Wierstra, Oriol Vinyals, Razvan Pascanu, and Timothy Lillicrap.  Relational recurrent neural networks.  _Advances in neural information processing systems_ , 31, 2018.
  * Schaul et al. (2015) Tom Schaul, Daniel Horgan, Karol Gregor, and David Silver.  Universal value function approximators.  In _International conference on machine learning_ , pp. 1312–1320. PMLR, 2015.
  * Singh (1992) Satinder Pal Singh.  Transfer of learning by composing solutions of elemental sequential tasks.  _Machine Learning_ , 8(3):323–339, 1992.
  * Sohn et al. (2020) Sungryull Sohn, Hyunjae Woo, Jongwook Choi, and Honglak Lee.  Meta reinforcement learning with autonomous inference of subtask dependencies.  _arXiv preprint arXiv:2001.00248_ , 2020.
  * Sohn et al. (2022) Sungryull Sohn, Hyunjae Woo, Jongwook Choi, Lyubing Qiang, Izzeddin Gur, Aleksandra Faust, and Honglak Lee.  Fast inference and transfer of compositional task structures for few-shot task generalization.  In _Uncertainty in Artificial Intelligence_ , pp. 1857–1865. PMLR, 2022.
  * Vaswani et al. (2017) Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin.  Attention is all you need.  _arXiv_ , 2017.
  * Veeriah et al. (2019) Vivek Veeriah, Matteo Hessel, Zhongwen Xu, Janarthanan Rajendran, Richard L Lewis, Junhyuk Oh, Hado P van Hasselt, David Silver, and Satinder Singh.  Discovery of useful questions as auxiliary tasks.  _Advances in Neural Information Processing Systems_ , 32, 2019.
  * Watkins & Dayan (1992) Christopher JCH Watkins and Peter Dayan.  Q-learning.  _Machine learning_ , 8(3):279–292, 1992.
  * Zhang et al. (2017) Jingwei Zhang, Jost Tobias Springenberg, Joschka Boedecker, and Wolfram Burgard.  Deep reinforcement learning with successor features for navigation across similar environments.  In _2017 IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS)_ , pp. 2371–2378. IEEE, 2017.
  * Zhu et al. (2017) Yuke Zhu, Daniel Gordon, Eric Kolve, Dieter Fox, Li Fei-Fei, Abhinav Gupta, Roozbeh Mottaghi, and Ali Farhadi.  Visual semantic planning using deep successor representations.  In _Proceedings of the IEEE international conference on computer vision_ , pp. 483–492, 2017.

##  Appendix A Learning and Generalization Challenges

Our experimental setup for §5.1 follows the same setup as Barreto et al.
(2018). For $n$ tasks with reward functions $\\{r_{i}\\}^{n}_{i=1}$, they
showed that discovered features which predict rewards, i.e.
$[\tilde{\phi}_{1},\ldots,\tilde{\phi}_{n}]=[\tilde{r}_{1},\ldots,\tilde{r}_{n}]$,
were effective features for transfer with SF&GPI. In their experimental setup,
training tasks were one-hot vectors, so regressing
$||r-\tilde{\phi}^{\top}w||$ effectively led $\phi_{i}$ to be a feature that
predicted reward $r_{i}$. We use this same setup, where now individual modules
are used for both $\tilde{\phi}^{(i)}$ and $\tilde{\psi}^{(i)}$.

Figure 8: We present a diagram providing intuition for the type of
generalization challenge required for our tasks. Green stars represent our
training tasks. Red diamonds represent test tasks. They are both outside of
training distribution for task encodings and can be far in task-space from
training task encodings.

While this is a simple method for learning features, it some difficulties.
First, there is an imbalance of rewarding vs. non-rewarding transitions, which
leads to challenges in learning $\phi$. In our setting, an optimal policy only
has $7\%$ of its transitions be rewarding transitions. Over the course of
training, the percentage is much lower. Second, Barreto et al. (2018) were
mainly able to show continual learning results with additional learning using
their discovered features and had limited success with zero-shot transfer.

Zero-shot transfer is challenges because training tasks are “out-of-
distribution” while additionally requiring behaviors that the agent was never
trained on. By out-of-distribution, we mean that the agent is only trained on
basis vectors (see Figure 8). However, at test time, the agent must perform
tasks that are are far in “task-space” from the training distribution. For
example, $w_{\tt test}=[-1,1,-1,-1]$ has an L2 distance of $1.73$ from closest
training task. This requires an agent that can generalize to test reward
functions that are quite different from training reward functions. While both
USFA-Learned-$\phi$ and MSFA get about the same training error for predicting
rewards, we see that MSFA gets a dramatically better reward prediction error
for test reward functions $\tilde{r}_{\tt test}=\tilde{\phi}^{\top}w_{\tt
test}$.

Reward prediction error on generalization tasks. To investigate this, we
collect 40 episodes of each generalization task and compute the reward
prediction error of MSFA and USFA-Learned-$\phi$ for each task. We present the
results in Table 1.

Task | MSFA Error | USFA-Learned-$\phi$ Error
---|---|---
1,1,0,0 | $0.0003\pm 0.0090$ | $0.1362\pm 0.3410$
1,1,.5,.5 | $0.0069\pm 0.0405$ | $0.3922\pm 0.7051$
1,1,1,1 | $0.0118\pm 0.0856$ | $2.5530\pm 4.0324$
-.5,1,-.5,-.5 | $0.0095\pm 0.0478$ | $0.0556\pm 0.1041$
-1,1,0,1 | $0.0014\pm 0.0264$ | $0.0652\pm 0.2470$
Table 1: MSFA has a smaller reward prediction error for test tasks. We present
the reward prediction error for MSFA and USFA-Learned-$\phi$. We present
results using 40 episodes for each test task.

##  Appendix B Combining knowledge of heterogeneous tasks

Aside from navigation tasks, researchers and practitioners will be interested
in generally combining solutions to different tasks. R5: Can MSFA enable
combining solutions to heterogeneous tasks?

Figure 9: We study whether MSFA can combine three separate settings (1)
limited visibility (2) the presence of monsters (3) the presence of
teleportation traps (see text for more). All successor feature based methods
best enable combining training tasks for generalization in this setting. (10
run)

Setup. We leverage the “Minihack” environment (Samvelyan et al., 2021). Here,
an agent is spawned at the top of a map with obstacles. It must navigate to
the a ladder at the bottom of the map. The agent experiences partial
observations of the environment. Actions: At each time-step the agent can
select from $8$ compass directions. Training tasks: The agent experiences
three training tasks (1) avoiding monsters which kill the agent, (2) avoiding
traps which teleport the agent away from the goal, and (3) experiencing a
small local view of the environment with everything else black. Generalization
requires that the agent generalize to (a) new configurations of objects and
(b) to the combination of all training conditions.

R5: MSFA enables combining solutions to heterogeneous tasks. Figure 9 shows
that SF-based methods best combine solutions to different tasks. In a small
room, both SF-based methods get 100% success rate. In a larger room, no method
generalizes perfectly: SF-based methods do best at 75%, followed by UVFA-FARM
and UVFA with 60%.

##  Appendix C Derivation of Modular Successor Features

$\displaystyle Q^{\pi}(s,a,w)$
$\displaystyle=\mathbb{E}_{\pi}\left[\sum^{\infty}_{j=0}\gamma^{j}R_{t}|S_{t}=s,A_{t}=a\right]$
(9)
$\displaystyle=\mathbb{E}_{\pi}\left[\sum^{\infty}_{j=0}\gamma^{j}\left(\sum^{n}_{k=1}{\phi_{t+j}^{(k)}}^{\top}w^{(k)}\right)|S_{t}=s,A_{t}=a\right]$
(10)
$\displaystyle=\sum^{n}_{k=1}\left(\mathbb{E}_{\pi}\left[\sum^{\infty}_{j=0}\gamma^{j}{\phi_{t+j}^{(k)}}^{\top}w^{(k)}|S_{t}=s,A_{t}=a\right]\right)$
(11)
$\displaystyle=\sum^{n}_{k=1}\left(\mathbb{E}_{\pi}\left[\sum^{\infty}_{j=0}\gamma^{j}{\phi_{t+j}^{(k)}}|S_{t}=s,A_{t}=a\right]^{\top}w^{(k)}\right)$
(12) $\displaystyle=\sum^{n}_{k=1}\psi^{\pi,k}(s,a)^{\top}w^{(k)}$ (13)
$\displaystyle=\sum^{n}_{k=1}Q^{\pi,k}(s,a,w^{(k)})$ (14)

##  Appendix D Inter-module attention and gating

At each time-step $t$, each module updates with both observation features
$z_{t}=f^{\theta}_{z}(x_{t})$ and with information from the previous module-
states $\\{s^{(k)}_{t-1}\\}$. MSFA shares information between modules with a
“query-key” system. Each module has a “query” representation of its module-
state that it uses to select from “key” representations of the other module-
states. MSFA scores how well a query matches a key by computing the dot-
product between the two representations. A module then updates with the
module-state corresponding to the key with the highest dot-product. To enable
flexible updating, each module uses a gating mechanism to decide the degree to
which information should be incorporated into an update.

More technically, the query vector is computed using the previous-module state
and action: $q^{(k)}_{t}=W^{\tt
query}[s^{(k)}_{t-1},a_{t-1}]\in\mathbb{R}^{d_{q}}$. Keys and values are
computed as $K_{t}=W^{\tt
key}[s_{t-1}^{1};\ldots;s_{t-1}^{n};0]\in\mathbb{R}^{n+1\times d_{q}}$ and
$V_{t}=W^{\tt value}[s_{t-1}^{1};\ldots;s_{t-1}^{n};0]\in\mathbb{R}^{n+1\times
d_{q}}$, respectively. Note that we add a zero-vector key and value to allow a
query to select “no information”. Each module selects “values” $v^{(k)}_{t}$
to update using dot-product attention (Vaswani et al., 2017):

$\displaystyle
v^{(k)}_{t}=\mathrm{softmax}\left(\frac{q_{t}^{(k)}K^{\top}_{t}}{d_{q}}\right)V_{t}$
(15)

Gating mechanism has been shown important for leveraging transformer-style
attention when updating state in RL agents (Parisotto et al., 2020). Thus,
each module uses a sigtanh gate when updating:

$u^{(k)}_{t}=q^{(k)}_{t}+\tanh(W^{g_{1}}v^{(k)}_{t})\odot\sigma(W^{g_{2}}v^{(k)}_{t}-b_{g}).$
(16)

MSFA then updates its module-stateas follows:

$s^{(k)}_{t}=s_{\theta_{k}}(u^{(k)}_{t},z_{t})$ (17)

##  Appendix E Learning modular functions

Our goal is to learn functions for producing cumulants $\phi$ and SFs $\psi$
that have consider only state information from their own modules. To be more
precise, we learn modules which collectively learn a set of $n$ module-state
representations: $\\{s^{i}\\}^{n=1}_{i}$. As an example, consider learning a
linear function/transformation $A$ for producing cumulants. A typical
“monolithic function” would be one like the following:

$\begin{bmatrix}\phi^{1}\\\ \vdots\\\
\phi^{n}\end{bmatrix}=\begin{bmatrix}A_{11}&\cdots&A_{1m}\\\
\vdots&A_{22}&\vdots\\\
A_{n1}&\cdots&A_{nm}\end{bmatrix}\begin{bmatrix}s^{1}\\\ \vdots\\\
s^{n}\end{bmatrix}.$

By "entangled," we mean that a cumulant $\phi^{i}$ is (potentially) a function
of all module-state information $\phi^{i}=\sum_{j}A_{ij}s^{j}$. By a "modular
function", we simply meant that we were learning a function like the following

$\begin{bmatrix}\phi^{1}\\\ \vdots\\\
\phi^{n}\end{bmatrix}=\begin{bmatrix}A_{1}&0&0\\\ 0&A_{2}&0\\\
0&0&A_{n}\end{bmatrix}\begin{bmatrix}s^{1}\\\ \vdots\\\ s^{n}\end{bmatrix}.$

In our setting, we had $A_{1}=\ldots=A_{n}$ be a shared MLP and each $s^{i}$
was produced by a module with its own parameters. Our experiments demonstrate
leveraging modules to learn $\phi^{(i)}$ and $\psi^{(i)}$ is an effective way
to learn $\\{\phi^{i}\\}_{i}$ that promote zero-shot composition of task
knowledge with the SF&GPI framework.

##  Appendix F Additional analysis

Figure 10: We present a visualization of module activity for task
generalization task $w=[1,1,1,1]$ from §5.1. On the top, (a) we show the
activity of each module across time. On the bottom, (b) we show example
transitions where modules activate to visualize what they respond to. We see
that individual modules are able to effectively respond to different object
categories, despite a strong data imbalance involved in learning to represent
each object category (see §). We find that modules don’t perfectly specialize,
and multiple modules will sometimes activate at the same state.  Figure 11:
MSFA best matches the behavior of USFA across generalization tasks. We present
additional analysis for the experiments in §5.1. Each heatmap displays how
often object categories were picked up by each method for a given
generalization task. Rows correspond to different methods. We show the final
episode return for each method along the y-axis. Columns describe the
transition-reward when an object category is picked up. USFA has access to
hand-designed cumulants that described whether an object category was picked
up. Despite learning cumulants, MSFA best matches USFA’s object collection
dynamics. USFA-Learned-$\phi$ equally collects all objects regardless of task,
except for when $r\leq-1$ when an object is picked up. In this setting, no
method except USFA or MSFA collect objects. As a reminder, negative rewards
are never experienced during training and these tasks are the furthest away
from training tasks in task-space. UVFA-based methods tend to collect
rewarding objects but not as effectively as MSFA or USFA.  Figure 12: We
present the cross-correlation between pairs of modules. This confirms that
modules tend to activate at different time-points on a statistical level.

##  Appendix G Full results

Figure 13:  All results for ablating disentanglement of $\phi$ and $\psi$. We
consistently find that disentangled functions for $\phi$ and $\psi$ have the
best generalization performance.  Figure 14: All results for ablating GPI.
First, we compare MSFA to USFA (which has hand-designed cumulants). We can see
that the two perform comparably across the board. Sometimes (e.g. $[1,1,1,1]$
and $[-1,1,-1,-1]$), MSFA outperforms USFA. Next we compare MSFA to USFA-
Learned-$\phi$. We see that MSFA performs as well as USFA-Learned-$\phi$ or
better in all settings except $[-1,1,0,1]$. Interestingly, MSFA without GPI
outperforms USFA-Learned-$\phi$ with GPI in some generalization settings (e.g.
$[1,1,0,0]$ and $[-1,1,-1,1]$).  Figure 15: All results for the Fruitbot
environment. When objects must be avoided, we see that no method does well,
though MSFA generalizes best and can pick up a couple of objects. When
combinations of objects must be collected, MSFA outperforms baselines by a
large margin.

##  Appendix H Hyperparameters

We shown hyperparameters in Tables 2 and 3. Tables 2 describes hyperparameters
that were shared across algorithms and Tables 3 describes hyperparameters that
were different across algorithms. We developed the algorithms using the
Hoffman et al. (2020) reinforcement learning codebase and unless stated
otherwise using default “R2D2” (Kapturowski et al., 2018) config values in the
codebase. We describe our development and hyperparameter search in more detail
below.

###  H.1 Development and hyperparameter search

The first thing we did was replicate the generalization gap between USFA and
UVFA that Borsa et al. (2019) produced in their paper using the BabyAI
gridworld (Chevalier-Boisvert et al., 2018). As a result, most hyperparameters
were tuned for either USFA or UVFA. We used the Atari Convolutional network
used by “Deep Q-Networks” (DQN) in Mnih et al. (2015) as our visual encoder
and an LSTM (Hochreiter & Schmidhuber, 1997) as our state representation
function.

Once we were able to replicate their generalization performance, we
implemented USFA-Learned-$\phi$ and MSFA. In this setting, we are learning
Q-values, successor features, and predicting rewards from cumulants. This
leads to the following coefficients for losses: $\alpha_{Q}$, $\alpha_{\psi}$,
and $\alpha_{\phi}$ respectivey. We fixed $\alpha_{\psi}=1$ and searched
$\alpha_{\phi}\in[.01,.1,1,10,100]$ and $\alpha_{Q}\in[.01,.05,.1,.5,1.0]$ for
both USFA-Learned-$\phi$ and MSFA. For MSFA and UVFA-FARM, we chose the number
of modules to be the number of training tasks. We set each module size so that
the number of parameters by all algorithms was approximately the same (with
USFA as a reference). We found that FARM hyperparameter (Carvalho et al.,
2021a) worked well for both UVFA-FARM and for MSFA (FARM modules are the basis
of MSFA state-modules). We highlight that properly masking all losses was
absolutely critical to good generalization performance. Without proper
masking, out-of-episode data is used for value estimation of in-episode data,
which can lead to inaccurate value estimates. We suspect that this hampers
GPI.

When doing experiments for Fruitbot, we took all implementation details from
(Cobbe et al., 2020). We first replicated the DQN performance from Cobbe et
al. (2020) with UVFA. Some important changes that we made included (1) adding
prioritized experience replay (2) leveraging the Impala ResNet vision torso
(Espeholt et al., 2018) (3) lowering the memory size to match the paper (4)
using the ${\tt SIGNED\\_HYPERBOLIC\\_PAIR}$ value transformation for UVFA-
based algorithms (5) increasing the capacity of the Q-value estimation MLPs.
We applied these settings to other baselines and redid our search over
$\alpha_{\phi}$ and $\alpha_{Q}$ for both MSFA and USFA-Learned-$\phi$. For
each method, we changed the size of the networks so they were approximately
the same (with UVFA as a reference). Since MSFA uses disentangled functions
for $\phi$ and $\psi$, it has fewer parameters.

When doing experiments in Minihack, we found that the Fruitbot changes were
important for good performance in this domain. All hyperparameters led to
strong training performance but poor generalization for UVFA. We found that
increasing the batch size was important for improving generalization of USFA-
Learned-$\phi$ and MSFA but was detrimental to UVFA-based methods.

Table 2: Hyperparameters shared across all algorithms. Algorithm |  |  |
---|---|---|---
Loss Hyperparameters | BabyAI | Fruibot | Minihack
discount | 0.99 | 0.99 | 0.99
burn_in_length | 0 | 0 | 0
trace_length | 40 | 40 | 40
importance_sampling_exponent | 0 | 0.6 | 0.6
max replay size | 100,000 | 70,000 | 70,000
max gradient norm | 80 | 80 | 80
Shared Network Components |  |  |
Vision torso | DQN ConvNet | Impala ResNet | Impala ResNet
Table 3: Hyperparameters for individual algorithms. Note: $T_{1}={\tt IDENTITY\\_PAIR}$ and $T_{2}={\tt SIGNED\\_HYPERBOLIC\\_PAIR}$.  | BabyAI | Fruibot | Minihack
---|---|---|---
MSFA |  |  |
Parameters (millions) | 2.1M | 1.66M | 1.7M
$\alpha_{\phi}$ | 1 | 1 | 1
$\alpha_{\psi}$ | 1 | 1 | 1
$\alpha_{Q}$ | 0.5 | 0.5 | 0.5
module_size | 150 | 60 | 80
nmodules | 4 | 4 | 3
attention heads | 2 | 2 | 1
batch size | 32 | 32 | 64
$\phi$ MLP hidden sizes | [256] | [256] | [256]
$\psi$ MLP hidden sizes | [128] | [512, 512] | [512, 512]
projection dim | 16 | 16 | 16
USFA |  |  |
Parameters (millions) | 2.35M | 1.98M | 1.95M
lstm size | 512 | 300 | 300
$\alpha_{\phi}$ | 1 | 1 | 1
$\alpha_{\psi}$ | 1 | 1 | 1
$\alpha_{Q}$ | 0.5 | 0.5 | 0.5
batch size | 32 | 32 | 128
$\phi$ MLP hidden sizes | [256] | [256] | [256]
$\psi$ MLP hidden sizes | [128] | [512, 512] | [512, 512]
UVFA |  |  |
Parameters (millions) | 2.09M | 1.96M | 1.95M
lstm size | 512 | 256 | 256
rlax.TxPair | $T_{1}$ | $T_{2}$ | $T_{2}$
batch size | 32 | 32 | 32
Q MLP hidden sizes | [128] | [512, 512] | [512, 512]
UVFA-FARM |  |  |
Parameters (millions) | 2.17M | 2.15M | 2.06M
module_size | 150 | 64 | 80
nmodules | 4 | 4 | 3
attention heads | 2 | 2 | 1
rlax.TxPair | $T_{1}$ | $T_{2}$ | $T_{2}$
batch size | 32 | 32 | 32
projection dim | 16 | 16 | 16
Q MLP hidden sizes | [128] | [512, 512] | [512, 512]

##  Appendix I Environments

We presented most environment information in the main text. Here, we specify
which levels we used from the corresponding environments or how the changed
environments for our experiments.

###  I.1 Procgen

We used the Fruitbot environment in Procgen. We made the following changes. We
set the maximum number of levels that an agent could complete within a
lifetime to $4$. We divided the 12 object types $\\{O_{1}\ldots O_{12}\\}$
into 4 categories:
$C_{1}=\\{O_{1},O_{2},O_{3}\\},\ldots,C_{4}=\\{O_{10},O_{11},O_{12}\\}$, where
each category provided postive reward for picking up any of its 3 category
types.

###  I.2 Minihack

We did not make any changes to the environment. Train tasks were: “MiniHack-
Room-Monster-Y-v0”, “MiniHack-Room-Trap-Y-v0”, and “MiniHack-Room-Dark-Y-v0”.
Test tasks were “MiniHack-Room-Ultimate-Y-v0”. We used $Y=5\times 5$ for the
“small room” experiments and $Y=15\times 15$ for the “large room” experiments.