Original: C OMPOSING T ASK K NOWLEDGE WITH M ODULAR S UCCESSOR F EATURE A PPROXIMATORS
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 and modular architectues for learning 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 (SFs).
Revised: MSFA is composed of a set of modules, which each learn their own state-features and corresponding predictive representations (SFs).

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 metagradients (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 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.
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 combining task knowledge.
Revised: Hierarchical RL (HRL) is one dominant approach for combining task knowledge.

Original: 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.
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 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).
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 for generalization by combining 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;
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;

Original: Haarnoja et al., 2018), and to novel language instructions (Logeswaran et al., 2021).
Revised: Haarnoja et al., 2018), and to novel language instructions (Logeswaran et al., 2021).

Original: Our architecture,
Revised: Our architecture,

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

Original: This integration enables combining 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: MSFA is most closely related to
Revised: MSFA is most closely related to

Original: Recurrent Independent Mechanisms (RIMS) (Goyal et al., 2019) and its follow-up, Feature-Attending
Revised: Recurrent Independent Mechanisms (RIMS) (Goyal et al., 2019) and its follow-up, Feature-Attending

Original: Recurrent Modules (FARM) (Carvalho et al., 2021a).
Revised: Recurrent Modules (FARM) (Carvalho et al., 2021a).

Original: RIMs showed that leveraging modules to learn a state function improved out-of-distribtion generalization.
Revised: RIMs showed that leveraging modules to learn a state function improved out-of-distribtion generalization.

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

Original: MSFA differs from both in that it employ modules for learning value functions in the form of SFs.
Revised: MSFA differs from both in that it employ modules for learning value functions in the form of SFs.

Original: This enables a principled way to compose task knowledge while additionally generalizing to novel environment conﬁgurations.
Revised: 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 e , A , X , R, p, f x (cid:105) .
Revised: (cid:104)S e , A , X , R, p, f x (cid:105) .

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

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

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

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

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

Original: 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 :
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 = Q
Revised: Q π,wt = Q

Original: [Equation]
Revised: [Equation]

Original: ,w ( s
Revised: ,w ( s

Original: [Equation]
Revised: [Equation]

Original: =
Revised: =

Original: E
Revised: E

Original: [Equation]
Revised: [Equation]

Original: [ (cid:80)
Revised: [ (cid:80)

Original: [Equation]
Revised: [Equation]

Original: wt ] .
Revised: wt ] .

Original: SF&GPI .
Revised: SF&GPI .

Original: During training, we can have the agent learn policies for tasks: e.g. “open drawer” and “open fridge”.
Revised: During training, we can have the agent learn policies for tasks: e.g. “open drawer” and “open fridge”.

Original: 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”.
Revised: 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: We can leverage MSFA to have different modules learn different “cumulants”, φ , and SFs, ψ .
Revised: We can leverage MSFA to have different modules learn different “cumulants”, φ , and SFs, ψ .

Original: For example, 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 = ψ πt (cid:62) 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.

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

Original: (cid:80)
Revised: (cid:80)

Original: i
Revised: i

Original: α
Revised: α

Original: i w i .
Revised: i w i .

Original: We can re-use the SFs we’ve learned from training tasks M train for transfer with GPI.
Revised: 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
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

Original: [Equation]
Revised: [Equation]

Original: ∀
Revised: ∀

Original: [Equation]
Revised: [Equation]

Original: s,
Revised: s,

Original: [Equation]
Revised: [Equation]

Original: ∈ ( S ×
Revised: ∈ ( S ×

Original: [Equation]
Revised: [Equation]

Original: (Barreto et al., 2017).
Revised: (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 with 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 learning cumulants and SFs with modules improves zero-shot composition of task knowledge 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: 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.
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: MSFA then produces modular 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 (see Figure 3).
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 − 1 , 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 modular 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
Revised: Note that cumulants update with the module-state from the next

Original: time-
Revised: time-

Original: step.
Revised: step.

Original: This is
Revised: This is

Original: ﬁne
Revised: ﬁne

Original: since
Revised: since

Original: cumulants
Revised: cumulants

Original: are only
Revised: are only

Original: used
Revised: used

Original: during
Revised: during

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

Original: [Equation]
Revised: [Equation]

Original: { ψ θ
Revised: { ψ θ

Original: [Equation]
Revised: [Equation]

Original: s
Revised: s

Original: [Equation]
Revised: [Equation]

Original: , a, z )
Revised: , a, z )

Original: [Equation]
Revised: [Equation]

Original: w test } , where M train are train task vectors.
Revised: w test } , where M train are train task vectors.

Original: 4. L EARNING A LGORITHM
Revised: 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: =
Revised: =

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 =
Revised: L =

Original: [Equation]
Revised: [Equation]

Original: Q +
Revised: Q +

Original: [Equation]
Revised: [Equation]

Original: +
Revised: +

Original: [Equation]
Revised: [Equation]

Original: L φ .
Revised: L φ .

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

Revised: 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: combinations of object types while avoiding other object types.
Revised: combinations of object types while avoiding other object types.

Original: This is the setting used by USFA which assumed hand-designed cumulants.
Revised: This is the setting used by USFA which assumed hand-designed cumulants.

Original: (b) In Procgen, we study extending this form of generalization to a visually diverse, procedurally generated environment.
Revised: (b) In Procgen, we study extending this form of generalization to a visually diverse, procedurally generated environment.

Original: (c) In Minihack, we go beyond combining object navigation skills.
Revised: (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: This comparison shows shows the transfer beneﬁts of SF&GPI.
Revised: This comparison shows shows the transfer beneﬁts of SF&GPI.

Original: 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 FARM module.
Revised: Each state-factor s ( k ) is the output of a FARM module.

Original: This comparison allows us to test the beneﬁts of leveraging modules for learning value functions.
Revised: This comparison allows us to test the beneﬁts of leveraging modules for learning value functions.

Original: (3) Modular Value
Revised: (3) Modular Value

Original: Function Approximator(MVFA) is an adaptation of (Haarnoja et al., 2018) where modules learn individual Q-values Q ( i ) θ
Revised: Function Approximator(MVFA) is an adaptation of (Haarnoja et al., 2018) where modules learn individual Q-values Q ( i ) θ

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

Original: This comparison shows us to test the beneﬁts of leveraging modules for learning value functions in the form of SFs.
Revised: This comparison shows us to test the beneﬁts of leveraging modules for learning value functions in the form of SFs.

Original: (4) 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 hand-designed 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 OMBINING OBJECT NAVIGATION TASK KNOWLEDGE
Revised: 5. C OMBINING OBJECT NAVIGATION TASK KNOWLEDGE

Original: MSFA learns modular 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 prior 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 modular functions for φ and ψ ?
Revised: How important is it to learn modular functions for φ and ψ ?

Original: R3 :
Revised: R3 :

Original: Without GPI, do learning modular 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) BabyAI environment.
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: Despite the apparent simplicity of this environment, it presents several challenges.
Revised: Despite the apparent simplicity of this environment, it presents several challenges.

Original: First, learning φ is challenging because rewarding transition are rare subsets of all transitions leading to a data imbalance.
Revised: 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: Second, test tasks are out-of-distribution and involve behaviors the agent has never learned such as avoiding objects.
Revised: Second, test tasks are out-of-distribution and involve behaviors the agent has never learned such as avoiding objects.

Original: We expand on these challenges in the Appendix.
Revised: 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.

Original: We 
Revised: We show and error generalization

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

Original: Generalization results degrade signiﬁcantly.
Revised: Generalization results degrade signiﬁcantly.

Original: (b) We ablate leveraging GPI for generalization from all SF-based methods.
Revised: (b) 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.
Revised: MSFA without GPI can outperform both USFA-Learned- φ with GPI and USFA without GPI.

Original: This shows the utility of modularity for generalization.
Revised: This shows the utility of modularity for generalization.

Original: R2:
Revised: R2:

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

Original: Learning an entangled function corresponds to learning a monolithic 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: Modular functions correspond to equation 5.
Revised: Modular functions correspond to equation 5.

Original: Figure(a) shows that without modular 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: 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 OMBINING OBJECT NAVIGATION TASK KNOWLEDGE WITH NOVEL 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: R4 :
Revised: R4 :

Original: Can MSFA enable combining task knowledge in a visually diverse
Revised: Can MSFA enable combining task knowledge in a visually diverse

Original: procedurally generated environment
Revised: 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: R4: MSFA enables combining task knowledge 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 OMBINING KNOWLEDGE 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: R5 :
Revised: R5 :

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

Original: et al
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 SF-based methods best 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 an agent’s ability to combine task knowledge 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 modular 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, learning modular cumulants and SFs improve generalization over monolithic SF architectures (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 improves 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.