Composing Task Knowledge with Modular Successor Feature Approximators
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 eﬀective 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 architectues for learning state representations.
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 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 conﬁgurations.
Returning to our example, our robot might need to additionally generalize to both novel dirty dishes on the table and to novel arrangements of the chairs around the table.
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 eﬃciently 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 conﬁgurations.
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.
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 statefeature 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 zeroshot 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.
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 diﬀer 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 deﬁnition).
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 successor features 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 mixtureof-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).
Our architecture, MSFA, is the ﬁrst to integrate MRL with SF&GPI.
This integration enables combining task solutions in novel environment conﬁgurations.
Generalizing to novel environment conﬁgurations is typically studied with procedurally generated environments (Cobbe et al., 2020; Chevalier-Boisvert et al., 2018; Samvelyan et al., 2021).
Some approaches to this problem leverage techniques from supervised learning such as regularization, curriculum strategies, hyper-parameter tuning, and using self-supervised objectives (Cobbe et al., 2020;
Laskin et al., 2020; Igl et al., 2019;
Anand et al., 2021).
MSFA is most closely related to Recurrent Independent Mechanisms (RIMS) (Goyal et al., 2019) and its follow-up, Feature-Attending Recurrent Modules (FARM) (Carvalho et al., 2021a).
RIMs showed that leveraging modules to learn a state function improved out-ofdistribtion generalization. FARM showed that a modiﬁed attention mechanism led to strong generalization improvements with RL.
MSFA diﬀers 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 conﬁgurations.
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 train tasks M train =
{M i } n train i =1 , sampled from a training distribution p train ( M ) .
During testing, the agent is evaluated on n test tasks, {M i } n test i =1 , sampled from a testing distribution p test ( M ) .
Each task M i is speciﬁed as Partially Observable Markov Decision Process (POMDP), M i =
(cid:104)S e , A , X , R, p, f x (cid:105) .
Here, S e , A and X are the environment state, action, and observation spaces.
p ( ·| s et , a t ) speciﬁes the next-state distribution based on taking action a t in state s et , and f x ( s et ) 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 wt = R ( s et , a t , s et +1 , w ) is the reward obtained for transition ( s et , a t , s et +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 θ ( x t , s t − 1 , a t − 1 ) .
Given this learned state, we want to obtain a behavioral policy π ( s t ) that best maximises the expected reward it will obtain when taking an action a t at a state s t : Q π,wt
[Equation]
Transfer with SF&GPI.
In order to leverage successor features (SFs)
(Barreto et al., 2017), one assumes an agent has access to state features known as “cumulants”, φ t = φ ( s t , a t , s t +1 ) .
Given a behavioral policy π ( a | s ) , SFs are a type of value function that use cumulants as pseudo-rewards:
[Equation]
If reward is (approximately) r wt = φ (cid:62) t w , then action-values can be decomposed as Q π,wt = ψ πt (cid:62) w .
This is interesting because it provides an easy way to reuse task-agnostic features ψ πt for new tasks.
We focus on transfer to test tasks w test that are in the span of the training tasks, w test =
(cid:80)
i α i w i . We can re-use the SFs we’ve learned from training tasks M train for transfer with GPI.
Assume we have learned (potentially optimal) policies { π i } n train i =1 and their corresponding SFs { ψ π i ( s, a ) } n train i =1 .
Given a test task w 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:
[Equation]
This is useful because the GPI theorem states that π will perform as well as all of the training policies, i.e. that Q π,w test ( s, a ) ≥ max i Q π i ,w i ( s, a
[Equation] ∀ ( s, a ) ∈ ( S × 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 deﬁned that value function.
However, SF&GPI relies on combining a ﬁxed 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 θ , Q θ ( s, a, w ) .
If Q θ is smooth with respect to w , then Q θ should generalize to test tasks nearby to train tasks in task space.
Borsa et al.
(2019) showed that one could combine the beneﬁts of both with “Universal Successor Feature Approximators”.
Since rewards r w , and therefore task vectors w , reference deterministic task policies π w , one can parameterize successor feature approximators with task-vectors ˜ ψ π w = ˜ ψ w ≈ ψ θ ( 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.
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 θ k } nk =1 that evolve with independent parameters θ k and have sparse inter-module information ﬂow.
MSFA then produces modular cumulants { ˜ φ ( k ) t } nk =1 and SFs { ˜ ψ π,kt } nk =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. 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 )
= (cid:80) k
R k ( s t , a t , s t +1 )
= (cid:80) 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 φ ( k ) ∈ R :
R ( k ) t = [Equation]
[Equation]
· 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
[Equation]
where we now have modular SFs { ψ π,k ( s, a ) } nk =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. Architecture
We learn a set of modules with states S t = { s ( k ) t } nk =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 θ ( s ( k ) t − 1 , S t − 1 ) .
Following prior work (Santoro et al., 2018; Goyal et al., 2019; Carvalho et al., 2021a), we have A θ 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 θ 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.
[Equation]
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:
[Equation]
We highlight that cumulants share parameters but diﬀer 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-speciﬁc cumulants and SFs to form the ﬁnal outputs:
˜ φ t = (cid:104) ˜
φ (1) t , . .
˜ φ ( n ) t (cid:105) and ˜ ψ wt = ψ θ ( s t , a t , w ) =
(cid:104)
˜ ψ w, 1 t , . . . ,
˜ ψ w,nt (cid:105) .
Note that cumulants update with the module-state from the next time-step.
This is ﬁne since cumulants are only used during learning.
4. Behavior
During training , actions are selected in proportion to Q-values computed using task SFs as π ( s t , w ) ∝ ˜ Q ( s t , a, w ) =
ψ θ ( s t , a, w ) (cid:62) w .
In practice we use an epsilon-greedy policy, though one can use other choices.
During testing , we compute policies with GPI as π ( s t , w test ) ∈ arg max a max z ∈ M train { ψ θ ( s t , a, z ) (cid:62) w test } , where M train are train task vectors.
4. Learning Algorithm
MSFA relies on three losses.
The ﬁrst loss, L Q , is a standard Q-learning loss, which MSFA uses to learn optimal policies for the training tasks.
The main diﬀerence here is that MSFA uses a particular parameterization of the Q-function Q π w ,w ( s, a )
=
ψ π w ( s, a ) (cid:62) w .
The second loss, L ψ , 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 wt = φ (cid:62) t w .
For this, we again apply standard Q-learning but using SFs as value functions and cumulants as pseudo-rewards.
The ﬁnal loss, L φ is a loss for learning cumulants that grounds them in the environment reward signal.
The losses are summarised as follows
[Equation]
where a (cid:48) =
arg max a ψ θ ( s t +1 , a, w ) (cid:62) w .
Selecting the next action via the combination of all modules ensures they individually convergence to optimal values (Russell & Zimdars, 2003).
The ﬁnal loss is L = α Q L Q + α ψ L ψ + α φ L φ .
We found α ψ = α φ = 1 .
and α Q = 0 .
5 to generally work well.
Experiments
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 conﬁgurations.
Baselines.
(1) Universal Value Function Approximator (UVFA)
(Schaul et al., 2015), which takes the task as input: Q θ ( s, w ) .
This comparison shows shows the transfer beneﬁts of SF&GPI.
(2) UVFA with Feature-Attending Recurrent Modules (UVFA+FARM) instead takes state-factors as input Q θ ( s (1) , . .
, s ( n ) , w ) .
Each state-factor s ( k ) is the output of a FARM module.
This comparison allows us to test the beneﬁts of leveraging modules for learning value functions. (3) Modular Value Function Approximator(MVFA) is an adaptation of (Haarnoja et al., 2018) where modules learn individual Q-values Q ( i ) θ ( s ( i ) , w ( i ) ) . This comparison shows us to test the beneﬁts 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- φ , 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 φ , ψ , 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 (Hoﬀman et al., 2020) for reinforcement learning.
5. Combining object navigation task knowledge
MSFA learns modular functions for computing φ and estimating ψ .
We hypothesize that this facilitates learning cumulants that respond to diﬀerent aspects of the environment (e.g. to diﬀerent object categories).
This leads to the following research questions.
R1 :
Can we recover prior generalization results that relied on hand-designed cumulants for diﬀerent object categories?
R2 :
How important is it to learn modular functions for φ and ψ ?
R3 :
Without GPI, do learning modular functions for φ and ψ still aid in generalization?
Setup.
We implement a simpliﬁed 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 ∈ R 4 with training tasks being the standard unit vectors.
For example, [0 , 1 , 0 , 0] speciﬁes the agent must obtain objects of category 2 .
Generalization tasks are linear combinations of training tasks.
For example, [ − 1 , 1 , − 1 , 1] speciﬁes the agent must obtain categories 2 and 4 while avoiding categories 1 and 3 .
Despite the apparent simplicity of this environment, it presents several challenges. First, learning φ is challenging because rewarding transition are rare subsets of all transitions leading to a data imbalance.
Borsa et al.
(2019) showed that USFA could generalize with hand-designed cumulants that described whether an object was picked up.
Second, test tasks are out-of-distribution and involve behaviors the agent has never learned such as avoiding objects. We expand on these challenges in the Appendix.
R1: MSFA is competitive with USFA, which uses oracle φ .
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- φ does slightly worse than MSFA.
However, with more objects to avoid, all methods except MSFA (including USFA-Learned- φ ) degrades comparison
R2:
Learning modular φ and ψ functions is critical for generalization.
Learning an entangled function corresponds to learning a monolithic function for ψ or φ where we concatenate module-states, e.g. ˜ ψ wt = ψ θ ( s (1) t , . . .
, s ( n ) t , a, w ) .
Modular functions correspond to equation 5.
Figure 6 (a) shows that without modular functions for φ and ψ , performance severely degrades.
This also highlights that a naive combination of USFA+FARM—with entangled functions for φ and ψ —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: π ( s ) =
arg max a ψ θ ( s t , a, w ) (cid:62) w .
Figure(b) shows that GPI is critical for generalization with USFA as expected.
USFA-Learned- φ beneﬁts less from GPI (presumably because of challenges in learning φ ).
Interestingly, MSFA can generalize relatively well without GPI, sometimes doing better than USFA without GPI.
5. Combining object navigation task knowledge with novel appearances and environment configurations
Beyond generalizing to combinations of tasks, agents will need to generalize to diﬀerent 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.
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.
5. Combining knowledge of heterogeneous tasks
Aside from navigation tasks, researchers and practitioners will be interested in generally combining solutions to diﬀerent tasks.
R5 :
Can MSFA enable combining solutions to heterogeneous tasks?
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 conﬁgurations of objects and (b) to the combination of all three training conditions.
: MSFA enables combining solutions to heterogeneous tasks
shows that SF-based methods best combine solutions to diﬀerent 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%.
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 ﬁrst 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).
Finally, we showed that MSFA can combine solutions of heterogeneous tasks (§5.3).
Our ablations show that learning modular cumulants and SFs is critical for generalization with GPI.
Interestingly, without GPI, learning modular cumulants and SFs improve generalization over monolithic SF architectures (Figure 6).
We compared MSFA to two baselines: (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.
A Derivation of Modular Successor Features
[Equation]
[Equation]
B Inter-module attention and gating
At each time-step t , each module updates with both observation features z t
= f θ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 modulestate 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 ﬂexible 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
[Equation]
Keys and values are computed as K t =
W key [ s 1 t − 1 ; . . .
; s n t − 1 ; 0] ∈ R n +1 × d q and V t =
W value [ s 1 t − 1 ; . . .
; s n t − 1 ; 0] ∈ R n +1 × 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):
[Equation]
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:
[Equation]
MSFA then updates its module-stateas follows:
[Equation]
C Learning and Generalization Challenges
Our experimental setup follows the same setup as Barreto et al.
For n tasks with reward functions { r i } ni =1 , they showed that discovered features which predict rewards, i.e. [˜ φ 1 , . . . ,
˜ φ n ] =
[˜ r 1 , . .
˜ r n ] , were eﬀective features for transfer with SF&GPI.
In their experimental setup, training tasks were one-hot vectors, so regressing || r − ˜ φ (cid:62) w || eﬀectively led φ i to be a feature that predicted reward r i .
We use this same setup, where now individual modules are used for both ˜ φ ( i ) and ˜ ψ
( i ) .
While this is a simple method for learning features, it some diﬃculties.
First, there is an imbalance of rewarding vs. non-rewarding transitions, which leads to challenges in learning φ .
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 9).
However, at test time, the agent must perform tasks that are are far in “task-space” from the training distribution.
For example, w 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 diﬀerent from training reward functions.
While both USFA-Learned- φ 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 ˜ r test =
˜
φ (cid:62) w test .
Reward prediction error on generalization tasks .
To investigate this, we collectepisodes of each generalization task and compute the reward prediction error of MSFA and
USFA-Learned- φ for each task.
We present the results in Table 1.
[Table]
D Learning modular functions
Our goal is to learn functions for producing cumulants φ and SFs ψ 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:
[Equation]

[Equation]
By "entangled," we mean that a cumulant φ i is (potentially) a function of all module-state information φ
i = (cid:80) j A ij s j .
By a "modular function", we simply meant that we were learning a function like the following
[Equation] In our setting, we had A 1 = . . . = 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 φ ( i ) and ψ ( i ) is an eﬀective way to learn { φ i } i that promote zero-shot composition of task knowledge with the SF&GPI framework.
E Additional analysis
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 diﬀerent methods.
We show the ﬁnal episode return for each method along the when an
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- φ equally collects all objects regardless of task, except for when r ≤ − 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 eﬀectively as MSFA or USFA.
F Full results
Sometimes (e.g. [1 , 1 , 1 , 1] and [ − 1 , 1 , − 1 , − 1] ), MSFA outperforms USFA. Next we compare MSFA to USFA-Learned- φ . We see that MSFA performs as well as USFA-Learned- φ or better in all settings except [ − 1 , 1 , 0 , 1] .
Interestingly, MSFA without GPI outperforms
USFA-Learned- φ with GPI in some generalization settings (e.g. [1 , 1 , 0 , 0] and [ − 1 , 1 , − 1 , 1] ).
When combinations of objects must be collected, MSFA outperforms baselines by a large margin
We shown hyperparameters in Tables 2 and 3. Tables 2 describes hyperparameters that were shared across algorithms and Tables 3 describes hyperparameters that were diﬀerent across algorithms.
We developed the algorithms using the Hoﬀman et al.
(2020) reinforcement learning codebase and unless stated otherwise using default “R2D2” (Kapturowski et al., 2018) conﬁg values in the codebase.
We describe our development and hyperparameter search in more detail below.
G. Development and hyperparameter search
The ﬁrst 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 (ChevalierBoisvert 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.
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 USFALearned- φ and MSFA.
In this setting, we are learning Q-values, successor features, and predicting rewards from cumulants.
This leads to the following coeﬃcients for losses: α Q , α ψ , and α φ respectivey.
We ﬁxed α ψ = 1 and searched α φ ∈ [ . 01 , .
1 , 1 , 10 , 100] and α Q ∈ [ . 01 , .
5 , 1 . 0] for both USFA-Learned- φ 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 ﬁrst 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 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 α φ and α Q for both MSFA and USFA-Learned- φ .
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 φ and ψ , 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- φ and MSFA but was detrimental to UVFA-based methods.
H 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.
H. 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 [Table]
[Table] [Table] [Table] 12 object types { O 1 . . .
O 12 } into 4 categories: C 1 =
{ O 1 , O 2 , O 3 } , . .
, C
= { O 10 , O 11 , O 12 } , where each category provided postive reward for picking up any of its 3 category types.
H. Minihack
We did not make any changes to the environment.
Train tasks were: “MiniHack-RoomMonster-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 × 5 for the “small room” experiments and Y = 15 × 15 for the “large room” experiments.