Improving Generalization with Approximate Factored Value Functions
Authors: Shagun Sodhani
Sergey Levine
Amy Zhang
Venue: Accepted by TMLR ()

Abstract: Reinforcement learning in general unstructured MDPs presents a challenging learning problem. However, certain MDP structures, such as factorization, are known to simplify the learning problem. This fact is often not useful in complex tasks with high-dimensional state spaces which do not usually exhibit such structure, and even if the structure is present, it is typically unknown. In this work, we instead turn this observation on its head. Instead of developing algorithms for structured MDPs, we propose a representation learning algorithm that approximates an unstructured MDP with one that has factorized structure. We then use these factors as a more convenient representation of the state for downstream learning. The particular structure that we leverage is reward factorization, which defines a more compact class of MDPs that admit factorized value functions. We empirically verify the effectiveness of our approach in terms of faster training (better sample complexity) and robust zero-shot transfer (better generalization) on the ProcGen benchmark and the MiniGrid environments.

Published in Transactions on Machine Learning Research (MM/YYYY)
Improving Generalization with Approximate Factored Value Functions
Shagun Sodhani Meta AT
sshagunsadheni@@mall. com
Sergey Levine
sedevinne@eeccs.eereeeey
University of California, Berkeley
Amy Zhang
amyzhang@mmeta.com
Meta Al
UT Austin
Reviewed on OpenReview: https: oppnreeveee. net/ forum? id= LuEWrnKyja
Abstract
Reinforcement learning in general unstructured MDPs presents challenging learning problem. However, certain MDP structures, such as factorization, are known to simplify the learning problem. This fact is often not useful in complex tasks with high-dimensional state spaces which do not usually exhibit such structure, and even if the structure is present, it is lypically unknown. In this work, we instead turn this observation its head. Instead of developing algorithms for structured MDPs, we propose representation learning algorithm that approximates an unstructured MDP with one that has factorized structure. We then use these factors as a more convenient representation of the state for downstream learning. The particular structure that we leverage is reward factorization, which defines a more compact class of MDPs that admit factorized value functions. We empirically verify the effectiveness of approach in terms of faster training (better sample complexity) and robust zero-shol transfer (better generalization) the ProcGen benchmark and the MiniGrid environments T
1 Introduction
Reinforcement Learning problems are often modeled as Markov Decision Processes (MDPs). While the MDP formulation is quite general, it is not always the most optimal. Reinforcement learning in such general, unstructured MDPs faces challenges like sample inefficiency and brittleness (Zhang et al. 2018a Song 2021). To address these challenges, recent works (Zhang et al. 2021a Sodhani et al. 2022b Zhang al. 2021bb have proposed introducing additional assumptions to leverage the structure
et al. 2019 Ghosh al.]
underlying the given task (s). While it is intuitive to leverage structure when it exists and is available, these
approaches may not be helpful when working with complex MDPs with high-dimensional state spaces where that structure often does not exist or is not known.
this work, we propose to leverage structure in the reverse direction. Instead of developing algorithms for structured MDPs, we propose to map an unstructuued MDP one that is approximate structured, where we can then exploit that structure for improved sample efficiency and generalization. In the scope of this work, focus class of structured MDP that call Factored Reward MDP, where the state and the rewards are factorized. Formulating benefit as these MDPs are more compact than standard MDPs. For example, consider case where the state space can be factorized into two variables, each taking m unique discrete values. The Factored Reward MDP formulation will use 2 dimensions, while the standard MDP formulation will use 22xm dimensions for
environment a Factored Reward MDP provides compulational
Project website: http:://tstttttteeeeeee////////e

---

Published in Transactions on Machine Learning Research (MM/YYYY
representing the state. From the cognitive science perspective, the factors can be thought of reusable mechanisms (Battaglia et al. 2018: Bapst et. al. 2019 Sanchez-Gonzalez et al. 2020) that can be shared across tasks, thus enabling the learning agent to transfer knowledge across tasks easily and quickly adapt to new tasks.
We also note that the Factored Reward MDPs are not merely a useful theoretical construct. This class of MDPs actually manifests in many real-world scenarios. When working with embedded systems, it is common to concatenate imput from different sensors into single observation vector, resulting in factorized state space. Suppose the training agent has access to this factorization. In that case, it can use specialized mechanisms for encoding the output of different kind of sensors or can selectively focus on the output of the most relevant sensors (for the task at hand). Similarly, when working on robotics applications, the different. joint angles, positions, and velocities are concatenated into a single, large state space Hodge et al. (2021) even though these factors could lead to independent rewards. For example, the agent could be rewarded for running at a particular velocity as well as for reaching a spocific position. Knowing the underlying lactorization could help the learning agent optimize for the overall lask.
show that Factored Reward MDPs exhibit factored value functions which may allow for better generalization to novel combinations of those factors. Existing RL algorithms can be extended to exploit the Factored Reward MDP's structure. Motivated by the advantages of Factored Reward MDPs, we propose a representation learning technique, called Approximately Factored Representations (AFaR), that can map an unstructured MDP into an approximate Factored Reward MDP. We show that this can lead to sample efficiency gains and improved generalization performance even when that mapping is approximate.
Contributions: (i) We propose represcntation learning technique to map given MDP into approximate Factored Reward MDP. ii) We show that the proposed representation learning algorithm can be easily plugged into standard actor-crite critic algorithms and the entire system can be trained end-to-end. (iii) We empirically verify the effectiveness of the proposed algorithm in terms of better sample complexity (faster training) and better generalization (robust zero-shot transfer).
2 Problem Setup
We first define the standard assumptions for a Markov decision process (MDP) and introduce a new form of structured MDP we call a Factored Reward MDP in Section 2.1 Next, in Section 2.2 we show that Factored Reward MDP exhibits the useful property emitting factored value functions. In Section 3 we use these concepts in reverse to design a representation learning approach for transforming standard MDPs into we describe how we use this factored representation and
approximate Factored Reward MDPs. In Section the factored value functions to derive policy.
4
2.1 Preliminaries
A Markov Decision Process (MDP) (Bellman 1957 Puterman 1995) is defined by a tuple (S,A, where S is the set of states, A the set of actions, R: x + is the reward function, T' 8 x A + Dist(S) is the environment transition probability function, and 7 E (0, 1) is the discount factor. At each time step, the learning agent perceives state 8; S, takes action E A drawn from a policy : S A - [0, 1], and with probability T(st41ls;, a.) enters next state St+1. receiving a numerical reward R(s;, a,) from the environment. The state value function of policy is defined as: (s(s) = Ex i 7 {(((s++))S) s]. The optimal state value function V* is the maximum value function over the class of stationary policics.
A Factored Reward Markov Decision Process (Factored Reward MDP) is a special class of MDP where the state and the reward can be factorized into variables or factors. For example, a state s, can be factored into & state factors (s)...... st). The set of possible values that the 4t state factor take are represented as St and the state space S is the cross product of the value spaces for the individual state factors, i.e S S St We rely on the following assumption:
Assumption 2.1 (Factored Rewards). For given full states s, E S, action at E A, and st denoting the tth factor of the state St, we have R(st, a.) rst a.) where p² S A IR are the local reward functions.
2

---

Published in Transactions on Machine Learning Research (MM/YYYY)
Factored Reward MDPs are useful abstraction for representing high-dimenniooal MDPs generated by systems with many weakly connected components. These MDPs can also be viewed as relaxed form of the Factored MDPs (Boutilier et al. 1995 1999 Kearns & Koller 1999). another class of MDPs that rely on the additional assumption factored transition function.
2.2 Factored Reward MDPs Emit Factored Value Functions
Given the above definition for Factored Reward MDP, we show that the state value functions and the state-action value functions are also factorized.
Proposition 2.2. Consider Factored Reward MDP where the state is factorized into k state factors
st..... of and the reward function is factorized as R(s,, a1) r az) where r St x A + R are the local reward functions. The state value function and the state-action value functions for any policy a can be wrillen as:
k (*((s) Σ
(6)) Q"(s,,
=
Σ Q7
(1)
where
(T = Σ
Σ
T(84+1884 or
)errr (sy, ax),
1'=t
9,441
= =
Σ
Σ
T(se+ on )r (sy.
v=1+1
are the state value and the state-action value functions for the 8th factor, respectively and the actions a₂ are oblained using the policy m.
The proof is available in Appendix A Our proposed approach (described in Section does not assume access to factored Q-functions. Instead, it applies these concepts in reverse to learn a representation that transforms standard MDPs into approximate Factored Reward MDPs.
3 AFaR: Learning Approximate Factored Representations
showed that Factored Reward MDPs emit. factored value functions. We now reverse this observation to propose representation learning technique that induces such factorization. Specifically, we approximate the value of a state as the sum of the values corresponding to the different factors In practice, this translates to computing the state-value function for a given state in terms of the sum of the state-value function for the factors of the given state. These factors are learnt using a. Factor Encoder module, denoted as d, and is described in detail in Section 3.1 Our hope is that learning the factorized representation could be helpful for improving the generalization performance and sample efficiency of RL agents on environments with factorizable observation spaces. We refer our proposed approach Approximately Factored Reprosentations, or AFaR. In Section 4 discuss leveraging this form of structured reprosentation for rownstream control.
In Section 2.2
3.1 Factor Encoder
Factor Encoder module, denoted by Φ, an encoder that maps given input state 8t set of k factor representations, denoted as ,2f,,...,f))) Here, & is hyperparameter. For each factor in the
Note that in practice, to the number of factors. Civen that the number of factors constant for given instantiation of the training agent, using multiplicative constant. Using mean, in place of sum, ensures that the maguitude of output of value
the operation instead of the operation to normalize the value function with respect
or is equivalent functions (and training losses) remains of the order as we vary the number of encoders, there by enabling us the same learning across experiments.
3

---

Published in Transactions on Machine Learning Research MM/YYYY)
(e)
(f)
(g)
V(=
>
[(.ss)
St
2
((2)
V(A
(d)
(a)
2
k
//
(b)
(c)
0
π
(k)
Figure Architecture the AFaR model: Given an input state se (shown as (a)), the Factored Encoder module ø (shown as (6)) is used to compute the representations for the k factors, denoted (. (shown as (c)). State value is computed for each factor (shown as (d, e)) and the value for the given state is obtained by summiing over the individual state values (shown (J. g)). The state values, along with factor representations, are used to compute the attention scores on (shown as (h, i)) which is used for aggregating the factor representations by the Aggregalion Module, 9, (shown as (k))- The aggregated feature representation is used to select the action from the universal policy (shown as (1)) Dashed lines indicates the components (or computations) through which gradient does not flow. The blue line and box show that the same factored representations are used in (c) and (j) stages.
factored representation. Following Equation (9), obtain the state value for the state st by summing over the state values corresponding to the individual factors (shown in step f in Figure 1
compute its corresponding state value, denoted as Vi V(z))Vi E {1.-... k)).
The Factor Encoder module is instantiated as a neural network, and we overload the notation to denote both the module and the learnable weights (of the encoder) using Φ. The Factor Encoder module trained using the critic loss alongside the value function. We note that while we use the subscript i in the definition of the value functions in Equation (9) (e-g. V.), the neural network (instantiating the value funetions) are shared across all factors (shown in component d in Figure D Using shared value function helps to ensure that the learned state factors lie in the same representational space.
4 Extracting a Policy from Factorized Value Functions
described our proposed representation learning approach, AFaR, that enables modeling an MDP as an approximate Factored Reward MDP. Since AFaR is representation learning algorithm, we must pair with a policy optimization algorithm for end-to-end reinforcement learning. In this section, we describe how to combine AFaR with actor-critic algorithms. We outline a standard actor-critic algorithm in Algorithm 1 and highlight the changes (in red color) for incorporating the AFaR for performing end-to-end reinforeement learning. In Section 5 we evaluate the effectiveness of our proposed approach empirically by testing AFaR in conjunction with three different actor-critic algorithms - Exploration RIDE (Raileanu & Rocktäsche] 2020), Data-regularized Actor-Critic or DrAC (Raileanu ct al. 2020) and Invariant Decoupled Advantage Actor Crroii or IDAAC (Raileanu & Fergus 2021 For additional details about these algororithms and why we specifically choose them as baseline algorithms, we direct the readers to Appendix E.1
In Section
Rewarding Impact-Driven
Previous works have shown that factorized state and value functions alone need not. lead to factorized policies (Liberatore] 2002 Allender al.) 2003). As such, we need to combine the factorized representation into single representation that can be fed into universal policy The learning agent aggregates the
Here, universal policy refers policy shared across the factors
4

---

Published in Transactions on Machine Learning Research (MM/YYY)
Algorithm 1 AFaR algorithm.
Require: Actor-Critic Components
Require: Factor Encoder 6
Require: Aggregation Module 0
Require: Attention Module
Black Text: Standard Actor-Critic Algorithm
Red Text: Proposed changes
Cyan Text: Comments
4: for each limestep / = I...7 do
5:
(()) = [()) Vi (1..... 1}
6:
(b), (c) in Figure
7:
Compute state value (*)
Compute state values V77 Vi (1,...,k) k) (d), (e) in Figure
9:
Compute state value (*(8)) = Ei= (z(z)) using Equation 9 (f) in Figure
10:
Compute a = 6(2); ... ,,,2, V (xt) ... (=())) using Equation 3 (h), (1) in Figure
11:
Compute 2; = [(2),
0. a4f) (k) Figure
12:
Sample a, ~ TT (2,)
13:
Sample reward rı and next state 8141
14:
Compute TD error 6, Te + y ($(++)) - (s;) using Equation 9
15:
Update value function and the Factor Encoder using the critic loss (TD error)
Update policy, Attention Module and Aggregation Module a using the actor loss (Policy-Gradion Theorem) 17: end for
16:
factored representations using an Aggregation Module, denoted by 0. Essentially, the Aggregation Module returns a distribution of weights (denoted as {w) ... }) over the factored representation, which are used to compule a weight sum of the factors.
In practice, we normalize the weights using the softmax operator, ensure that they sum up The resulting weights, {a! ...afaaf can be interpreted as some form of attention weights as they control how much the policy allends to any given factor. The aggregated representation (that will be fed to the policy network) is computed as (z)],... at ... .at)) Ei= 2f x af where denotes the scalar product. This factor aggregation operation is shown as the component in Figure 1
So far, we have described the high-level mechanism to aggregate the factored representations before passing to the policy network. In the next section (Section 4.1 we describe how we can obtain the attention scores for aggregating the representations.
4.1 Attending to the Factorized Representations
Given collection of factor representations (represented jointly as ar) that capture the relative importance of the different factors for the given task. These attention scores would enable the policy to determine the most relovant factors and attend to them more strongly than the less relevant ones.
]+ we want to learn attention scores [a], ... af
There are several ways of computing ar. In the simplest case, we can of = 1. ie. assign equal weight to all the factors. We note that this simple instantiation of the proposed AFaR approach is competitive with the baseline approaches for simpler environments. However, it does not perform well on more complex environments. This limitation likely stems from the fact that not all factors are relevant at all times, making the equal-weight strategy sub-optimal. For example, if the agent is searching for a key to open a door, it may not need to attend to the door till it finds the key.
A more systematic and expressive approach to learn the attention scores using an Attention Module, denoted using sa that takes as input the factor representations, [=],-.. 2 and outputs the attention scores
II practice we concatenate the different. factored representations. Other choices, like averaging over the factored representation, before feeding the Attention Module should also work.
5

---

Published in Transactions on Machine Learning Research (MM/YYYY)
la)....,,af ,af]. The learnt attention scores enable the Aggregation Module to condition only the important factors, while ignoring the irrelevant factors.
The Attention Module can be instantiated as a feedforward network that is trained using the actor loss, along with the policy network. Now we can compute the attention scores as:
0₁ x( ,
(2)
In the next section (Section 4.2 when aggregating the factored representations.
we describe how we can leverage the structure of Reward Factored MDP
4.2 Leveraging the structure of Factored Reward MDP
Actor-critic algorithms commonly train the state value function, V, as one of the many components. The output of this state value function captures the value of a given state, as measured in terms of the expected returns. Comparing the state value functions for two states can determine which state will likely lead to higher returns for the learning agent. Extending this argument the state factors, a factor with a higher value of the state value function will lead a higher expected return for the agent than a factor with a lower value. This observation motivates the use of the state value of the factors input to the Attention Module. We compute the atlention scores as:
ag = c ". ,,, v(+)), VV(V)(
(3)
In practice, we concatenate the different factored representations, and the output of the state value function, corresponding to the different factors before feeding the Attention Module should also work. The factored representations and the state values are detached from the computation graph before passing to the Attention Module. This ensures that only the critic loss updates the value function and the (factored) representations. Effectively, the critic cont what the different factors learn (via the critic loss), and the actor controls which factors are used for computing an action (via the policy loss). 5
5 Experiments
So far, we have described how to learn factor representations and use them when training RL algorithms. In this section, we empirically verify the usefulness of the AaRR technique by designing experiments to answer the following questions: (i) Will the RL policy Icarned using the factored representation be useful? (ii) Does AFaR improve generalization to new environments where factors vary in novel ways? Here, "new environments" refer to environments with different layouts (or sizes) of different level of difficulty than the original environments. "Factors vary in novel way" refer to change in properties of the factors that the agent saw during training. For example, in the BigFish environment (from ProcGen), the number/size/shenn of fishes varies across the environments. Similarly, in the MiniGrid environment, the number of doors and their types (i.e. locked or unlocked) varies across environments, (iii) How robust is AFaR to the number of factors that the agent learns?
We use the Procgen benchmark (Cobbe al. 2020) and the MiniGrid environments Chevalier-Boisverr et al. 2018) to evaluate the effectiveness of the proposed AFaR algorithm. Both environments are designed to assess RL algorithms for generalization and sample efficiency. Additionally, they satisfy critical criteria like high diversily across the environment instances, optimized implementation for fast. training, and shared obsorvation and action spaces. These environments are further described in Section 5.3 and Section 5.4
FFllowing the same line of reasoning, the factored representations detached from the computational graph even in the of Equation ?

---

Published in Transactions on Machine Learning Research (MM/YYY)
5.1 Choice of actor-Critic Methods
In Section | we noted that AFaR is representation learning algorithm, and need to pair it with a policy optimization algorithm for end-to-end reinforcement learning. In Algorithm [ we highlighted the changes (in red color) for incorporating the AFaR with standard actor-critic algorithms. In this section, we describe the different actor-critic algorithms that we use in conjunction with AFaR and outline the rationale for choosing these specific baselines.
Our first criterion for selecting the baseline actor-critic algorithms was strong performance on the selected environments. For the Procgen environments, use two baselines: (i) Data-regularized Actor-Crite (DrAC) Raileanu et al. (2020), and (ii) Invariant Decoupled Advantage Actor-Critic (IDAAC) Raileanu &
Fergus [021). The key difference between these two approaches is that IDAAC is the state-of-the-at among mothods that use separate policy and value function, while DrAC is the stale-of-the-arl among methods that use a shared network for policy and value functions. For MiniGrid environments, we use Rewarding Impact-Driven Exploration (RIDE) Raileanu Rocktäschel (2020) method, a statc-of-tho-art method for lasks on the MiniGrid environments.
There are additional benefits when using these methods as the underlying actor-critic methods for the AFaR approach. These methods incorporate different inductive biases and focus on different challenges in training RL agents. For example, DrAC uses dala--ugmentattin techniques, IDAAC uses disjoint policy and value functions, and RIDE is designed to address exploration-related challenges. Moreover, these methods also differ in terms of how the value functions and the encoders are trained. For example, IDAAC learns a value function to provide the learning signal for an advantage function, which, in turn, is used for training the policy. RIDE uses two sets of encoders one for training the policy and the other for computing the intrinsic rewards (used for computing the exploration bonuses). Integrating AFaR with such diverse methods and demonstrating that integrating with AFaR can improve over these baselines shows that our proposed algorithm is helpful for a variety of actor actor-critic algorithms.
For additional details about the underlying actor-critic methods, refer to Appendix
5.2 Integrating with Actor-Criti Methods
Integrating AFaR with existing actor-critic methods requires replacing the given encoder (from the underlying actor-critic method) with a Faclor Encoder. A simple approach for instantiating a Factor Encoder is to start with the given encoder and scale the output dimension by factor of k, where
the number of factors that we want to learn. Since the factored representations are aggregated using weighted-sum, the size of the subsequent (and even the previous) layers does not have to be changed. Using AFaR increases the number of trainable parameters in the network. For the RIDE baseline, the number of parameters increase from 18M to 18M + 1.05 (k 1),M, ie. increase of about 5.8% parameters for each new factor. For the DrAC baseline, the increase is from 626K to 626K + 66K x (k- 1), i.c. an increase of about 10.5% parameters for each new factor. For the IDAAC baseline, the increase is from 1.25M to 1.25M + 66K x (k- 1), i.e. an increase of about 3.2% paramcters for each new factor.
When training the combined models, we take care not change any hyperparameters (like learning rate, weight decay, exploration bomuses etc) for the underlying actor-critic algorithms. This choice ensures a fair- comparison as the combined model is not likely to perform better because of use of "better" hyperparameters. Any improvement in the performance has to come because of use of AFaR approach. This is also useful for demonstrating that AFaR can be used with the existing baselines without having to retune all the hyperparameters from scratch, which could have raised the bar of adoption for the proposed approach. Wo run all experiments with 10 seeds and report the mean and the standard error for all the results.
5.3
Progen Benchmark
5.3.1
Environment
The Procgen benchmark comprises 16 procedurally generated environments shown in Figure Each environment represents
7

---

Published in Transactions on Machine Learning Research (M//YYYY)
distribution of levels where each level corresponds to new environment, allowing for strict partition between train and test environments. Each environment has multiple objects and entitics that the agent can interact with. For example, in the BigFish environment, the agent comes across many fishes (of different shapes, sizes and color). In this case, the AFaR agent can learn factors (representations) corresponding to the different fishes.
Following the setup in Raileamu et al. (020), we train the agent on a fixed set of 200 levels while testing on the full distribution of levels. In practice, this is simulated by sampling levels at
random during evaluation. For additional details about the
environment, refer Appendix B.1
Environment.
DrAC
AFaR
5.3.2 Baselines
Bigfish
8.88 0.89
1..5
DrAC builds on the idea of using data augmentation for improving gencralization in RL (Laskin al., 2020 Srini vas et al., 2020 Kostrikov
Bossfight
0.21
7.79
Cavellyer
0.21
5.23
al.) 2020,, and introduces regularization terms that make the use of data augmen- tation theoretically sound for actor-critic algorithms. We replace the shared observation encoder with our Factored Encoder. IDAAC proposes using separate networks to model the policy and the value functions by suggesting that using a shared representation (for actor and critic) can lead to overfitting. The learnt value function is used for training an advantage function which acts the critic for the policy. Further, the advantage function shares parameters with the actor network, as is the case for standard actor- actor-critic architectures. So we factorize the obsorvation encoder of the policy (that shared with the advantage function). For the IDAAC architecture, "critic" (as used in Algorithm ! and Figure D refers the advantage function. For additional details on the baselines, refer to Appendix
Chaser
6.26
7.07
Climber
6.16
7.14
Coinrun
8.62
#.62:0009
Dodgeball
4.28+
5.24
Fruitbot
27.85
2.12
Heist
3.96
0.15
Jumper
5.14
+.09
Leaper
0.3
5.19
Maze
6.33 0.11
6.72 0.07
Miner
9.113
0.09
Ninja
5.49 + 0.11
Plunder
8.12 0.48
9.27 + 0.34
Starpilot.
29.67 + 0.85
28.89 + 0.65
Table We compare AFaR with DrAC (as the base- line approach) on the evaluation environments after training for 25M steps. We report the mean the standard error over 10 runs. The results marked bold denote the results that are statistically sig- nificantly better. Additional details about testing for statistical significance can be found in Appendix
5.3.3 Results
We train individual policies on each of the 16 environments using 200 levels (per environment) and report the score the evaluation environments, after training for 25M (environment) steps. The agent is evaluated on the full distribution of the levels (generated by sampling levels at random during evaluation). Along side reporting the mean and the standard error metrics (computed over 10 seeds), we also report environments where the improvements are statistically significant. Details about testing for statistical significance can be found in Appendix
Dodgetoll
Plunder
DrAC In Table I compare the performance of DrAC (as baseline) with AFaR. (using DrAC for learning the policy). We observe that the proposed AFaR approach leads to stalistically significant improvement in 7 out of 16 environments. The corresponding training and evaluation plots are presented Figure Cnd Figure 7 (in Appendix). We also perform an ablation where we train the systems using just 10 levels (instead of 200 levels) which severely
Frittbot
20
Atali
(x200)
147

---

Published in Transactions on Machine Learning Research
limits the diversity in the training data and can hurt the generalization performance significantly. In Figure g (in Appendix), we show that while the generalization performance degrades for both the approaches, AFaR still outperforms (or is par) the baseline approach on all the environments.
We compare the performance of AFaR with DrAC (as the baseline approach) on the evaluation environments from Proegen Benchmark, when training with 10 levels (per environment) for 25M steps. The solid curves correspond to the mean performance (score) while the shaded regions represent the standard error (computed over 10 runs). For corresponding curves when training on 200 levels per environment, refer to Figure
AFaR introduces new hyper-parameter k, or the number of representations that the Factor Encoder produces. We perform ablation experiment to understand how sensitive is the performance of AFaR algorithm to this hyperparameter. In Figure g we note that in some cases, the choice of k does not make much difference in the convergence performance for the agent. However, in general, the performance can be improved by tweaking
IDACC In Table 2 we compare the performance of IDAAC (as baseline) with AFaR (using IDAAC for learning the policy) using the same experiment, protocol before. Recall that TDAAC is actor-critic method that separate policy and value functions. We observe that the proposed AFaR approach improves over the baseline in 14 out of 16 environments (slatistically significant improvement in 4 environments). We also consider two additional baselines. In the first case, we factorize only the actor and not the critic. The actor is factorized as before but the critic produces a single state-value, using the concatenated factors as the feature. The agent assigns weight to the factors by using another neural network. This agent is refered to as the Factorized Actor. In the second case, we disable the AMention Module and assign equal weights all the factors while aggregating them. We denote this ablation as AFaR-mean. As we observe in Table / Factorized Actor sometimes improves over the IDAAC baseline but gencrally lags behind IDAAC. AFaR-mean is generally stronger than Factorized Actor but lags behind AFaR.
5.4 MiniGrid
5.4.1
Environment
MiniGrid is a procedurally generated grid-world environment where the "world" is a N x M grid of tiles. Each tile may contain atmost one object and the learning agent has to navigate the grid and find and interact with objects to complete different lasks. In these environments, the AFaR agent can learn factors (representations) corresponding to the different objects. We use the following environments from MiniGrid:
1. MiniGrid-KeyCorreei Kororrdo-kyyyrrrrorkkkky environments where x denotes the size of room and denotes the number of rows for compactness we will denote these by KeyCorridorSxRy. The agent's 's task is to pick up an object behind a locked door. The key to this locked door is hidden in another room, and the agent has to explore the environment to find the hidden key.
2. Minirrid-Kult maximum size of a room - for compactness we will denote these by Mul MultiRoom The agent is placed in the first room and must navigate to green goal square in the most distant room from the agent
environments where denotes the mumber of room and y denotes the
For additional details about the environment, refer Appendix B.2

---

Published in Transactions on Machine Learning Research (MM/YYY)
Environment
IDAAC
Factorized Actor
AFaR-mean
AFaR
Bigfish
17.2 0.8
16.99 + 1.2
1..4
20.9 + 1.5
Bossfight
9.5
10.23 + 0.5
1..
10.77 - 0.8
Caveflyer
5.25
5.37 0.8
5.75
5.5 0.5
Chaser
1.0
0.8
7.2114
8.1
Climber
0.4
0.6
0.5
9.4
Coinrun
9.4
9.27
9.25
9.7
Dodgeball
0.3
0.3
3.15
0.5
Fruitbot
2.9
2.94
2.5
3.1
Heist
0.4
0.5
3
3.75
Jumper
0.0
6.45
6.75
7.7
Leaper
1.75
0.7
0.9
8.5
Maze
0.3
0.6
5.0.88
6.7 0.3
Miner
0.6
1.3
8.9
1..16
Ninja 0.9
7.37
7.75
8.25
Plunder
1.1
22.71
21.7
24.39 1.2
Starpilot 41.65 2.1
37.91 1.9
38.82
42.35 2.2
Table 2:
We compare AFaR with IDAAC on the evaluation environments after training for 25M steps. We report the mean and slandard error over rums. The results marked bold denote results that are (i) best performing among all the approaches and (ii) are statistically significautly better than IDAAC or AFaR or both (whichever applicable). Additional details about testing statistical significance can found Appendix D
5.4.2 Baselines
MiniGrid environments use a sparse reward setup which makes exploration key challenge. RIDE proposed novel intrinsic reward mechanism in which the agent is rewarded for actions that affect its learned state representation. On the MiniGrid environments, RIDE improves sample efficiency compared to approaches like count-based exploration Bellemare et al. (2016), Random Network Distillation Burda et al. (2019). Intrinsic Curiosity Module Pathak et. al. 2017) and IMPALA (Espeholt et al. 2018).
We note that RIDE proposed using two sets of encoders one for the actor-critic method and the other for computing the intrinsic reward. The first encoder is trained along the lines of the standard RL works, while the second encoder trained using forward and backward dynamics loss. Since the Factor Encoder, é, is proposed in the context of actor-critic algorithms, Encoder and do not change the design of the second encoder (that is used for the computation of the intrinsic reward). This design choice also enables a fairer comparison between RIDE and the proposed approach as we do not change how the intrinsic rewards are computed (which core contribution of the RIDE approach). Further, we modify the actor-critic components of RIDE as per Algorithm for maintaining a history over the observations. We share the weights of the recurrent network across the different factors (with every factor having their own hidden state).
only replace the first encoder with the Factor
I RIDE uses recurrent network
5.4.3 Results
With the MiniGrid environments, we look at two forms of generalization. In the first case, we train and evaluate the agents on given environment. Note that every time the environment reset, a new instance of the environment is sampled, with a new randomized layout and new arrangement of objects. The agent also spawns at random location in the grid. This ensures that the learning agent can not solve the task during evaluation by simply memorizing the training instances, not only all the training instances unique (which makes memorization very challenging), the test instances are disjoint from the training instances. We perform this evaluation with the both the cases, we report the number of environment interactions (in millions) that learning agent nceds to converge. For KeyCorriddrs3R3-v0 environment, the success rate at convergence 90% while for MultiroomU7S4-v0 environment, the convergence performance 75%. Note that for both the environments, the convergence performance of the agents is the same though the sample complexity different. Tn Table
environment and the MultiroomN7S4 environment.
10

---

Published in Transactions on Machine Learning Research (MM/YYY)
note that the AFaR approach is about 18% more sample efficient that RIDE for KeyCorr1dorS3R3--0 and about 33% more sample efficient for the MultiroomM/SS-vD environment (Tn Table 4
Similar to the IDACC setup, we consider two additional baselines. In the first case, we factorize only the actor and not the critic. The actor is factorized as before but the critic produces a single state-value, using the concatenated factors as the feature. The agent assigns weight. to the factors by using another neural network. This agent is referred to as the Factorized Actor. In the second case, we disable the Attention Module and assign equal weights to all the factors while aggregating them. We denote this ablation as m We note that while mean provides a marginal improvement over RIDE, it lags behind AFaR. Factorized Actor needs significantly more steps to solve the tasks. The corresponding evaluation curves can be found in Figure
Similar to the IDACC setup, we consider two additional baselines. In the first case, we factorize only the actor and not the critic. The actor factorized as before but the critic produces a single state-value, using the concatenated factors as the feature. The agent assigns weight to the factors by using another neural network. This agent is referred to the Factorized Actor. In the second case, we disable the Attention Module and assign equal weights to all the factors while aggregating them. We denote this ablation AeaR-mean. As we observe in Table 2 AFaR outperforms all the basclines. AFaR- mean is sometimes (slightly) better tha RIDE while under-performing on one environment. Factorized Actor lags behind all the other nethods.
In the second case, we train the agent on one environ- ment and evaluate it on a different environment in a zero-shot manner. Specifically, we perform zero-shot evaluation, at regular intervals, on KeyCorr1dorS3R1-o and KeyCorriddrs3R2--0 enviromments, while train- ing on the KeyCorridorS3R3-v0 environment and on Multiroom//1054-00 environment while training on the Multiroom//7$4-v0 enviromment. Following the protocol in the previous set of experiments; we report the number of environment interactions (in millions) that a learning agent needs to converge. For syyCorridorS3R1--- KeyCorridor 1or23322-00 environments, the zero-shol success rate at convergence is 90% while for MultiroomNN004--v environment, the zero-shot the convergence performance is 75%. This second setup is more challenging than the first setup as now even the size of the rooms and/o the grid changes, along with the layout. The correspond- ing results are shown in the second and third rows of Table 3 and in the second row of Table 4 where we record the number of environment interactions (in millions) on KeyCorridorSan3-v and Mul Multiroom/7744-00 respectively. We note that the AFaR approach requires much fewer samples to generalize to the unseen environments. The corresponding evaluation curves can be found in Figure D
10
Figure 4: We plot how the state-value for the dif- ferent factors, averaged across steps in an episode and across a batch of environments, changes over time for the KeyCorridorS3R3-v0 environment and the MultiroomN774- environments (left to right).
and
We want to ensure that: (i) the factors are not degenerate, i.e., most or all-but-one factors are not collapsing to 0 value, and (ii) AFaR is not learning to simply average the state-value of the factors to compute the overall state value. In Figure 4 we plot the stato-value for the different factors, averaged across steps in an episode and across a batch of environments, changes over time for the KeyCorridorS3R3-v environment and the MultiroomM7744-00 environments (left to right). As we can see in the figure, all the factors are assigned non-zero (state) value (i.e., the factors do not collapse to 0 value), and all the factors are not contributing uniformly to the value function (i.c., the overall state value is not simply a mean of the value corresponding to the different factors). Moreover, the distribution of the state value (across factors) seems to depend on the task. For the KeyCorridors3R3--0 task, the agent has to interact with multiple values, and the value assigned to the different factors is quite close. The MultiroomM7744-00 task does not focus on interacting with multiple objects, so the distribution (of state values) is more uneven.
we plot how the stato-value for the different factors, averaged across steps in episode and across batch of environments, changes over tine for the KeyCorriddosS33--00 environment and the Mult MultiroomN754-1 environments (left to right).
In Figure 4
11

---

Published in Transactions on Machine Learning Research
Environment
AFaR
RIDE
A-aR- mean
Factorized
Actor
Number of environment interations (in millions)
MiniGrid- SumnirridkeeyCooridooSSSSS.
8.5
10.4
10.3
13.6
MiniGrid- MiniCriddKKyyCrrrdosSSiii KeyCorrido:S3R11-
7.6
8.9
8.9
9.05
MiniGrid- Reyn
8.5
9.2
10.1
10.2
Table 3: We compare the performance of AFaR with RIDE (baseline approach) KeyCorridorS3R3-v0 KeyCorridorSSR1-p0 and KeyCorriddrS3R2--0 environments (top to bottom) in terms of the number of environment interactions (in millions) that the agent needs converge the 90% of the AFaR algorithm where we disable the Attention Module and assign equal weights to all the factors while aggre- gating them. We denote this ablation as AFaR-mean. Note that while the agent evaluated (KeyCorriddrS333-v0 JeyCorridoS33R. environment, thus evaluation on KeyCorridorS3R2-00 and KeyCorridorS3R1-0 enviromments is done zero- shot manner. For all the three environments, AFaR obtains the best sample efficiency. The corresponding evaluation curves can be found in Figure
rate. We also include ablation version
three environments KeyCorridorS3R1-00) it was trained only on the KeyCorridorS3R3-v
Environment
AFaR Number of environment interations (in millions)
RIDE
NiniCGrii-Multtioonn/SS4vv MinnGrii-Muutiioo..NOSSvvv
6.8
9.03
7.6
9.7
Table 4: We compare the performance AFaR with RIDE (haseline approach) on Maltiroom:/754-00 and MultiroomM10044-00 environments (top to bottom) in terms of the number of environment interactions (in mil- lions) that the agent needs to converge to the 75% success rate. Note that while the agent is evaluated on two environments (Multirromm/754-v00 MultiroomM10044-0, thus the evaluation on MultiroomM1004- environment is done in a zero-shot manner. For both the environments, A AFaR obtains the best sample efficiency.
trained only the Multirooml7554-00 environment,
12

---

Published in Transactions on Machine Learning Research (MM/YYYY)
6 Related Work
Factored Reward MDPs are closely related to Factored MDPs and Successor Features, which have rich history of prior work. Our proposed approach is also related to the prior work on linear approximation of value functions and mixture of experts. Given the focus of our approach is on improving generalization in RI, we briefly summarize the related efforts in the literature.
Factored MDPs (Boutilier et al., 1995 1999 Kearns Koller 1999 Cui & Khardon 2016 are a special
class of MDPs where the state, dynamics, and the reward can be factored into variables. Kearns & Koller
(1999) proposed modeling the factorized transition function as a Dynamic Bayesian Network (DBN) where the
underlying structure is known upfront, but the parameters are to be inferred. Guestrin et al. (2003) proposed
two algorithms, based on linear programming and policy iteration, for learning in Factored MDPs. More
recently, Factored MDPs have been used in conjunction with deep RL approaches in environments known to
have factored structure and shown to improve sample efficiency and transfer (Li & Czarnecki 2018 Wang
et al. 2018 Balaji et al. 2020). However, factored transitions are a strong assumption and do not necessarily give rise to factored value functions. Our work is also related to the work on Exogenous MDPs (Dietterich et al. 2018 which proposes factorizing an MDP into exogenous Markov Reward Process and an endogenous Markov Decision Process. It further shows that if the reward function is assumed to decompose additively into the two factors, the Bellman equation for the original MDP decomposes into two equations: one for the exogenous Markov reward process (Exo-MRP) and the other for the endogenous MDP (EndoMDP). AFaR does not restrict the decomposition to exogenous and endogenous factors and supports using a much larger number of factors.
We operate on Factored Reward MDPs, which, as shown in Scction 2.2 admit factorizod value functions. Such factorization is related to the work on Successor Features Dayan (993); Barreto et al. (2017) There, the key idea is to represent the value functions the linear combination of a special class of basis functions. These basis functions, referred to as the successor features, encode the state such that under a policy, the representation of the state is similar to the representation of the successor states. However, Successor Features are commonly used to obtain fast policy evaluation to aid transfer across different reward functions rather than to improve generalization in the single-task settings Lehnert et al. (017): Barreto et al, (2018). Other works that represent the value functions as a linear combination of basis functions include Parr et al. (2008) that assumes that the transition function is known and uses linear fixed-point methods, Mahadevan & Maggioni (2006), Lagoudakis & Parr (2003) and Konidaris et al. 2011] that used Polynomial basis, proto-value functions and Fourier basis respectively. Koller & Parr (1999 is quite close to our work as they also proposed imposing the additive structure on the value function. However, they use hand-designed basis functions while we learn factorization as part of the training process. Approximating the value function (of a state) as a sum of the value function of its factors can be seen as an application of the mean-field principle (Anderson & Peterson 1987: Peterson 1987 Parisi 1988 Jordan et al.] 1999;
Wainwright & Jordan 2008 Peyrard & Sabbadin 2006). Mean field principle has also been used in context of multi-agent RL (Yang et al. 2018 Subramanian et al. 2020), where the interactions within the population of agents are approximated those between single agent and the average effect from the population of neighboring agents.
Our work is also related to the work on Mixture of Experts (MoE) (Jacobs et al. 1991 Jordan & Jacobs
1994 Chen et al. 1999 Yuksel et al. 2012) where the high level idea is to learn experts that specialize to different aspects of task/data distribution. Aljundi et al. (011 Rannen et al. (017 proposed learning per-task experts for avoiding catastrophic forgetting in the lifelong learning set (McCloskey & Cohhen 1989
Silver et al. 2013 Mitchell et al.] 2018 Sodhani et al., 2022a). MoE architectures
also prominently used in
natural language processing (Shazeer et al. 2017 Lepikhin et al. 2020 Fedus et al. 2021 Lewis et al. 2021
and computer vision (Ahmed et al.) 2016 Gross et al. 2017 Yang et al. 2019 Wang et al. 2020). In the RL
literature, MoE has been used in the context of i-ask learning (Yang et al.) 2020: Sodhani et al. 2021
hierarchical reinforcement learning (Andreas et al. 2017 Goyal et al. 2019]. mulli--ooda policy (Ren et al.,
2021) and multi-agent lcarning (He et al.] (2016). In contrast to these works, we use a MoE-like architecture to map a given MDP to an approximate Factored Reward MDP and train RL agents on this approximate Factored Reward MDP.
13

---

Published in Transactions on Machine Learning Research (MM/YYYY)
Generalization. Given that Factored Reward MDPs are more compact than standard MDPs and exhibit factored value functions, we hypothesize that learning Factored Reward MDP will improve the general- ization performance of downstream RL algorithms. Previous works have focused on different aspects of
generalization: Bellemare ct al. [013)) Tassa et al.] (2018) focused on the singleton environment setup where the train and the Lest. environments are the same while Cobbe et al. (2020); Küttler el al. (2020); Sannvelynn ct al. (2021); Zhang et al. (2021a) focused on the setup where the train and the evaluation distributions
are the samee, even though the instances of train and evaluation environments are different. [Ahmd al.
(2020); Stone et al. (2021) focused on the setup where the train and evaluation distributions differ. Along
different axis, (Sddanii et al. (202bb) focused on the zero-shot generalization setup while Rakelly et al. (2019) allow few gradient updates (or limited interaction with the environment) before evaluating the generalization performance. In this work, we consider both selups - where the train and the evaluation distribution may or may not be the same. We focus on the zero-shot generalization setup as it is a more challenging and realistic sctup. For example, an RL agent deployed in the real world may not know when the underlying task distribution has changed and may have to generalize the new distribution in zero-shot manner.
7 Conclusion
In this work, we propose to study a form of structured environment with additively factorized rewards. We call this setup Factored Reward MDP. nice feature of this structured MDP is that the value function also factorizes. We design an algorithm, AFaR, that learns to map given MDP to an approximate Factored Reward MDP. We show that AFaR can be easily combined with existing RI. algorithms, leading to improved sample efficiency and generalization performance in both Procgen benchmark and the MiniGrid environments. An interesting future work would be to extend AFaR to learn factorized policies that can act on the individual factors that may lead to benefits like better exploration.
Another interesting follow up direction would be to incorporate the factorized structure into a universal policy using hicrarchical approaches. If the learned factored representation correspond to the factors of voriation in an environment, the factors can be used as input to low-level policies, or options (Sutton al. 1999) which can act on these factors independently thus aiding the interpretability of options and enabling improved exploration strategies.
References
Karim Ahmed, Mohammad Haris Baig, and Lorenzo Torresani. Network of experts for large-scale image categorization. In European Conference on Computer Vision, pp. 316-532. Springer, 2016.
Ossama Ahmed, Frederik Träuble, Anirudh Goyal, Alexander Neitz, Yoshua Bengio, Bernhard Schölkopf, Manuel Wüthrich, and Stefan Bauer. Causalworld: A robotic manipulation benchmark for causal structure and transfer learning. arXiv preprint arXiv:2010.0429 2020.
Rahaf Aljundi, Punarjay Chakravarty, and Tinne Tuytelaars. Expert gate: Lifelong learning with a network of experts. In Proceedings of the IEFE Conference on Computer Vision and Pattern Recognition, pp- 3366 3375, 2017.
Eric Allender, Sanjeev Arora, Michael Kearns, Cristopher Moore, and Alexander Russell. A note on the representational incompatibility of function approximation and factored dynamics. Advances in Neural Information Processing Systems, pp. 447-454, 2003.
James R Anderson and Carsten Peterson. A mean field theory learning algorithm for neural networks.
Complez Systems, 1:995 1019, 1987.
Jacob Andreas, Dan Klein, and Sergey Levine. Modular multitask reinforcement learning with policy sketches. In International Conference on Machine Learning, pp. 166-175. PMLR, 2017.
Bharathan Balaji, Petros Christodoulou, Xiaoyu Lu, Byungsoo Jeon, and Jordan Bell-Masterson. Factoredrl: Leveraging factored graphs for deep reinforcement learning. 2020.
14

---

Published in Transactions on Machine Learning Research (MM/YYYY)
Victor Bapst, Alvaro Manchez-Gonzaler Carl Doersch, Kimberly Stachenfeld, Pushmeet Kohli, Peter Battaglia, and Jessica Hamrick. Structured agents for physical construction. In International Conference on Machine Learning, pp. 464 474. PMLR, 2019.
Andre Barreto, Will Dabney, Remi Munos, Jonathan J Hunt, Tom Schaul, Hado P van Hasselt, and David Silver. Successor Features for Transfer in Reinforcement Learning. In I. Guyon, U. V. Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 30. Curran Associates, Tnc., 2017. URL. httpss///proceedings. hewpppooceednng.....rrp......... wwwrips.cc/pppppeeepp
Paters-appr.ppppp pdf
Andre Barrcto, Diana Borsa, John Quan, Tom Schaul, David Silver, Mattco Hessel, Daniel Mankowitz, Augustin Zidek, and Remi Munos. Transfer in deep reinforeement learning using successor features and generalised policy improvement. In International Conference on Machine Learning, pp. 501-510. PML.R, 2018.
Peter W Battaglia, Jessica B Hamrick, Victor Bapst, Alvaro Sanchez-Govzalez, Vinicius Zambaldi, Mateusz Malinowski, Andrea Tacchetti, David Raposo, Adam Santoro, Ryan Faulkner, et al. Relational inductive biases, deep learning, and graph networks. arXiv preprint orXiv:1806. 01261,, 2018.
M. G. Bellemare, Y. Naddaf, J. Veness, and M. Bowling. The arcade learning environment: An evaluation platform for general agents. Journal of Artificial Intelligence Research, 47:253-20 279, jun 2013.
Mare G. Bellemare, Sriram Srinivasan, Georg Ostrovski, Tom Schaul, David Saxton, and Rémi Munos. Unifying count-based exploration and intrinsic motivation. In NIPS, 2016.
Richard Bellman. Dynamic Programming. Princcton University Press, Princeton, NJ, USA, 1 edition, 1957.
Craig Boutilier, Richard Dearden, Moisés Goldszmidt, et al. Exploiting structure in policy construction. In IJCAI, volume 14, pp. 1104-1113, 1995.
Craig Boutilier, Thomas Dean, and Steve Hanks. Decision-theoretie planning: Structural assumptions and computational leverage. Journal of Artificial Intelligence Research, 11:1-94, 1999.
Yuri Burda, Harri Edwards, Deepak Pathak, Amos Storkey, Trevor Darrell, and Alexei A. Efros. Large-scale study of curiosity-driven learning. In International Conference on Learning Representations, 2019. URL
htttss:/openreevie wet.:oormm/.drrrl///aaooao http:s//ppppuuueoe...eeeee/orrrr
Ke Chen, Lei Xu, and Huisheng Chi. Improved learning algorithms for mixture of experts in multiclass classification. Newral networks, 12(9):1229-122)) 1999.
Maxime Chevalier-Boisverr Lucas Willems, and Suman Pal. Minimalistic gridworld environment for openai gym. http:///ggttuuub.cooaaaineeeeeee
2018.
Karl Cobbe, Chris Hesse, Jacob Hilton, and John Schulman. Leveraging procedural generation to benchmark reinforcement learning. In International conference on machine learning, pp. 2048 2056. PMLR, 2020.
Hao Cui and Roni Khardon. Online symbolic gradient-based optimization for factored action mdps. In IJCAT, pp. 3075 3081, 2016.
Poter Dayan. Improving generalization for temporal difference learning: The successor representation. Neural Compulation, 5(4):613 624, 1993.
Thomas Dietterich, George Trimponias, and Zhitang Chen. Discovering and removing exogenous state variables and rewards for reinforcement learning. Tn International Conference on Machine Learning, pp- 1262-1270. PMLR, 2018.
Lasse Espcholt, Hubert Soyer, Remi Munos, Karen Simonyan, Vlad Mnih, Tom Ward, Yotam Doron, Vlad Firoiu, Tim Harley, lain Dunning, et al. Impala: Scalable distributed deep-rl with importance weighted actor-learner architectures. In International Conference Machine Learning, pp- 1407-1416: PML.R, 2018.
15

---

Published in Transactions on Machine Learning Research (MM/YYYY)
William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. arviv preprint arXiv:2001.022 003961, 2021.
Dibya Ghosh, Jad Rahme, Aviral Kumar, Amy Zhang, Ryan P Adams, and Sergey Levine. Why generalization in rl difficult: Epistemic pomdps and implicit partial observability. Advances in Neural Information Processing Systems, 34, 2021.
Anirudh Goyal, Shagun Sodhani, Jonathan Binas, Xue Bin Peng, Sergey Levine, and Yoshua Bengio. Reinforcement learning with competitive ensembles of informattion-constrainee primitives. arXiv preprinl arXiv:1906. 2019.
Sam Gross, Mare' Aurelio Ranzato, and Arthur Szlam. Hard mixtures of experts for large scale weakly supervised vision. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pp. 6865-6873, 2017.
Carlos Guestrin, Daphne Koller, Ronald Parr, and Shobha Venkataraman. Efficient solution algorithms for factored mdps. Journal of Artificial Intelligence Research, 19:399 468,,68888 2003.
Charles R. Harris, K. Jarrod Millman, St'efan J. van der Walt, Ralf Gommers, Pauli Virtanen, David Cournapeau, Eric Wieser, Julian Taylor, Sebastian Berg, Nathaniel J. Smith, Robert Kern, Matti Picus, Stephan Hoyer, Marten H. van Kerkwijk, Matthew Brett, Allan Haldane, Jaime Fern andee del R'io, Mark Wiebe, Pearu Peterson, Pierre G' G'erard-Marchant, Kevin Sheppard, Tyler Reddy, Warren Weckesser, Hameer Abbasi, Christoph Gohlke, and Travis Oliphant. Array programming with NumPy. Nature, 585(7825):351 362, September 2020. doi: (0.1038/s41556-022 1215866-20-2649-11
URL https: http:///doooogg//.../ otg/10.1038..
He He, Jordan Boyd-Graber, Kevin Kwok, and Hal Daumé III. Opponent modeling in deep reinforcement learning. In International conference on machine learning, pp. 1804-1813. PMLR, 2016.
Victoria J Hodge, Richard Hawkins, and Rob Alexander. Deep reinforcement learning for drone navigation using sensor data. Neural Computing and Applications, 3366) 2021.
Robert A Jacobs, Michael Jordan, Steven Nowlan, and Geoffrey E Hinton. Adaptive mixtures of local experts. Neural computation, 3(1):79-87, 1991.
Michael I Jordan and Robert A Jacobs. Hierarchical mixtures of experts and the algorithm. Neural computation, 6(2):181-214 1994.
Michael I Jordan, Zoubin Ghahramani, Tommi S Jaakkola, and Lawrence K Saul. An introduction to variational methods for graphical models. Machine learning, 38(2):183-25)) 233, 1999.
Michael Kearns and Daphne Koller. Efficient reinforcement learning in factored mdps. In IJCAL, volume 16, pp. 740 747, 1999.
Daphne Koller and Ronald Parr. Computing factored value functions for policies in structured mdps. In IJCAI, volume 99, pp. 1332 1339, 1999.
George Konidaris, Sarah Osentoski, and Philip Thomas. Value function approximation in reinforcement learning using the fourier basis. In Twenly-fifth AAAT conference on artificial intelligence, 2011.
Nya Kostrikov, Denis Yarats, and Rob Fergus. Image augmentation is all you need: Regularizing deep reinforcement learning from pixels. arXiv preprint 2tt:20044.16443 2020.
Heinrich Küttler, Nantas Nardelli, Alexander H. Miller, Roberta Raileam,, Marco Selvatici, Edward Grefen- stette, and Tim Rocktäschel. The NetHack Learning Environment. In Proceedings of the Conference on Neural Information Processing Systems (NewrIPS), 2020.
Michail G Lagoudakis and Ronald Parr. Least-squares policy iteration. The Journal of Machine Learning Research, 4:1 4:1107-11440 2003.
16

---

Published in Transactions on Machine Learning Research (MM/YYYY)
Michael Laskin, Kimin Lee, Adam Stooke, Lerrel Pinto, Picter Abbeel, and Aravind Srinivas. Reinforcement learning with augmented data. arXie preprint orXiv:2004. 14990, 2020.
Lucas Lehnert, Stefanie Tellex, and Michael L. Littman. Advantages and limitations of using successor features for transfer in reinforcement learning. arXiv preprint orXiv:1708.001 00100, 2017.
Dmitry Lepikhin, HyoukJoong Lee, Yuanzhong Xu, Dehao Chen, Orhan Firat, Yanping Huang, Maxim Krikun, Noam Shazeer, and Zhifeng Chen. Gshard: Scaling giant models with conditional computation and automatic sharding. ar preprinl urXiv:2006.1666, 2020.
Mike Lewis, Shruti Bhosale, Tim Dettmers, Naman Goyal, and Luke Zettlemoyer. Base layers: Simplifying training of large, sparse models. Xiv preprinl orXiv:2003. (0n:2203.1676,,, 2021.
Changjian Li and Kraysztof Czarnecki. Urban driving with multi-objective deep reinforcement learning. arXiv preprint orXii:1811.0880 2018
Paolo Liberatore. The size of mdp factored policies. In AAAI/IAAI, pp. 267 272, 2002.
Sridhar Mahadevan and Mauro Maggioni. Value function approximation with diffusion wavelets and laplacian eigenfunctions. Advances in neural information processing systems, 18:843, 2006.
Michael McCloskey and Neal J Cohen. Catastrophic, interference in connectionist networks: The sequential learning problem. In Psychology of learning and motivation, volume 24, pp. 109 165. Elsevier, 1989.
Tom Mitchell, William Cohen, Estevam Hruschka, Partha Talukdar, Bishan Yang, Justin Betteridge, Andrew Carlson, Bhavana Dalvi, Matt Gardner, Bryan Kisiel, et al. Never-ending learning. Communications of the ACM, (1(5):103-11) 2018.
The pandas development team. pandas-dev/pandass Pandas, February 2020. URL https://doi.org/// 5281/zenodo. 3509134 3509134
G Parisi. Statistical field (adoro,,, 1988.
Ronald Parr, Lihong Li, Gavin Taylor, Christopher Painter-Wakefiee and Michael L. Litiman. An analysis of linear models, linear value-function approximation, and feature selection for reinforcement learning. In Proceedings of the 25th international conference on Machine learning, pp. 752 759, 2008.
Adam Paszke, Sam Gross, Francisco Massa, Adamn Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, al. Pytorch: An imperative style, high-performanee deep learning library. arXiv preprint Xiv: 1912.01 703, 2019.
Deepak Pathak, Pulkit Agrawal, Alexei A Efros, and Trevor Darrell. Curiosity-driven exploration by self-supervised prediction. In International Conference on Machine Learning, pp. 2778 2787. PMLR, 2017.
Carsten Peterson. A mean field theory learning algorithm for neural networks. Complex systems, 1:995 1019, 1987.
Nathalie Peyrard and Régis Sabbadin. Mean field approximation of the policy iteration algorithm for graph- based markov decision processcs. In Proceedings of the 2006 Conference on ECAI 2006: 17th Europeon Conference Artificial Intelligence August 29 September 1, 2006, Riva Del Garda, Italy, pp- 595-599, NLD, 2006. IOS Press. ISBN 1586036424.
Martin L Puterman. Markov decision processes: Discrete stochastic dynamic programming. Journal of the Operational Research Society, 1995.
Roberta Railcanu and R. Fergus. Decoupling value and policy for generalization in reinforcement learning. ArXiv, abs/2102.103300, 2021.
Roberta Raileanu and Tim Rocktäschel. Ride: Rewarding impact-driven exploration for sroeedurally-gnnerate environments. Xiv preprint ar Xiv:2000. 12292, 2020.
17

---

Published in Transactions on Machine Learning Research (MM/YYYY)
Roberta Raileanu, Maxwell Goldstein, Denis Yarats, Ilya Kostrikov, and Rob Fergus. augmentation for generalization in reinforcement learning. 2020.
Automatic data
Kate Rakelly, Aurick Zhou, Chelsea Finn, Sergey Levine, and Deirdre Quillen. Efficient off-policy meta- reinforcement learning via probabilistic context variables. In International conference on machine learning, pp. 5331-5340 PMLR, 2019.
Amal Rannen, Rahaf Aljundi, Matthew B Blaschko, and Tinne Tuytelaars. Encoder based lifelong learning. In Proceedings of the IEEE International Conference on Computer Vision, pp. 1320-1328, 2017.
Jie Ren, Yewen Li, Zihan Ding, Wei Pan, and Hao Dong. Probabilistic mixture-of-experts for efficient deep reinforcement learning. arXiv preprint arXi::2104 2021.
Mikayel Samvelyan, Robert Kirk, Vitaly Kurin, Jack Parker-Holder, Minqi Jiang, Eric Hambro, Fabio Petroni, Heinrich Kuttler, Edward Grefenstette, and Tim Rocktäschel. Minihack the planct: A sandbox for open-ended reinforcemenl learning research. 2021.
Alvaro Gonzale, Jonathan Godwin, Tobias Pfaff, Rex Ying, Jure Leskovee, and Peter Battaglia. Learning to simulate complex physics with graph networks. In International Conference on Machine Learning, pp. 8459-8468. PMLR, 2020.
Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc Le, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely-gated mixture-of-experts layer. or preprint 1XXi:1701.06388 2017.
Daniel L Silver, Qiang Yang, and Lianghao Li. Lifelong machine learning systems: Beyond learning algorithms. In 2013 AAAI spring symposium series, 2013.
Shagun Sodhani, Amy Zhang, and Joelle Pineau. Multi-task reinforcement learning with context-based representations. In International Conference on Machine Learning (ICML), 2021.
Shagun Sodhani, Mojtaba Faramarzi, Sanket Vaibhav Mehta, Pranshu Malviya, Mohamed Abdelsalam, Janarthanan Janarthanan, and Sarath Chandar. An introduction to lifelong supervised learning. arXiv preprinl ar xrXiv:2207.4333 2022a.
Shagun Sodhani, Franziska Meier, Joelle Pincau, and Amy Zhang. Block contextual MDPs for continual learning. In Proceedings of the 4th Annual Conference Learning for Dynamics and Control, L4DC 2022, Stanford, CA, USA, 23-24 June 2020, Proceedings of Machine Learning Research. PMLR, 2022b.
Xingyou Song, Yiding Jiang, Stephen Tu, Yilun Du, and Behnam Neyshabur. Observational overfitting in reinforcement learning. arXio preprint or Xin: 2. 02975, 2019.
Aravind Srinivas, Michael Laskin, and Pieter Abbeel. Curl: Contrastive unsupervised representations for reinforcement learning. arXiv preprint orXiv:2004.04130 2020.
Austin Stone, Oscar Ramirez, Kurt Konolige, and Rico Jonschkowski. The distracting control suite a challenging benchmark for reinforcement learning from pixels. arXiv preprint arX1::100.002722, 2021.
Student. The probable error of mean. Biometríko, pp. 1-25, 1908.
Sriram Ganapathi Subramanian, Matthew Taylor, Mark Crowley, and Pascal Poupart. Partially observable mean field reinforcement learning. Xi preprint arXiv:2012. 15791, 2020.
Richard S. Sutton, Doina Precup, and Satinder Singh. Between MDPs and semi-MDPs: framework for temporal abstraction in reinforcement learning. Artificial Intelligence, 112(1):181 -11, 1999. ISSN 0004-3702. doi: http:://ooaaooog///....oo0000000 sccaacee/aticcee//1///4400000550
URL hwt.ittp:://ww.eceenneeeeee
Yuval Tassa, Yotam Doron, Alistair Muldal, Tom Erez, Yazhe Li, Diego de Las Casas, David Budden, Abbas Abdolmalcki, Josh Merel, Andrew Lefrancq, al. Deepmind control suite, arXiv preprint 101600080..00690, 2018.
18

---

Published in Transactions on Machine Learning Research (MM/YYYY)
Martin J Wainwright and Michael Irwin Jordan. Graphical models, erponential families, and variational inference. Now Publishers Inc, 2008.
Tingwu Wang, Renjie Liao, Jimmy Ba, and Sanja Fidler. Nervenet: Learning structured policy with graph neural networks. In International Conference on Learning Representations, 2018.
Xin Wang, Fisher Yu, Lisa Dunlap, Yi-An Ma, Ruth Wang, Azalia Mirhoseini, Trevor Darrell, and Joseph E Gonzalez. Deep mixture of experts via shallow embedding. In Uncertainty in Artificial Intelligence, pp. 552-562. PMLR, 2020.
Omry Yadan. Hydra framework for clegantly configuring complex applications. Github, 2019. URL https://github. http///acccbboooesraraaaaaaaaaea
Brandon Yang, Gabriel Bender, Quoc V Lc, and Jiquan Ngiam. Condconv: Conditionally paramcterized convolutions for elficient inference. ar preprind orXix:190..4499 2019.
Ruihan Yang. Huazhe Xu, Yi Wu, and Xiaolong Wang. Multi-task reinforcement learning with soft modular- ization. Xrviv preprint or 200: 13661, 2020.
Yaodong Yang, Rui Luo, Minne Li, Ming Zhou, Weinan Zhang, and Jun Wang. Mean field multi-agent reinforcement learning. In International conference on machine learning, pp. 5571 5580. PMLR, 2018.
Seniha Esen Yuksel, Joseph N Wilson, and Paul Gader. Twenty years of mixture of experts. IEEE transactions on neural networks and learning systems, 23(8)11177 1193, 2012.
Amy Zhang, Nicolas Ballas, and Joelle Pineau. A dissection of overfitting and generalization in continuous reinforcement learning. arXiv preprint Xrvi 079707737, 2018a.
Amy Zhang, Shagun Sodhani, Khimya Khetarpal, and Joelle Pineau. Learning robust state abstractions for hidden-parametter block MDPs. In International Conference Learning Representations, 2021a. URL
https:://penrrrriew http://wweeeeoorrrrrreeeeoootaaa
Chiyuan Zhang, Oriol Vinyals, Remi Munos, and Samy Bengio. study on overfitting in deep reinforcement learning. ar preprint 1xXx::884.0688) 2018b.
Weitong Zhang, Jiafan He, Dongruo Zhou, Amy Zhang, and Quanquan Gu. Provably efficient representation learning in low-rank markov decision processes. arXiv preprint Niw:2106.11195 2021b.
A Q-Functions Under Factorized Rewards
We show that Factored Reward MDP has state value functions and state-action value functions that can be written as a sum of factors.
Proposition 2.2 Consider Factored Reward MDP where the state st is factorized into k state factors st, ...... s and the reward function is factorized as R(s,, 0z) = Σt= r(st, where S x A - T.. are local reward functions. The state value function and the state-action value functions for any policy can be written as:
4 V=(81) = (F())
(4)
1=1
(F(s
Σ
(7).
(5)
19

---

Published in Transactions on Machine Learning Research (MM/YYN)
where
20
(*((s)
Σ
?
(7 a1) = r(s), a1) +
Σ
T(s ar br (r)rrsk
e=t+1
can be seen as state value and the state-action value functions for the 8th factor, respectively.
Proof. The state-value function is defined as:
(*((s) = Σ e' Σ
T(s++1 10se+11l,,,00 ar) R(sv, ar).
(6)
M=t
62'41
Applying the reward factorization, we get
DC
(*((t) = Σx Σ
T(se41188 0e) Σ ac)
(7)
p=
Re-arranging the terms, we get
(*t(s) Σ
Σ 8,r41
T(sv-ils,, (e)r,,,,
(8)
1=1 p=t
Equation
can be re-written as:
k V*(ss) Σ (st)
(9)
where (s)) = Ef Eny,, 4(s ae (r( ap) can be seen as the state value function for the 7th factor. Similarly, can obtain the following factorization for the stale-action value function:
(F(s = Σ Q (z)
(10)
where
(s), = r(5), ax) +
Σ
Σ T(s84118m ar (r
e't+11
and can be seen as the state-action value function for the 5th factor.
B Environments
We use the Procgen benchmark (Cobbe al. 2020) and MiniGrid environment (Chevalier Boccaalie-BBoivertt al. 2018) to evaluate the effectiveness of the proposed algorithm.
20

---

Published in Transactions on Machine Learning Research
Figure Screenshols for KeyCorridorS3R1-10 KeyCorridorS3R2-10, and KeyCorridorS3R3-00 environments (left. to right). Taken from Chevalier- Boisvert et al. (2018)
B.1 Procgen Benchmark
The Procgen benchmark comprises 16 procedurally generated environments designed to evaluate RL algorithms for generalization and sample elficiency. Each enviromment represents distribution of levels where each level corresponds to a new environment, allowing for a strict partition between train and test environments. The benchmark satisfics several important critoria:
High diversity in the environments
Optimized environments that enable fast evaluation of algorithums
Control over parameters like difficulty and need for exploration. All environments support modes casy and hard.
4. Emphasize on visual recognition and motor control
5. Level solvability with about 99% levels solvable for each environment
Shared observation and action space
B.1.1 Observation Space
The observation space is an RGB image of size 64 64 3.
B.1.2
Action Space
All environments use discrete action space with 15 actions. In some environments, a no-op action is also present. All environments use a discrote 15-dimensional action space and produce 64 x 64 x 3 RGB observations. We use Procgen's 's easy setup, so for each game, agents are trained on 200 levels and lested on the full distribution of levels. More details about our experimental setup and hyperparameters can be found in Appendix E.
B.2 MiniGrid Environment
Minimalistic Gridworld Environment (MiniGrid) is an open source grid world Gym environment 0 enables the creation of gridworld environments with adjustable parameters (like size and layout of the grid) with different objects of different types, colors, and behaviors.
21

---

Published in Transactions on Machine Learning Research (MM/YYYY)
Table Hyperparameter values that are common across all the experiments with Procgen Benchmark
Hyperparameter
Hyperparameter
values
optimizer
RMSProp
learning rate
5 x 10-1
alpha for optimizer
0.99
epsilon for optimizer discount.
10-5
0.999
gae lambda, parameter entropy coefficient
0.95
0.01
value loss coefficient max grad norm
0.5
0.5
number training processes number of forward steps in A2C
64
256
number of PPO epochs number of minibatches for PPO
3
8
clip parameter for PPO number of euvironment steps to train on distribution mode for environment number of levels
0.2
25 106
easy
10, 100, 200,
300
augmentation type
crop
coefficient for the augmentation loss
0.1
Table 6: Hyperparameter values that specific the AFaR method, when trained Procgen Benchmark
Hyperparameter
Hyperparameter values
number encoders
2, 4, 6, 8
B.2.1 Structure of the world
The "world" is a N x M grid of tiles. Each tile contains zero or one object. Each object has color (discrete value) and type (discrete value). Objects include wall, door, key, box etc. The agent can perform interactions like "pick" or 'carry" an object with some constraints. For example, the agent can carry at most one object at a time. Agent must carry a key (of matching color) to open a door.
B.2.2 Observation Space
Each tile (in the grid world) is encoded as 3 dimensional tuple of integers of the form (object_ id, color_oora id, state). The obsorvations (for the agent) are egocentric and partial. i.c., the agent views the 3 x 3 sized grid in front of it. Moreover, the agent can not see across obstacles (like walls or closed doors).
B.2.3 Action Space
The action space discrete, with actions: turn left, turn right, move forward, pick up an object, drop the object being carried, toggle (open doors, interact with objects) and done (task completed, optional).
22

---

Published in Transactions on Machine Learning Research (MM/YYYY)
Table Best performing values for the number of encoders for the different environments and baselines, when trained
Procgen Benchmark
Environment
AFaR with DrAC baseline
AFaR with IDAAC baseline
Bigfish
4
Bossfight
6
6
Caveflyer
2
2
Chaser
4
4
Climber
4
4
Coinrun
4
4
Dodgeball
6
6
Fruitbot
6
Jumper
4
4
Leaper
6
6
Maze
4
4
Miner
6
6
Ninja
4
4
Plunder
2
2
Starpilot
6
6
B.2.4 Reward Function
The agent is given a sparse reward for reaching the green goal tile. If the agent does not reach the tile within a specified time, gets zero rewards. In case the agent completes the task, it gets a reward of 0.9 step count max steps ie., the agent is rewarded more when completes the task faster.
Table 8: Hyperparameter values that common across all the experiments with MiniGrid environments
Hyperparameter
Hyperparameter
values
total frames
30000000
entropy cost
0.0005
intrinsic reward coefficient number of imput frames
0.1
1
number of actors batch size unroll length max grad norm entropy cost
40
32
100
40.0
0.001
baseline cost.
0.5
learning rate alpha (for optimizer)
0.0001
0.99
momentum (for optimizer)
0.0
epsilon (for optimizer)
10-5
forward loss coefficient (for explo-
10.0
ration)
inverse loss coefficient (for explo- 0.1
ration)
intrinsic reward coefficient (for explo- 0.5
ration)
random loss coefficient (for explo- 0.1
ration)
23

---

Published in Transactions on Machine Learning Research (MM/YYYY)
Table Hyperparameter values that are specific the AFaR method, when trained on MiniGrid environments
Hyperparameter
Hyperparameter values
mumber of encoders
2, 4, 6, 8
Bossfight
Caveflyer
Chaser
Bigfish
pliit
$5.0
5205
Stare
2.5
2.0
1.2
267
Climber
Coinrun
Dedgeball
Fruithot
D-AC::
DiAC Afaft
Afull
Spore
Emtc
1.0
3.5
200
255
3.0
0.5
25
Staps 10P)
Staps (K11
Stapo
Stago (x12P)
Heist
Jumper
Leaper
Maze
Denc
APail
0.0
0.5
1.0
2.5
2-3
255
0.0
3.5
3,0
3.5
2.0
2.
3.0
0.5
3.0
2.2
2.5
0.3
0.5
1.3
1.5
3.0
35
Steas (x.)
Staps
16100")
Stops
204
Steas
(x129)
Miner
Ninja
Plunder
Starpilot
DrAC
DeAC Als.h
Drsic
Scon
0.0
3.2
2.5
167
Figure 6: We compare the performance of AFaR with DrAC (as the baseline approach) on the training environments from Procgen Benchmark when training with 200 levels (per euviroument) for 25M steps. The solid curves correspond to the mean performance (score), while the shaded regions represent the standard error (computed over ten runs). There is a substantial performance improvement in 4 environments, and the performance of the proposed approach is close to the baseline approach in the remaining environments. For the performance on the evalnation environments, refer to Figure
In tables 6 and we describe the space of values that explored for &. We picked the best value of k by comparing the performance over 3 runs (seeds) for different values (of k). Then re-ran 10 experiments (sceds) with the best value of and reported the corresponding results in the paper. The best performing values of l are given in the table 7 for the ProcGen environments and was 4 for the MiniGrid environments. We note the values are very consistent across the implementations, with the & = 4 being the most common value. The case where we nced more than encoders are environments like DodgeBall where the agent can interact with several 'objects" (or factors). We also note that in several environments (like FruitBot Leaper), the performance across different values of l is quite close.
24

---

Published in Transactions on Machine Learning Research (MM/YYY)
Bigfish
Bossfight
Caveflyer
Chaser
DJAC
DrAC Aiax
Herr)
2.0
2.0
1.2
12
Steps b=107)
Shape bx107)
Steps (x121)
Stepe (53
Climber
Coinrun
Dodgehall
Fruitbot
Deyic Aftul.
Cenc
Scara
A/a4
0.5
1.0
2.5
2,0
2.0
2.5
2.0
0.5
2.5
2.0
Stupa (x:00))
Stops (x1177)
Stopo
201
Steps: 1-10f1
Jumper
Leaper
D:CC
Ceti
4m
Afrht
Scre
My lnd Hayer
2.0
0.5
2.0
2.3
0.3
1.5
3.5
2.5
2.0
00
05
3.0
2.
Siops
Stago
(x109)
Stapo
Si4pa
Ninja
Starpilot
DVAC
10
DeAC
Aflan
Araft
Altil.
SCOR
4.0-
A/a
0.0
0.5
3.0
3.5
2.3
2.5
6.3
2-5
3.5
2.5
0.0
0.5
3.0
2.5
2.0
25
2.0
0.5
3.5
23
Stops
Stago (+30%)
0x10%)
Stopa
Figure 7: We compare the performance of AFaR with DrAC (as the baseline approach) on the evaluation environments from Procgen Benchmark when training with 200 levels (per environment) for 25M steps. The solid curves correspond to the mean performance (score), while the shaded regions represent the standard error (compuled over ten runs). There is substantial performance improvement in euvironments, and the performance of the proposed approach is close to the baseline approach in the remaining euvironments. For the performance (on the evaluation environments the end of the training, refer to Table T and for performance the training environments, refer Figure
25

---

Published in Transactions on Machine Learning Research
Bigfish
Boostight
4.0
Caveflyer
Afrs.
20
0.3
20
0.0
3.5
3,0
3.5
2.0
23
0.0
0.5
1.5
2.2
0.5
1.0
2.0
25.
Steps
(x10P)
Stepe 1x50%
Coinrun
Dodgeball
Di-cc
best
1.0
Aral
2.0
0.5
2.0
20
2.5
2.0
235
0.0
0.5
1.5
2.2
2.0
0.5
3.5
23
25
Sispa
Stope (H12P)
S14g4
Jumper
Leaper
D-AC
CIAC AraB
CHAC Aral
Aral.
3.4
28
DvAC
3.0
41.0
2.5
Miner
Ninia
Plunder
Starpilot
Afo/t
DV00
4.6
4.6
Aryak
Hew
3.5
2.0
2.5
Steps 0x:009
(x107)
367
(x127)
1.7
Stege (s3
Figure 8: We compare the performance of AFaR with DrAC (as the baseline approach) the evaluation environments from Procgen Benchmark when training with levels (per environment) for 25M steps. The solid curves correspond to the mean performance (score), while the shaded regions represent the standard error (computed ten runs). For corresponding curves when training on 200 levels per environment, refer to Figure
26

---

Published in Transactions on Machine Learning Research (MM/YYY)
WindrrikkeyCCrriddrSS3---
MintGridkkeyCoriddoSSRLL-
Minlord-KeyCordddSSRR22-
RIDE
RIDE
0.8
ARaR
0.8
Mean Episode Retur
Afall-mean
Afaltmenn
Afall-men
0.6
0.6
Mean Eptsode
Mean Fpicode
0.4
0.2
0.0
0.0
0.2
0.4
0.6
1.0
1.2
0.4
0.5
0.8
1.0
Steps
167
Figure 9: We compare the performance of AFaR with RIDE (baseline approach) KeyCorriddrSSR3-00 KeyCorridorS3R112 and KeyCorr1dor99R2-00 environments (left to right) with the number of environment in- teractions (in millions) along the x-axis and the mean episodic return along the y-axis. We also include ablation version of the AFaR algorithm where we disable the Atlention Module and assign equal weights to all the factors while uggregating them. We denote this ablation AFaltmmean. Note that while the agent evaluated three enviromments (KeyCorridors333-vd KeyCorriag KeyCorridor $332r KeyCorridors3R1-00, it was trained only on the KeyCorridorS3R3- environment, thus the evaluation KeyCorridorS3R2-vo and KeyCorridors3R1-0 environments is done zero-shot manner. For all the three environments, A FaR obtains the best sample efficiency.
C
Additional Implementation Details
C.1 Libraries
We use the following open-source libraries:
1. PyTorch
Paszke et al.
2019
Hydra Yadan 2019
3. Numpy Harris ct al.
2020 0
4. Pandas pandas development leam
2020 10
RIDE Implementation
Raileanu Rocktäschel
2020
6. DrAC Implementation Raileanu et al.
2020 12
IDAAC Implementation Raileanu & Fergus 2021
13
In case feedforward layers are used process the output convolutional encoders, we can additionally vary the size of the intermediate layers in the feedforward networks (corresponding to the different factors). Using variable-sized networks can further help the network to model various factors more effectively.
Testing for statistical significance
We perform two-tailed, Student's t-distribution test Student 1908) under equal sample sizes, unequal variance setup (also called Welch's is 1-test). The null hypothesis is: that the mean performance of the two models (AFaR and any baseline) are equal. The significance level (p) is set to 0.05.
https://pytorch. org/
https://github com/ wwwp:tttuu..oommcaeebteeeeeeeeee comf/Eecebokkreeearrc//yyyz
https://numpy.
https://pandas.
https:///tthbb..cmm/rrttt
httpp::/ggthub..cmmrraaaaaa
27

---

Published in Transactions on Machine Learning Research (MM/YYY)
E Baselines
E.1 Choice of Actor-Critic Methods
In Section we noted that AFaR is representation learning algorithm, and we need to pair it with a policy oplimization algorithm for end-to-end reinforcement learning. We consider three actor-critic algorithms that are shown to obtaín state-of-the-art performance in the given environments. For the Procgen environments, we use two basclines: (i) Data-regularized Actor-Critic (DrAC) Railcanu et al. (2020), and (ii) Invariant Decoupled Advantage Actor-Critie (IDAAC) Raileanu & Fergus (2021). The key dilference bet these two approaches is that IDAAC is the state-of-the-art among methods that use a separate policy and value function, while DrAC is the state-of-the-art among methods that use shared network for policy and value functions. For MiniGrid environments, we use Rewarding Impact-Driven Exploration (RIDE) Raileanu & Rocktäschel 2020 method, a state-of-the-art method for tasks on the MiniGrid environments.
There are several benefits of using these diverse methods as a baseline for AFaR. These methods focus on different challenges in training RL agents and use different inductive biases and approaches to achieve good performance, using exploration, disjoint policy and value functions, and data augmentation. Moreover, these methods also offer diversity in how the encoders and value functions value function used lo provide the learning signal for an advanlage function, which is used for training the policy (thus acting as the critic). The advantage function also shares parameters with the actor network as is done in standard actor-critic architectures. i.c. in IDAAC, the actor and the advantage function use the same encoder. Since the advatnage function plays the role of critic, we factorize the advantage function. RIDE uses two sets of encoders one for training the policy and the her for computing the intrinsic rewards. We factorize only the encoder used for training the policy, Integrating AFaR with diverse baselines and demonstrating that AFaR can improve over these baselines shows that our proposed algorithm is helpful for various actor-critic algorithms.
trained. For example, IDAAC learns
Next, we describe the different actor-critic algorithms in detail.
DrAC builds on the idea of using data augmentation for improving generalization in RL (Laskin et al. 2020 Srinivas et al. 2020 Kostrikov et al. 2020), and proposes two novel regularization terms for the policy and value function
make the use of data augmentation theoretically sound for actor-critic algorithms. Specifically, given an image transformation f, the following two regularization terms are added to the actor-critic loss: G2 KL(To(aass) (e( and G, (Vo(s) va() where v are the parameters of f(). These additional terms ensure that the learning agent's 's policy and value functions are invariant to the transformations induced by the augmentations.
MiniGrid environments use the sparse reward setup, making exploration a key challenge. RIDE proposed a novel intrinsic reward where the agent rewarded for taking actions that affect its learned state representation. Specifically, the intrinsic reward is computed as the L2 norm of the difference between the representation of the consecutive states, i.e. || (8++1) ((s)]l-- Episodic state visitation counts discount this reward to ensure that the agent does not simply go back and forth between the same two states. The overall intrinsic reward is given as: rf [((+++))=((+))] where (s) is the number of time state visited in episode e. The (N1p( agent is trained to maximize the weighted sum of the intrinsic reward (rt) and the extrinsic reward (obtained from the environment). On the MiniGrid environments, RIDE is reported to improve sample efficiency as compared to approaches like count-based exploration Bellemare et al. (2016), Random Network Distillation (uses prediction error of a random network as the exploration bonus) Burda et al. (2019), Intrinsic Curiosity Module (uses a curiosity basod intrinsic reward) Pathak et al. (2017) and IMPALA (an actor-critic approach that uses only external rewards) (Espeholt el. al.) 2018).
28


Reviews:
[{"review_id": "9nCL6AKKUK", "reviewer": "Reviewer_fNHM", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths: \n- paper tackles an important and yet understudied problem in RL\n- paper is easy to follow; technical contributions are laid out clearly\n- Evaluation is thorough - algorithm is tested on a range of challenging benchmark problems and it consistently seems to do better than its baseline counterparts. \n\nWeaknesses:\n- It is not clear to me how we can be confident that the proposed architecture is actually learning a factored representation.\n- can the aggregation function on the actor side be summarized as scaled dot-product attention? Or is different in a way that I am missing?\n- The discussion about successor features seemed a little weak to me - are there any reasons why we cannot use successor features in Procgen or Minigrid in the same way that we are using the proposed algorithm? \n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper studies the important problem of learning factored representations in RL; rather than assuming factored MDPs, this paper attempts to learn a factored approximation of the underlying MDP. They propose an attention-based architecture that learns factored value functions that when summed together, yield the agent’s current value function $\\hat{V*}$. The learned factors in conjunction with their corresponding values are used as the input representation to a policy network, which is trained using a classic actor loss. This modular architecture outperforms baseline actor-critic algorithms on procgen and minigrid.", "claims_and_evidence_yn": "", "claims_explanation": "", "audience_yn": "", "audience_explanation": "", "requested_changes": "- Ablations: actor-critic baseline with the same number of parameters as AFAR and actor-critic with attention. Such an ablation will strengthen the paper.\n- A discussion about Hybrid Reward Architecture (HRA) from Van Seijen et al and how it fits into the factored MDP landscape; this will also strengthen the paper.\n- Some experiments (maybe qualitative) that showcase that the proposed architecture is learning meaningful factors that wouldn’t automatically be learned by a vanilla actor-critic agent; I think such experiments will significantly improve the paper", "broader_impact_concerns": "I dont have any concerns about the broader impact of this work", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "wDgtMigpda", "reviewer": "Reviewer_pydn", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "# Strengths:\n- S1. The paper is generally well constructed, and the overall motivation and main line of research is clear to the reader.\n- S2. The algorithm description and architecture diagram were useful in understanding the approach proposed.\n- S3. The paper attempts to do proper significance testing (using Welch's t-test) for some of the comparisons.\n\n\n# Weaknesses:\n\nGenerally, I like to organize my critiques of a paper based on three criteria: overall clarity of the paper, consistency between the different parts of the paper, well organized and motivated empirical evaluation of the stated research questions, and conclusions which are supported by the provided work.\n\n## Clarity:\n\nOverall, the paper is clear but there are a few doubts on how the algorithm works in detail.\n\nCL-1). Step d of your architecture and line 8 of the algorithm are a bit confusing. For each factored state $z^i$ is there a separate corresponding function $V_i^\\pi(z^i)$ to construct the factored value function $V_i^\\pi$, or is there a single parameterized function $V^\\pi(z^i)$ which constructs each value function $V_i^\\pi$? Perhaps a more concrete way of asking is \"are the weights for the mapping $V: Z -> \\mathbb{R}$ shared for each factored state or different for each factored state\"?\n\nCL-2). The second paragraph of section 3 feels a bit too broad, and muddles the motivation for the new architecture. Some questions I think are worth considering if the authors want to motivate their approach through the mean-field principle:\n   CL-2.1). Is the factorization of the value function of a un-factorized state into a sum of value functions of factorized states the only possible application of the mean-field principle in this setting?\n   CL-2.2). Are we testing the mean-field principle/approximation generally in this setting, or are we testing one particular factorization? If we are testing one particular factorization, I think making the connection to the mean-field principle in the intro/related works would be an interesting addition with a more concrete discussion.\n\nCL-3). You state the hypothesis: \"If the factorization is optimized to match the value function, the mean-field approximation may perform well in practice\" in section 3. Orthogonal the confusion I have above w/ the connection to the mean-field principle, the hypothesis stated here is imprecise and should be refined into something that is testable. The main issue with the statement is \"may perform well in practice\". While the authors may have a intuitive understanding of what they mean, this is by no means a testable/communicable hypothesis. Potentially this was written to try and soften the claims the paper makes, but hypotheses can be strong as long as they are testable and the authors test the validity of the hypothesis rather than form a hypothesis to explain their results.\n\nI would suggest the authors try and specify exactly what it means to perform well in practice, or better yet consider what hypothesis would be interesting to pose. Some potentially useful questions:\n  - Are you more focused on generalization (which seems to be the focus later in section 5) or generally to perform better than all other approaches?\n  - Do you expect the factorization to be beneficial broadly or in certain environments?\n  - Do you have some idea of what performing well means? Do the baselines \"perform well\"?\n  - How do we determine if the factorization is or isn't optimized to \"match the value function\"? Is this only due to the loose conditioning on the training from the out-going objective? Is there a stronger constraint I'm missing? Can we definitively say other networks don't implicitly learn this kind of factorization?\n\nCL-4). It is unclear how RIDE is incorporated into AFaR or how it is changed to make it a baseline. I'm specifically talking about the second paragraph of 5.4.2. This is incredibly unclear, and should be expanded on (at least in the appendix). \n\n\n## Consistency:\n\nCO-1). The main consistency issues I have are shared in some of the clarity questions above (specifically CL-2, and CL-3). Generally, I don't think positioning the main hypothesis around the mean-field principle and a loose definition of \"may perform well in practice\" is actually what the authors really mean (and actually suggest later in the paper). This needs to be made consistent.\n\n## Empirical Evaluation:\n\nMy main concern with this paper is the hanging questions which are avoided in favor of primarily testing downstream performance. The questions posed seem to be constructed as a consequence of the benchmarking done rather than novel questions posed and then tested. While this is generally not problematic, my issue is the avoidance of a main thrust underlying all of your claims. First, I want to go over the missing research question that needs answered for the other questions and the hypothesis to be tested and discussed. Second, I want to discuss the questions actually posed by the authors, and consider if these are appropriately interesting and if they should be posed differently. Finally, I will discuss some quibbles I have with how the experiments are laid out and some edits.\n\n### Missing Research Question\n\nThe main underlying claim the authors use as motivation is that their architecture factorizes the value function in accordance to a linear reward factored MDP. This claim, while maybe intuitive to the authors from their architectural choices, hasn't being evaluated or posed as an empirical question. I think testing this empirically is critical for the claims in the paper to be adequately supported. Below I've listed several recommendations/possible directions which might be interesting to fully analyze the network.\n\nEE-1). An evaluation of the individual value functions (i.e. $V(Z^i_t)$) and how they compare to the total value function.\n     - Are all the value functions contributing to the final value function equally? (i.e. are they just learning $\\frac{1}{k} V(s_t)$, or something else).\n     - If we compare instead to a feed-forward network which learns the value function directly (rather than like AFaR) and the output of a hidden layer is used in the same capacity as $z$.\n\nEE-2). How do the different states $(z^i)$ represent the underlying state? Do they look entirely different?\n\nEE-3). If AFaR is used on an MDP which has a clear reward factoraization will the network learn that factorization? I would really like to see some experiments where we have control over the factorizations, but I'm unaware of some specific examples you could use. Maybe something hand constructed?\n\n### Posed Questions\n\nNext I'll turn to the empirical questions posed in the paper. I'll go in order as they are posed.\n\nEE-4). (i). \"Can we use factored representations learned by the AFaR algorithm to train an RL policy?\"\n\nThis question suffers from some of the issues mentioned in CL-3. Specifically, it being interpreted as \"is it possible to use factored representations...\". The answer to this is yes, as your architecture uses it. Another way I could interpret this question is \"Will the policy learned using only a factored representation be useful?\" I think this is the interpretation we are wanting, but the current question doesn't say this clearly.\n\nEE-5). (ii). \"Does AFaR improve generalization to new environments where factors vary in novel ways?\"\n\nThere are two parts of this question that are unclear. What constitutes a new environment and what do you mean by \"where factors vary in novel ways?\". I think the authors need to clarify what they mean in some discussion \n\n\nEE-6). (iii). \"How robust is AFaR to hyper-parameters like the number of factors that the agent learnings?\"\n\nIn the empirical section, you only seem to check the number of factors and not other hyper-parameters. So you should focus this question specifically to that hyper-parameter.\n\n### Quibbles\n\nEE-7). A huge quibble I have with how the results are presented are through (what I honestly believe) is a misrepresentation of the actual results shown. If your performance gain is not statistically significant, your null hypothesis (i.e. the means are the same) is true. This means the results (which you've bolded but not added an astrix in your tables) are not improvements over the baselines. There are two changes required, you need to adjust the claims made in the text of the paper and you need to follow what is standard across the field (at least to my knowledge) which is to bold only statistically significant best performers. \n\nEE-8). I'm hesitant to trust the shaded region in Figure 3. If you added all the individual learning curves in the appendix, this would help add evidence to your claim that the lines are significant.\n\nEE-10). Did you only run the results in table 3 with one seed? If so we can't really make any claims wrt \n\nEE-11). It is unclear how the different baselines (and your method) compare on number of tunable parameters. This should be very clearly described and discussed, and if this is controlled for in your empirical evaluation you should mention this more clearly.\n\nEE-12). You included the state-action value function factorization in prop 2.2. I think an empirical evaluation over how this factorization would work in a DQN type network would be really interesting. I think this would more directly support your hypothesis that the specific factorization is helpful for a control algorithm.\n\nEE-13). What are your network architectures specifically? It is unclear given the current state of the paper. For the empirical evaluation in MiniGrid, do you use recurrence to solve the problem? If so how is this put into your networks?\n\n\n### Supported Conclusions:\n\nI've talked about this quite a bit above. So I'll just refer to the above review for this criteria.\n\n### Misc Questions:\n\nQ-1). In Figure 3, there are some weird results that require explanation (assuming EE-8 is appropriately addressed):\n     - Why does the factors set to 4 perform so much worse than those set to 2?\n     - What is being learned by each factored value function?\n     - It isn't enough to state that \"we note that in some cases the choice of k does not make much difference...However, in general, the performance can be improved by tweaking k\". I would expect a clear explanation of the behavior given your posed question (discussed in EE-6).\n\nQ-2). Would a MoE network be a good baseline to your approach? Why not use a more basic MoE network as compared to AFaR?\n\n### Misc Comments:\n\nMC-1). I would recommend re-writting the last paragraph of section 7. Currently, it is confusing. I think it is supposed to be focused on future work, but this isn't clear as currently stated.\n\nMC-2). The second paragraph in the introduction is a bit dense, and I'm not sure the cognitive science perspective is adding to the motivation much. It might be worthwhile being further explored in the related work section.\n\nMC-3). If you are referencing Figures 6 and 7 in the main text you should include them in the main paper (you still had a page of space!).\n\n### Minor Edits:\nThese are edits that are just minor mistakes, and don't factor into my thoughts on the paper.\n\n- The equations in both section 2.2 and appendix A often have hanging parenthesis at the end.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The primary focus of this paper was a novel architecture for representations learning based on reward factored MDPs. As I see it there are two main contributions:\n1. The novel architecture (AFaR) and describing this architecture in terms of reward factored MDPs.\n2. A performance comparison of AFaR with other schemes to improve generalization. Specifically, DrAC and IDAAC in the procgen environment, and RIDE in the MiniGrid-KeyCorridorSxRy environments.\n\nThe proposed architecture uses factored states learned implicitly to represent a value function. These states are then passed to the policy network (without gradient connections) and used to learn a policy in tandem with learning the states. The paper sets out to answer three questions related to the generalization performance of the proposed architecture, and some ancillary properties of the architecture (i.e. does it work, sensitivity to the number of factors).", "claims_and_evidence_yn": "", "claims_explanation": "", "audience_yn": "", "audience_explanation": "", "requested_changes": "I've presented a lot of comments, but there are some more serious than others. Below I'll list the most critical suggestions/comments that need to be addressed.\n\n*Need to be addressed:*\n- CL-(1, 2, 3, 4)\n- EE-(1, 2, 3, 4, 5, 6, 7, 10, 13)\n\nI would like to see all the points addressed or discussed in a rebuttal, but understand there is only so much time.\n", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "vjiC8V1lKW", "reviewer": "Reviewer_vDzx", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths\n* The paper is in general clearly written.\n\nMajor weaknesses:\n* The baseline results of DrAC and IDAAC does not match (are noticeably lower than) those reported in their original papers, which may render the comparison unfair.\n* Even if only considering the scores in this paper, I do not think the improvement is significant. Besides, I suggest that the authors report the normalized score over all 16 environments and use this new evaluation metric [1].\n* Experiments on MiniGrid are limited. Please consider running experiments on other environments, such as MultiRoom and ObstructMaze.\n* Some ablation experiments are missing. First, all ablation experiments should be conducted on both MiniGrid and Procgen, as the focuses of these two benchmarks are different. Second, the authors should add experiments with the attention mechanism in Eqn.(2). Third, since the highlight of this paper is reward factorization, it would be better to have ablation experiments on it.\n\nMinor weaknesses:\n* Some notations in Section 2.1 are inconsistent or incorrect.\n  * In Assumption 2.1 the reward function depends on $s_t$, $a_t$ and $s_{t+1}$, but in Proposition 2.2 it only depends on $s_t$, $a_t$.\n  * The expression $r:\\mathcal{S}\\times\\mathcal{A}\\to\\mathbb{R}$ in Proposition 2 is incorrect. For different factor $i$, $r_i$ is different since the sub state spaces are different. It should be $r_i:\\mathcal{S}^i\\times\\mathcal{A}\\to\\mathbb{R}$.\n  * Regarding the factor-wise value functions (two equations below Eqn.(1)), $\\pi$ appears on the left-hand-side but does not appear on the right-hand-side.\n\n\n[1] Agarwal, Rishabh, et al. \"Deep reinforcement learning at the edge of the statistical precipice.\" Advances in neural information processing systems 34 (2021): 29304-29320.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper introduces a new representation learning method that leverages reward factorization, and shows that this method (combined with baseline actor-critic methods) can improve generalization performance on Procgen and sample efficiency on MiniGrid.", "claims_and_evidence_yn": "", "claims_explanation": "", "audience_yn": "", "audience_explanation": "", "requested_changes": "See weaknesses above.", "broader_impact_concerns": "None", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper is easy to read and proposes a novel and useful idea. Automatically discovering factorizations is a useful direction and though this paper does not yet give much insight into what factorizations emerge and when this approach might be useful, I believe it will help others continue down this direction. 

The decision is Accept with Minor Revisions because I agree with the reviewers that there remains an open question about exactly what is providing performance differences with AFaR. The suggestion to only use attention, without factorization, should be feasible, as this is just a standard MOE approach without explicit critic updates. If that is not a possible ablation, then this could be explained in the work. 

A few other changes are also important.
1. The authors stated: "We are working to add a plot in the paper where show the output of the different value functions over time (averaged across timesteps and batch)". This would help understand the learned factorization.
2. The table only seems to bold items if AFaR is better. But typically bolding is used to show which method was statistically significantly better. You should also bold if other method were best. Note if you are willing to add more runs (at least in a few settings), then you could get more clear results.
3. It would be useful to explain for the experiment what factorizations might emerge. 
4. I actually could not understand exactly what was meant about learning the factorizations only with the critic, and attention for the policy part. A more explicit algorithm could clarify this, or an objective function could be provided in the section introducing the framework.

Not necessary changes below but just two comments that might be useful for you:
First, the sample efficiency improvements seem to be more important here than final performance results. I wonder if more of the results should focus on this aspect, as again the improvement from AFaR is not that large in terms of final performance. Adding this inductive bias should only improve efficiency, and so that is the natural thing to report. 

Second, such a reward factorization has been previously considered for Exogenous MDPs. It only factors into two parts of the state, so it by no means accomplishes what is done here. But, they also note the factorization of the value function. It could be useful to look at and cite this work: "Discovering and Removing Exogenous State Variables and Rewards for Reinforcement Learning", Dietterich et al.
Decision: Accept (poster)