Original: C OMPOSITIONAL T ASK G ENERALIZATION WITH D ISCOVERED S UCCESSOR F EATURE M ODULES
Revised: C OMPOSING T ASK K NOWLEDGE WITH M ODULAR S UCCESSOR F EATURE A PPROXIMATORS

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: 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.
Revised: 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.

Original: SF&GPI works by having an agent learn predictive representations (SFs) that can be combined for transfer to new tasks with GPI.
Revised: SF&GPI works by having an agent learn predictive representations (SFs) that can be combined for transfer to new tasks with GPI.

Original: However, to be effective this approach requires state features that are useful to predict, and these state-features are typically hand-designed.
Revised: However, to be effective this approach requires state features that are useful to predict, and these state-features are typically hand-designed.

Original: In this work, we present a novel neural network architecture, “Modular Successor Feature
Revised: In this work, we present a novel neural network architecture, “Modular Successor Feature

Original: Approximators” (MSFA), where modules both discover what is useful to predict, and learn their own predictive representations.
Revised: Approximators” (MSFA), where modules both discover what is useful to predict, and learn their own predictive representations.

Original: We show that MSFA is able to better generalize compared to baseline architectures for learning SFs anda modular network that discovers factored state representations.
Revised: We show that MSFA is able to better generalize compared to baseline architectures for learning SFs and modular architectues for learning state representations.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Consider a household robot that needs to learn tasks including picking up dirty dishes and cleaning up spills.
Revised: Consider a household robot that needs to learn tasks including picking up dirty dishes and cleaning up spills.

Original: Now consider that the robot is deployed and encounters a table with both a spill and set of dirty dishes.
Revised: Now consider that the robot is deployed and encounters a table with both a spill and set of dirty dishes.

Original: Ideally this robot can combine its training behaviors to both clean up the spill and pickup the dirty dishes.
Revised: Ideally this robot can combine its training behaviors to both clean up the spill and pickup the dirty dishes.

Original: We study this aspect of generalization: combining knowledge from multiple tasks.
Revised: We study this aspect of generalization: combining knowledge from multiple tasks.

Original: 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.
Revised: 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.

Original: This challenge is exacerbated when an agent also needs to generalize to novel appearances and environment conﬁgurations.
Revised: This challenge is exacerbated when an agent also needs to generalize to novel appearances and environment conﬁgurations.

Original: 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.
Revised: 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.

Original: Successor features (SFs) and Generalized Policy Improvement (GPI) provide a mechanism to combine knowledge from multiple training tasks (Barreto et al., 2017; 2020).
Revised: Successor features (SFs) and Generalized Policy Improvement (GPI) provide a mechanism to combine knowledge from multiple training tasks (Barreto et al., 2017; 2020).

Original: SFs are predictive representations that estimate how much state-features (known as “cumulants”) will be experienced given a behavior.
Revised: SFs are predictive representations that estimate how much state-features (known as “cumulants”) will be experienced given a behavior.

Original: By assuming that reward has a linear relationship between cumulants and a task vector, an agent can efﬁciently compute how much reward it can expect to obtain from a given behavior.
Revised: By assuming that reward has a linear relationship between cumulants and a task vector, an agent can efﬁciently compute how much reward it can expect to obtain from a given behavior.

Original: 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).
Revised: 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).

Original: However, SF&GPI commonly assume hand-designed cumulants and don’t have a mechanism for generalizing to novel environment conﬁgurations.
Revised: However, SF&GPI commonly assume hand-designed cumulants and don’t have a mechanism for generalizing to novel environment conﬁgurations.

Original: Modular architectures are a promising method for generalizing to distributions outside of the training distribution (Goyal et al., 2019; Madan et al., 2021).
Revised: Modular architectures are a promising method for generalizing to distributions outside of the training distribution (Goyal et al., 2019; Madan et al., 2021).

Original: Recently, Carvalho et al.
Revised: Recently, Carvalho et al.

Original: (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).
Revised: (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).

Original: In this work, we hypothesize that modules can further be leveraged to discover state-features that are useful to predict.
Revised: In this work, we hypothesize that modules can further be leveraged to discover state-features that are useful to predict.

Original: We present “Modular Successor Feature Approximators” (MSFA), a novel neural network for discovering, composing, and transferring predictive knowledge and behavior via SF&GPI.
Revised: We present “Modular Successor Feature Approximators” (MSFA), a novel neural network for discovering, composing, and transferring predictive knowledge and behavior via SF&GPI.

Original: MSFA is composed of a set of modules, which each learn their own state-features and corresponding predictive representations.
Revised: MSFA is composed of a set of modules, which each learn their own state-features and corresponding predictive representations (SFs).

Original: We show that an architectural bias of having state-features be produced by their own modules is useful inductive bias for reward-driven discovery of useful features to use with SF&GPI.

Original: 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.
Revised: 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.

Original: We exemplify this with a simple state-feature discovery method presented in Barreto et al.
Revised: We exemplify this with a simple state-feature discovery method presented in Barreto et al.

Original: (2018) where the dot-product between state-features and a task vector is regressed to environment reward.
Revised: (2018) where the dot-product between state-features and a task vector is regressed to environment reward.

Original: This method enabled transfer with SF&GPI in a continual learning setting but had limited success in the zero-shot transfer settings we study.
Revised: This method enabled transfer with SF&GPI in a continual learning setting but had limited success in the zero-shot transfer settings we study.

Original: 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).
Revised: While there are other methods for state-feature discovery, they add training complexity with mutual information objectives (Hansen et al., 2019) or metagradients (Veeriah et al., 2019).

Original: With MSFA, by adding only an architectural bias for modularity , we discover statefeatures that (1) support zero-shot transfer competitive with hand-designed features, and (2) enable zero-shot transfer in visually diverse, procedurally generated environments.
Revised: 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.

Original: We are hopeful that our architectural bias can be leveraged with other discovery methods in future work.
Revised: We are hopeful that our architectural bias can be leveraged with other discovery methods in future work.

Original:  R ELATED W ORK ON G ENERALIZATION IN RL
Revised:  R ELATED W ORK ON G ENERALIZATION IN RL

Original: Hierarchical RL (HRL) is one dominant approach for tackling compositional generalization.
Revised: Hierarchical RL (HRL) is one dominant approach for combining task knowledge.

Original: The basic idea is that one can sequentially compose policies in time by having a “meta-policy” that sequentially activates “low-level” policies for protracted periods of time.
Revised: 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.

Original: By leveraging handdesigned 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).
Revised: By leveraging hand-designed or pretrained 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).

Original: We differ in that we focus on combining policies concurrently in time as opposed to sequentially in time.
Revised: We differ in that we focus on combining policies concurrently in time as opposed to sequentially in time.

Original: To do so, we develop a modular neural network for the SF&GPI framework.
Revised: To do so, we develop a modular neural network for the SF&GPI framework.

Original: 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).
Revised: 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).

Original: By combining them with Generalized Policy
Revised: By combining them with Generalized Policy

Original: 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).
Revised: 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).

Original: However, these works tend to require hand-designed cumulants which are cumbersome to design for every new environment.
Revised: However, these works tend to require hand-designed cumulants which are cumbersome to design for every new environment.

Original: In our work, we integrate successor features with Modular RL to facilitate reward-driven discovery of cumulants and improve successor feature learning.
Revised: In our work, we integrate successor features with Modular RL to facilitate reward-driven discovery of cumulants and improve successor feature learning.

Original: Modular RL (MRL) (Russell & Zimdars, 2003) is a framework that achieves compositional generalization through the composition of value functions.
Revised: Modular RL (MRL) (Russell & Zimdars, 2003) is a framework for generalization by combining value functions.

Original: Early work dates back to (Singh, 1992), who had a mixture-of-experts system select between separately trained value functions.
Revised: Early work dates back to (Singh, 1992), who had a mixture-of-experts system select between separately trained value functions.

Original: 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).
Revised: 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).

Original: Our architecture, MSFA, is the ﬁrst to integrate MRL with SF&GPI.
Revised: Our architecture, MSFA, is the ﬁrst to integrate MRL with SF&GPI.

Original: This integration enables composing task solutions in novel environment conﬁgurations.
Revised: This integration enables combining task solutions in novel environment conﬁgurations.

Original: 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).
Revised: 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).

Original: 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;
Revised: 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;

Original: Laskin et al., 2020; Igl et al., 2019;
Revised: Laskin et al., 2020; Igl et al., 2019;

Original: Anand et al., 2021).
Revised: Anand et al., 2021).

Original: Most relevant tous, Carvalho et al.
Revised: MSFA is most closely related to

Original: (2021a) recently showed that  amodular architecture for learning a factored state representation, Feature-Attending Recurrent Modules (FARM), could improve generalization in procedurally generated object-centric environments.
Revised: 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-of-distribtion generalization.

Original: FARM exploited modularity to combine state representations.
Revised: FARM showed that a modiﬁed attention mechanism led to strong generalization improvements with RL.

Original: We instead exploit modularity to combine value functions and show that this improves an agent’s ability to both combine task solutions and generalize to novel environment conﬁgurations.
Revised: 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 conﬁgurations.

Original:  P ROBLEM S ETTING AND B ACKGROUND
Revised:  P ROBLEM S ETTING AND B ACKGROUND

Original: We study a reinforcement learning agent’s ability to transfer knowledge between tasks in an environment.
Revised: We study a reinforcement learning agent’s ability to transfer knowledge between tasks in an environment.

Original: During training, the experiences n train tasks M train =
Revised: During training, the experiences n train tasks M train =

Original: {M i } n train i =1 , sampled from a training distribution p train ( M ) .
Revised: {M i } n train i =1 , sampled from a training distribution p train ( M ) .

Original: During testing, the agent is evaluated on n test tasks, {M i } n test i =1 , sampled from a testing distribution p test ( M ) .
Revised: During testing, the agent is evaluated on n test tasks, {M i } n test i =1 , sampled from a testing distribution p test ( M ) .

Original: Each task M i is speciﬁed as Partially Observable Markov
Revised: Each task M i is speciﬁed as Partially Observable Markov

Original: Decision Process (POMDP), M i =
Revised: Decision Process (POMDP), M i =

Original: (cid:104)S  , A , X , R, p, f x (cid:105) .
Revised: (cid:104)S e , A , X , R, p, f x (cid:105) .

Original: Here, S  , A and X are the  state, action, and observation spaces.
Revised: Here, S e , A and X are the environment state, action, and observation spaces.

Original: p ( ·| s t , a t ) speciﬁes the next-state distribution based on taking action a t in state s t , and f x ( s t ) maps the underlying environment state to an observation x t .
Revised: 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 .

Original: We focus on tasks where rewards are parameterized by a task vector w , i.e. r wt = R ( s t , a t , s t +1 , w ) is the reward obtained for transition ( s t , a t , s t +1 ) given task vector w .
Revised: 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 .

Revised: Since this is a POMDP, we need to learn a state function that maps histories to agent state representations.

Revised: We do so with a recurrent function: s t =

Revised: s θ ( x t , s t − 1 , a t − 1 ) .

Original: We want to obtain a behavioral policy π ( s  ) that best maximises the expected reward it will obtain when taking an action a t at a state s t :
Revised: 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 :

Original: Q π,wt 
Revised: Q π,wt = Q

Original: [Equation]
Revised: [Equation]

Original: drawer” and “open fridge”.
Revised: ,w ( s [Equation] = E

Revised: [Equation]

Revised: [ (cid:80)

Revised: [Equation]

Original:  Each task leads the agent to experience different “cumulants” φ —e.g., “fork” during “open drawer” and “apple” during “open fridge”.
Revised: wt ] . 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”.

Original: The agent leverages different modulesto learn these φ  andto estimate SFs, ψ .
Revised: We can leverage MSFA to have different modules learn different “cumulants”, φ , and SFs, ψ .

Original: For example,across tasks, module 1 ( θ 1 ) can estimate SFs for apple cumulants.
Revised: For example, module 1 ( θ 1 ) can estimate SFs for apple cumulants.

Original: Module SFs are combined to form the SF for a policy.
Revised: Module SFs are combined to form the SF for a policy.

Original: 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.
Revised: 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.

Original: The highest Q-value is then used to select actions.
Revised: The highest Q-value is then used to select actions.

Original: Transfer with SF&GPI.
Revised: Transfer with SF&GPI.

Original: In order to leverage successor features (SFs)
Revised: In order to leverage successor features (SFs)

Original: (Barreto et al., 2017), one assumes an agent has access to state features known as “cumulants”, φ t = φ ( s t , a t , s t +1 ) .
Revised: (Barreto et al., 2017), one assumes an agent has access to state features known as “cumulants”, φ t = φ ( s t , a t , s t +1 ) .

Original: Given a behavioral policy π ( a | s ) , SFs are a type of value function that use cumulants as pseudo-rewards:
Revised: Given a behavioral policy π ( a | s ) , SFs are a type of value function that use cumulants as pseudo-rewards:

Original: [Equation]
Revised: [Equation]

Original: If reward is (approximately) r wt = φ (cid:62) t w , then action-values can be decomposed as Q π,wt [Equation] w .
Revised: If reward is (approximately) r wt = φ (cid:62) t w , then action-values can be decomposed as Q π,wt = ψ πt (cid:62) w .

Original: This is interesting because it provides an easy way to reuse task-agnostic features ψ πt for new tasks.
Revised: This is interesting because it provides an easy way to reuse task-agnostic features ψ πt for new tasks.

Revised: We focus on transfer to test tasks w test that are in the span of the training tasks, w test =

Revised: (cid:80)

Original:  We can re-use the SFs we’ve learned from training tasks M train for transfer with GPI.
Revised: i α i w i . We can re-use the SFs we’ve learned from training tasks M train for transfer with GPI.

Original: Assume we have learned (potentially optimal) policies { π i } n train i =1 and their corresponding SFs { ψ π i ( s, a ) } n train i =1 .
Revised: Assume we have learned (potentially optimal) policies { π i } n train i =1 and their corresponding SFs { ψ π i ( s, a ) } n train i =1 .

Original: 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.
Revised: 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.

Original: This operation is summarized as follows:
Revised: This operation is summarized as follows:

Original: [Equation]
Revised: [Equation]

Original: 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 ) ∀ ( s, a ) ∈ ( S × A )
Revised: 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] ∀ [Equation]

Revised: s,

Original:  (Barreto et al., 2017).
Revised: [Equation] ∈ ( S × [Equation] (Barreto et al., 2017).

Original: 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.
Revised: 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.

Original: However, SF&GPI relies on combining a ﬁxed set of SFs.
Revised: However, SF&GPI relies on combining a ﬁxed set of SFs.

Original: 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 ) .
Revised: 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 ) .

Original: If Q θ is smooth with respect to w , then Q θ should generalize to test tasks nearby to train tasks in task space.
Revised: If Q θ is smooth with respect to w , then Q θ should generalize to test tasks nearby to train tasks in task space.

Original: Borsa et al.
Revised: Borsa et al.

Original: (2019) showed that one could combine the beneﬁts of both with “Universal Successor Feature Approximators”.
Revised: (2019) showed that one could combine the beneﬁts of both with “Universal Successor Feature Approximators”.

Original: Since rewards r w , and therefore task vectors w , reference deterministic task policies π w , one can parameterize successor feature approximators with task-vectors ˜ ψ π w = ˜
Revised: Since rewards r w , and therefore task vectors w , reference deterministic task policies π w , one can parameterize successor feature approximators with task-vectors ˜ ψ π w = ˜

Original: ψ w ≈ ψ θ ( s, a, w ) .
Revised: ψ w ≈ ψ θ ( s, a, w ) .

Original: However,
Revised: However,

Original: USFA assumed hand-designed cumulants.
Revised: USFA assumed hand-designed cumulants.

Original: We introduce an architecture for reward-driven discovery of cumulants and improved function approximation of universal successor features.
Revised: We introduce an architecture for reward-driven discovery of cumulants and improved function approximation of universal successor features.

Original: Modules then share information withan attention mechanism.
Revised: Modules then share information with mechanism.

Original: We hypothesize that isolated module computations facilitate learning cumulants that suppot generalization with GPI.
Revised: We hypothesize that isolated module computations facilitate learning cumulants that suppot generalization with GPI.

Original:  M ODULAR S UCCESSOR F EATURE A PPROXIMATORS
Revised:  M ODULAR S UCCESSOR F EATURE A PPROXIMATORS

Original: We propose a new architecture Modular Successor Feature Approximators (MSFA) for approximating
Revised: We propose a new architecture Modular Successor Feature Approximators (MSFA) for approximating

Original: SFs, shown in Figure 3.
Revised: SFs, shown in Figure 3.

Original: Our hypothesis is that learningdisentangled cumulants and SFs will improve compositional generalization with SF&GPI.
Revised: Our hypothesis is that learning cumulants and SFs with modules improves zero-shot composition of task knowledge with SF&GPI.

Original: We accomplish this by learning ndisentangled state modules { s θ k } n k =1 that evolve with independent parameters θ k and have sparse inter-module information ﬂow.
Revised: 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.

Original: We produce disentangled cumulants { ˜ φ ( k ) t } nk =1 and SFs { ˜ ψ π,kt } nk =1 by having their computations depend only on individual module-states.
Revised: MSFA then produces modular cumulants { ˜ φ ( k ) t } nk =1 and SFs { ˜ ψ π,kt } nk =1 by having their computations depend only on individual module-states.

Original: 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.
Revised: 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.

Original: This is in contrast to prior work, which learns a single monolithic prediction module for computing cumulants and SFs.
Revised: This is in contrast to prior work, which learns a single monolithic prediction module for computing cumulants and SFs (see Figure 3).

Original: The rest of section is structured as follows.
Revised: The rest of section is structured as follows.

Original: In section §4.1, we derive Modular Successor Feature
Revised: In section §4.1, we derive Modular Successor Feature

Original: Learning within the Modular RL framework.
Revised: Learning within the Modular RL framework.

Original: We then describe our architecture, MSFA, for learning modular successor features in §4.2.
Revised: We then describe our architecture, MSFA, for learning modular successor features in §4.2.

Original: In §4.3, we describe how to generate behavior with MSFA.
Revised: In §4.3, we describe how to generate behavior with MSFA.

Original: Finally, we describe the learning algorithm for MSFA in section §4.4.
Revised: Finally, we describe the learning algorithm for MSFA in section §4.4.

Original: 4. M ODULAR S UCCESSOR F EATURE L EARNING
Revised: 4. M ODULAR S UCCESSOR F EATURE L EARNING

Original: Following the Modular RL framework (Russell & Zimdars, 2003), we assume that reward has an additive structure R ( s t , a t , s t +1 )
Revised: Following the Modular RL framework (Russell & Zimdars, 2003), we assume that reward has an additive structure R ( s t , a t , s t +1 )

Original: = (cid:80) k
Revised: = (cid:80) k

Original: R k ( s t , a t , s t +1 )
Revised: R k ( s t , a t , s t +1 )

Original: = (cid:80) k R ( k ) t , where R ( k ) t is the reward of the k -th module.
Revised: = (cid:80) k R ( k ) t , where R ( k ) t is the reward of the k -th module.

Original: We enforce that every module decomposes reward into an inner-product between its own task-description w ( k ) and task-agnostic cumulants φ
Revised: We enforce that every module decomposes reward into an inner-product between its own task-description w ( k ) and task-agnostic cumulants φ

Original: ( k ) ∈ R : R ( k ) t =φ
Revised: ( k ) ∈ R : R ( k ) t =

Original: [Equation]
Revised: [Equation]

Original: · w ( k ) .
Revised: · w ( k ) .

Original: Here, we simply break up the task vector into n pieces so individual modules are responsible for subsets of the task vector.
Revised: Here, we simply break up the task vector into n pieces so individual modules are responsible for subsets of the task vector.

Original: This allows us to decompose the action-value function as
Revised: This allows us to decompose the action-value function as

Original: [Equation]
Revised: [Equation]

Original: where we now have modular SFs { ψ π,k ( s, a ) } nk =1 (see the Appendix for a derivation).
Revised: where we now have modular SFs { ψ π,k ( s, a ) } nk =1 (see the Appendix for a derivation).

Original: Rather than hand-designing modules or cumulants, we aim to discover them from the environment reward signal.
Revised: Rather than hand-designing modules or cumulants, we aim to discover them from the environment reward signal.

Original: 4. A RCHITECTURE
Revised: 4. A RCHITECTURE

Original: We learn a set of modules with states S t = { s ( k ) t } nk =1 .
Revised: We learn a set of modules with states S t = { s ( k ) t } nk =1 .

Original: 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  , S t − 1 ) .
Revised: 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 ) .

Original: 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.
Revised: 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.

Original: Since A θ is not the main contribution of this paper, we describe these computations in more detail with our notation in the Appendix.
Revised: Since A θ is not the main contribution of this paper, we describe these computations in more detail with our notation in the Appendix.

Original: We summarize the high-level update below.
Revised: We summarize the high-level update below.

Original: [Equation]
Revised: [Equation]

Original: We learn disentangled cumulants and SFs by having sets of cumulants and SFs depend on individual module-states.
Revised: We learn modular cumulants and SFs by having sets of cumulants and SFs depend on individual module-states.

Original: Module cumulants depend on the module-state from the current and next time-step.
Revised: Module cumulants depend on the module-state from the current and next time-step.

Original: Module SFs depend on the current module-state and on their subset of the task description.
Revised: Module SFs depend on the current module-state and on their subset of the task description.

Original: We summarize this below:
Revised: We summarize this below:

Original: [Equation]
Revised: [Equation]

Original: We highlight that cumulants share parameters but differ in their input.
Revised: We highlight that cumulants share parameters but differ in their input.

Original: 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).
Revised: 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).

Original: We show evidence for this hypothesis in Figure 6.
Revised: We show evidence for this hypothesis in Figure 6.

Original: We concatenate module-speciﬁc cumulants and SFs to form the ﬁnal outputs:
Revised: We concatenate module-speciﬁc cumulants and SFs to form the ﬁnal outputs:

Original: ˜ φ t = (cid:104) ˜
Revised: ˜ φ t = (cid:104) ˜

Original: φ (1) t , . .
Revised: φ (1) t , . .

Original: ˜ φ ( n ) t (cid:105) and ˜ ψ wt = ψ θ ( s t , a t , w ) =
Revised: ˜ φ ( n ) t (cid:105) and ˜ ψ wt = ψ θ ( s t , a t , w ) =

Original: (cid:104)
Revised: (cid:104)

Original: ˜ ψ w, 1 t , . .. ,
Revised: ˜ ψ w, 1 t , . .

Original:  ˜ ψ w,nt (cid:105) .
Revised: . , ˜ ψ w,nt (cid:105) .

Original: Note that cumulants update with the module-state from the next time-step.
Revised: Note that cumulants update with the module-state from the next time- step.

Revised: This is

Revised: ﬁne

Original: This is ﬁne since cumulants are only usedduring learning.
Revised: since cumulants are only used

Revised: during

Revised: learning.

Original: 4. B EHAVIOR
Revised: 4. B EHAVIOR

Original: During training , actions are selected in proportion to Q-values computed using task SFs as π ( s t , w ) ∝ ˜ Q ( s t , a, w ) =
Revised: During training , actions are selected in proportion to Q-values computed using task SFs as π ( s t , w ) ∝ ˜ Q ( s t , a, w ) =

Original: ψ θ ( s t , a, w ) (cid:62) w .
Revised: ψ θ ( s t , a, w ) (cid:62) w .

Original: In practice we use an epsilon-greedy policy, though one can use other choices.
Revised: In practice we use an epsilon-greedy policy, though one can use other choices.

Original: 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.
Revised: During testing , we compute policies with GPI as π ( s t , w test ) ∈ arg [Equation] { ψ θ [Equation]

Revised: s

Revised: [Equation]

Original:  4. L EARNING A LGORITHM
Revised: , a, z ) [Equation] w test } , where M train are train task vectors. 4. L EARNING A LGORITHM

Original: MSFA relies on three losses.
Revised: MSFA relies on three losses.

Original: The ﬁrst loss, L Q , is a standard Q-learning loss, which MSFA uses to learn optimal policies for the training tasks.
Revised: The ﬁrst loss, L Q , is a standard Q-learning loss, which MSFA uses to learn optimal policies for the training tasks.

Original: The main difference here is that MSFA uses a particular parameterization of the Q-function Q π w ,w ( s, a )
Revised: The main difference here is that MSFA uses a particular parameterization of the Q-function Q π w ,w ( s, a )

Original: =
Revised: =

Original: ψ π w ( s, a ) (cid:62) w .
Revised: ψ π w ( s, a ) (cid:62) w .

Original: 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 .
Revised: 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 .

Original: For this, we again apply standard Q-learning but using SFs as value functions and cumulants as pseudo-rewards.
Revised: For this, we again apply standard Q-learning but using SFs as value functions and cumulants as pseudo-rewards.

Original: The ﬁnal loss, L φ is a loss for learning cumulants that grounds them in the environment reward signal.
Revised: The ﬁnal loss, L φ is a loss for learning cumulants that grounds them in the environment reward signal.

Original: The losses are summarised as follows
Revised: The losses are summarised as follows

Original: [Equation]
Revised: [Equation]

Original:  where a (cid:48) =
Revised: ) where a (cid:48) =

Original: arg max a ψ θ ( s t +1 , a, w ) (cid:62) w .
Revised: arg max a ψ θ ( s t +1 , a, w ) (cid:62) w .

Original: Selecting the next action via the combination of all modules ensures they individually convergence to optimal values (Russell & Zimdars, 2003).
Revised: Selecting the next action via the combination of all modules ensures they individually convergence to optimal values (Russell & Zimdars, 2003).

Original: The ﬁnal loss is
Revised: The ﬁnal loss is

Original: L = α Q L Q + α ψ L ψ + α φ L φ .
Revised: L = [Equation] Q + [Equation]

Original:  We found α ψ = α φ = 1 . 

Revised: + [Equation] L φ . We found α ψ = α φ = 1 .

Original: and α Q = 0 .
Revised: and α Q = 0 .

Original: 5 to generally work well.
Revised: 5 to generally work well.

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original:  beyond combining object navigation skills.
Revised: 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.

Original: Here, there are three training settings: (1) avoiding teleportation traps, (2) avoiding monsters, and (3) partial visibility around the agent.
Revised: Here, there are three training settings: (1) avoiding teleportation traps, (2) avoiding monsters, and (3) partial visibility around the agent.

Original: We study generalization to the combination of these settings.
Revised: We study generalization to the combination of these settings.

Original: We study generalization when training behaviors must be combined concurrently in time in the presence of novel object appearances and layouts.
Revised: We study generalization when training behaviors must be combined concurrently in time in the presence of novel object appearances and layouts.

Original: The need to combine training behaviors tests how well MSFA can leverage SF&GPI.
Revised: The need to combine training behaviors tests how well MSFA can leverage SF&GPI.

Original: Generalization to novel object appearances and layouts tests how well MSFA’s modular construction supports generalization to novel environment conﬁgurations.
Revised: Generalization to novel object appearances and layouts tests how well MSFA’s modular construction supports generalization to novel environment conﬁgurations.

Original: Baselines.
Revised: Baselines.

Original: (1) Universal Value Function Approximator (UVFA)
Revised: (1) Universal Value Function Approximator (UVFA)

Original: (Schaul et al., 2015), which takes the task as input: Q θ ( s, w ) .
Revised: (Schaul et al., 2015), which takes the task as input: Q θ ( s, w ) .

Original: Comparing against UVFA enables us to measure the transfer beneﬁts of SF&GPI.
Revised: This comparison shows shows the transfer beneﬁts of SF&GPI.

Original: (2) UVFA with Feature-Attending Recurrent Modules (UVFA+FARM) instead takes state-factors as input Q θ ( s (1) , . . 
Revised: UVFA with Feature-Attending Recurrent Modules (UVFA+FARM) instead takes state-factors as input Q θ ( s (1) , . . .

Original: , s ( n ) , w ) .
Revised: , s ( n ) , w ) .

Original: Each state-factor s ( k ) is the output of a 
Revised: Each state-factor s ( k ) is the output of a FARM module.

Original: FARM module.
Revised: 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 ) ) .

Original: Comparing against UVFA+FARM enables us to probe the importance of having modules learn cumulants and successor features.
Revised: This comparison shows us to test the beneﬁts of leveraging modules for learning value functions in the form of SFs.

Original: (3) Universal Successor Function Approximator (USFA)
Revised: (4) Universal Successor Function Approximator (USFA)

Original: (Borsa et al., 2019) leverages a single monolithic function for successor features with handdesigned cumulants .
Revised: (Borsa et al., 2019) leverages a single monolithic function for successor features with hand-designed cumulants .

Original: USFA is an upper-bound baseline that allows us to test the quality of cumulants and successor features that MSFA learns.
Revised: USFA is an upper-bound baseline that allows us to test the quality of cumulants and successor features that MSFA learns.

Original: We also test a variant of USFA with learned cumulants, USFA-Learned- φ , which shows how the architecture degrades without oracle cumulants.
Revised: We also test a variant of USFA with learned cumulants, USFA-Learned- φ , which shows how the architecture degrades without oracle cumulants.

Original: Implementation.
Revised: Implementation.

Original: We implement the state modules of MSFA with FARM modules (Carvalho et al., 2021b).
Revised: We implement the state modules of MSFA with FARM modules (Carvalho et al., 2021b).

Original: For UVFA and USFA, we learn a state representation with an LSTM (Hochreiter &
Revised: For UVFA and USFA, we learn a state representation with an LSTM (Hochreiter &

Original: Schmidhuber, 1997).
Revised: Schmidhuber, 1997).

Original: We implement all φ , ψ , and Q functions with Mutli-layer Perceptrons.
Revised: We implement all φ , ψ , and Q functions with Mutli-layer Perceptrons.

Original: We train UVFA and UVFA+FARM with n-step Q-learning (Watkins & Dayan, 1992).
Revised: We train UVFA and UVFA+FARM with n-step Q-learning (Watkins & Dayan, 1992).

Original: When learning cumulants, USFA and MSFA have the exact same losses and learning alogirthm.
Revised: When learning cumulants, USFA and MSFA have the exact same losses and learning alogirthm.

Original: They both learn Qvalues and SFs with n-step Q-learning.
Revised: They both learn Qvalues and SFs with n-step Q-learning.

Original: We use n = 5 .
Revised: We use n = 5 .

Original: When not learning cumulants, following (Borsa et al., 2019), USFA only learns SFs with n-step Q-learning.
Revised: When not learning cumulants, following (Borsa et al., 2019), USFA only learns SFs with n-step Q-learning.

Original: All agents are built with JAX (Bradbury et al., 2018) using the open-source ACME codebase (Hoffman et al., 2020) for reinforcement learning.
Revised: All agents are built with JAX (Bradbury et al., 2018) using the open-source ACME codebase (Hoffman et al., 2020) for reinforcement learning.

Original: 5. C OMPOSITIONAL GENERALIZATION ACROSS OBJECT NAVIGATION TASKS
Revised: 5. C OMBINING OBJECT NAVIGATION TASK KNOWLEDGE

Original: MSFA learns disentangled functions for computing φ and estimating ψ .
Revised: MSFA learns modular functions for computing φ and estimating ψ .

Original: We hypothesize that this facilitates learning cumulants that respond to different aspects of the environment (e.g. to different object categories).
Revised: We hypothesize that this facilitates learning cumulants that respond to different aspects of the environment (e.g. to different object categories).

Original: This leads to the following research questions.
Revised: This leads to the following research questions.

Original: R1 :
Revised: R1 :

Original: Can we recover priorcompositional generalization results that relied on hand-designed cumulants for different object categories?
Revised: Can we recover prior generalization results that relied on hand-designed cumulants for different object categories?

Original: R2 :
Revised: R2 :

Original: How important is it to learn disentangled functions for φ and ψ ?
Revised: How important is it to learn modular functions for φ and ψ ?

Original: R3 :
Revised: R3 :

Original: Without GPI, do learning disentangled functions for φ and ψ still aid in generalization?
Revised: Without GPI, do learning modular functions for φ and ψ still aid in generalization?

Original: Setup.
Revised: Setup.

Original: We implement a simpliﬁed version of the object navigation task of (Borsa et al., 2019) in the (Chevalier-Boisvert et al., 2019) grid-world.
Revised: 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.

Original: The environment consists of 3 instances of 4 object categories.
Revised: The environment consists of 3 instances of 4 object categories.

Original: Observations are partial and egocentric.
Revised: Observations are partial and egocentric.

Original: Actions : the agent can rotate left or right, move forward, or pickup an object.
Revised: Actions : the agent can rotate left or right, move forward, or pickup an object.

Original: When it picks up an object, following (Borsa et al., 2019), the object is respawned somewhere on the grid.
Revised: When it picks up an object, following (Borsa et al., 2019), the object is respawned somewhere on the grid.

Original: Task vectors lie in w ∈ R 4 with training tasks being the standard unit vectors.
Revised: Task vectors lie in w ∈ R 4 with training tasks being the standard unit vectors.

Original: For example, [0 , 1 , 0 , 0] speciﬁes the agent must obtain objects of category 2 .
Revised: For example, [0 , 1 , 0 , 0] speciﬁes the agent must obtain objects of category 2 .

Original: Generalization tasks are linear combinations of training tasks.
Revised: Generalization tasks are linear combinations of training tasks.

Original: For example, [ − 1 , 1 , − 1 , 1] speciﬁes the agent must obtain categories 2 and 4 while avoiding categories 1 and 3 .
Revised: For example, [ − 1 , 1 , − 1 , 1] speciﬁes the agent must obtain categories 2 and 4 while avoiding categories 1 and 3 .

Original: This is challenging because the agent has never learned to avoid objects.
Revised: 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.

Original: Borsa et al.
Revised: Borsa et al.

Original: (2019) showed that USFA could generalize with hand-designed cumulants that described whether an object was picked up.
Revised: (2019) showed that USFA could generalize with hand-designed cumulants that described whether an object was picked up.

Original: train tasks.
Revised: 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.

Original: R1: MSFA is competitive with USFA, which uses oracle φ .
Revised: R1: MSFA is competitive with USFA, which uses oracle φ .

Original: Figure 5 shows USFA with a similar generalization trend to (Borsa et al., 2019).
Revised: Figure 5 shows USFA with a similar generalization trend to (Borsa et al., 2019).

Original: Tasks get more challenging as they are further from train tasks or involve avoiding more objects.
Revised: Tasks get more challenging as they are further from train tasks or involve avoiding more objects.

Original: For simply going to combinations of objects, USFA-Learnedφ does slightly worse than MSFA.
Revised: For simply going to combinations of objects, USFA-Learnedφ does slightly worse than MSFA.

Original: However, with more objects to avoid, all methods except MSFA (including USFA-Learned- φ ) degrades signiﬁcantly.
Revised: However, with more objects to avoid, all methods except MSFA (including USFA-Learned- φ ) degrades signiﬁcantly.

Original: For comparison, we show performance by an oracle bread-ﬁrst-search policy with access to ground-truth state ( BFS State Oracle ).
Revised: For comparison, we show performance by an oracle bread-ﬁrst-search policy with access to ground-truth state ( BFS State Oracle ).

Original: All methods have room for improvement when objects must be avoided.
Revised: All methods have room for improvement when objects must be avoided.

Original: In the appendix, we present heat-maps for how often object categories were picked up during different tasks.
Revised: In the appendix, we present heat-maps for how often object categories were picked up during different tasks.

Original: We ﬁnd that MSFA most matches
Revised: We ﬁnd that MSFA most matches

Original: USFA, while USFA-Learned- φ commonly picks up all objects regardless of task. 
Revised: USFA, while USFA-Learned- φ commonly picks up all objects regardless of task. We

Revised: (a) We ablate having modular functions for φ θ and ψ θ .

Revised: Generalization results degrade signiﬁcantly.

Original:  R2:
Revised: (b) We ablate leveraging GPI for generalization from all SF-based methods. MSFA without GPI can outperform both USFA-Learned- φ with GPI and USFA without GPI. This shows the utility of modularity for generalization. R2:

Original: Learning disentangled φ and ψ functions is critical for generalization.
Revised: Learning modular φ and ψ functions is critical for generalization.

Original: Learning an entangled function corresponds to learning a monolothic function for ψ or φ where we concatenate modulestates, e.g. ˜ ψ wt = ψ θ ( s (1) t , . . .
Revised: Learning an entangled function corresponds to learning a monolithic function for ψ or φ where we concatenate modulestates, e.g. ˜ ψ wt = ψ θ ( s (1) t , . . .

Original: , s ( n ) t , a, w ) .
Revised: , s ( n ) t , a, w ) .

Original: Disentangled functions correspond to equation 5.
Revised: Modular functions correspond to equation 5.

Original: Figure(a) shows that without disentangled functions for φ and ψ , performance severely degrades.
Revised: Figure(a) shows that without modular functions for φ and ψ , performance severely degrades.

Original: This also highlights that a naive combination of USFA+FARM—with entangled functions for φ and ψ —does not recover our generalization performance.
Revised: This also highlights that a naive combination of USFA+FARM—with entangled functions for φ and ψ —does not recover our generalization performance.

Original: We ablate leveraging GPI for generalization from all SF-based methods.

Original: MSFA without GPI can outperform both USFA-Learned- φ with GPI and USFA without GPI.

Original: This shows the utility of modularity for generalization.

Original: R3: Modularity alone improves generalization.
Revised: R3: Modularity alone improves generalization.

Original: For all SF-based methods, we remove GPI and select actions with a greedy policy: π ( s ) =
Revised: For all SF-based methods, we remove GPI and select actions with a greedy policy: π ( s ) =

Original: arg max a ψ θ ( s t , a, w ) (cid:62) w .
Revised: arg max a ψ θ ( s t , a, w ) (cid:62) w .

Original: Figure 6 (b) shows that
Revised: Figure 6 (b) shows that

Original: GPI is critical for generalization with USFA as expected.
Revised: GPI is critical for generalization with USFA as expected.

Original: USFA-Learned- φ beneﬁts less from GPI (presumably because of challenges in learning φ ).
Revised: USFA-Learned- φ beneﬁts less from GPI (presumably because of challenges in learning φ ).

Original: Interestingly, MSFA can generalize relatively well without GPI, sometimes doing better than USFA without GPI.
Revised: Interestingly, MSFA can generalize relatively well without GPI, sometimes doing better than USFA without GPI.

Original: 5. C OMPOSITIONAL GENERALIZATION ACROSS OBJECT NAVIGATION TASKS WITH NOVELVISUAL APPEARANCES AND ENVIRONMENT CONFIGURATIONS
Revised: 5. C OMBINING OBJECT NAVIGATION TASK KNOWLEDGE WITH NOVEL APPEARANCES AND ENVIRONMENT CONFIGURATIONS

Original: Beyond generalizing to combinations of tasks, agents will need to generalize to different layouts and appearances of objects.
Revised: Beyond generalizing to combinations of tasks, agents will need to generalize to different layouts and appearances of objects.

Original: Here we ask, R4 : Can MSFA enable compositional generalization in a visually diverse, procedurally generated environment?
Revised: R4 : Can MSFA enable combining task knowledge in a visually diverse procedurally generated environment

Original: Setup.
Revised: Setup.

Original: We leverage the “Fruitbot” environment within ProcGen (Cobbe et al., 2020).
Revised: We leverage the “Fruitbot” environment within ProcGen (Cobbe et al., 2020).

Original: Here, an agent controls a paddle that tries to obtain certain categories of objects while avoiding others.
Revised: Here, an agent controls a paddle that tries to obtain certain categories of objects while avoiding others.

Original: When the agent hits a wall or fence, it dies and the episode terminates.
Revised: When the agent hits a wall or fence, it dies and the episode terminates.

Original: If the agent collects a non-task object, nothing happens.
Revised: If the agent collects a non-task object, nothing happens.

Original: Observations are partial.
Revised: Observations are partial.

Original: Actions : At each time-step the agent moves one step forward and can move left or right or shoot pellets to open fences.
Revised: Actions : At each time-step the agent moves one step forward and can move left or right or shoot pellets to open fences.

Original: Training and generalization tasks follow the same setup as §5.1.
Revised: Training and generalization tasks follow the same setup as §5.1.

Original: there are objects to avoid ( w = [1 , 0 , 0 , − 1] ) in this setting (see text for more).

Original: However, as more objects need to be collected MSFA best generalizes (4 runs).

Original: R4: MSFA improves compositional generalization of object navigation tasks in a visually diverse, procedurally generated environment.
Revised: R4: MSFA enables combining task knowledge of object navigation tasks in a visually diverse, procedurally generated environment.

Original: Figure 7 shows that when an agent has to generalize to collecting more objects, modular architectures generalize best, with MSFA doing best.
Revised: Figure 7 shows that when an agent has to generalize to collecting more objects, modular architectures generalize best, with MSFA doing best.

Original: When objects have to be avoided, we see that no architecture does well, though MSFA tends to do better.
Revised: When objects have to be avoided, we see that no architecture does well, though MSFA tends to do better.

Original: We observe that avoiding objects leads agents to hit walls.
Revised: We observe that avoiding objects leads agents to hit walls.

Original: Since the agent always moves forward at each time-step in Fruibot, this makes avoiding objects a particularly challenging type of generalization.
Revised: Since the agent always moves forward at each time-step in Fruibot, this makes avoiding objects a particularly challenging type of generalization.

Original: 5. C OMPOSITIONAL GENERALIZATION OF HETEROGENEOUS TASKS
Revised: 5. C OMBINING KNOWLEDGE OF HETEROGENEOUS TASKS

Original: Aside from navigation tasks, researchers and practitioners will be interested in generally combining solutions to different tasks.
Revised: Aside from navigation tasks, researchers and practitioners will be interested in generally combining solutions to different tasks.

Original: We ask R5 :
Revised: R5 :

Original: Can MSFA enable combining solutions to heterogeneous tasks?
Revised: Can MSFA enable combining solutions to heterogeneous tasks?

Original: Setup. We leverage the “Minihack” environment (Samvelyan et al., 2021). Here, an agent is spawned at the top of a map with obstacles.
Revised: et al

Original: It must navigate to the a ladder at the bottom of the map.
Revised: It must navigate to the a ladder at the bottom of the map.

Original: The agent experiences partial observations of the environment.
Revised: The agent experiences partial observations of the environment.

Original: Actions : At each time-step the agent can select from 8 compass directions.
Revised: Actions : At each time-step the agent can select from 8 compass directions.

Original: Training tasks :
Revised: Training tasks :

Original: 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.
Revised: 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.

Original: Generalization requires that the agent generalize to (a) new conﬁgurations of objects and (b) to the combination of all three training conditions.
Revised: Generalization requires that the agent generalize to (a) new conﬁgurations of objects and (b) to the combination of all three training conditions.

Original: R5: MSFA enables combining solutions to heterogeneous tasks.
Revised: R5: MSFA enables combining solutions to heterogeneous tasks.

Original: Figure 8 shows that successor feature based methodsare bestable to successfully combine solutions to different tasks.
Revised: Figure 8 shows that SF-based methods best combine solutions to different tasks.

Original: In a small room, both SF-based methods get 100% success rate.
Revised: In a small room, both SF-based methods get 100% success rate.

Original: In a larger room, no method generalizes perfectly: SF-based methods do best at 75%, followed by UVFA-FARM and UVFA with 60%.
Revised: In a larger room, no method generalizes perfectly: SF-based methods do best at 75%, followed by UVFA-FARM and UVFA with 60%.

Original:  D ISCUSSION AND C ONCLUSION
Revised:  D ISCUSSION AND C ONCLUSION

Original: We have presented “Modular Successor Feature Approximators”, a modular neural network for learning cumulants and SFs produced by their own modules.
Revised: We have presented “Modular Successor Feature Approximators”, a modular neural network for learning cumulants and SFs produced by their own modules.

Original: We ﬁrst showed that MSFA is competitive with prior object navigation generalization results that relied on hand-designed cumulants (§5.1).
Revised: We ﬁrst showed that MSFA is competitive with prior object navigation generalization results that relied on hand-designed cumulants (§5.1).

Original: Afterwards, we showed that MSFA improves compositional generalization in a visually diverse, procedurally generated environment (§5.2).
Revised: Afterwards, we showed that MSFA improves an agent’s ability to combine task knowledge in a visually diverse, procedurally generated environment (§5.2).

Original: Finally, we showed that MSFA can combine solutions of heterogeneous tasks (§5.3).
Revised: Finally, we showed that MSFA can combine solutions of heterogeneous tasks (§5.3).

Original: Our ablations show that learning disentangled cumulants and SFs is critical for generalization with GPI.
Revised: Our ablations show that learning modular cumulants and SFs is critical for generalization with GPI.

Original: Interestingly, without GPI, disentangled cumulants and SFs improve generalization overa monolithic architecture (Figure 6).
Revised: Interestingly, without GPI, learning modular cumulants and SFs improve generalization over monolithic SF architectures (Figure 6).

Original: 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.
Revised: 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.

Original: Our results show that when learning cumulants, MSFA improves generalization with SF&GPI compared to USFA.
Revised: Our results show that when learning cumulants, MSFA improves generalization with SF&GPI compared to USFA.

Original: Additionally, without GPI, MSFA as an architecture improvescompositional generalization as compared to both FARM and USFA.
Revised: Additionally, without GPI, MSFA as an architecture improves generalization as compared to both FARM and USFA.

Original: Limitations .
Revised: Limitations .

Original: While we demonstrated reward-driven discovery of cumulants for transfer with SF&GPI, we focused on relatively simple task encodings.
Revised: While we demonstrated reward-driven discovery of cumulants for transfer with SF&GPI, we focused on relatively simple task encodings.

Original: Future work can extend this to more expressive encodings such as language embeddings.
Revised: Future work can extend this to more expressive encodings such as language embeddings.

Original: Another limitation is that we did not explore more sophisticated state-feature discovery methods such as meta-gradients (Veeriah et al., 2019).
Revised: Another limitation is that we did not explore more sophisticated state-feature discovery methods such as meta-gradients (Veeriah et al., 2019).

Original: 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.
Revised: 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.

Original: Future directions .
Revised: Future directions .

Original: SFs are useful for exploration (Janz et al., 2019;
Revised: SFs are useful for exploration (Janz et al., 2019;

Original: 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).
Revised: 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).

Original: We hope that future work can leverage
Revised: We hope that future work can leverage

Original: MSFA for improved state-feature discovery and SF-learning in all of these settings.
Revised: MSFA for improved state-feature discovery and SF-learning in all of these settings.