Factored Policy Gradients: Leveraging Structure for Efficient Learning in MOMDPs
Year: 2021

Abstract:
Policy gradient methods can solve complex tasks but often fail when the dimensionality of the action-space or objective multiplicity grow very large. This occurs, in part, because the variance on score-based gradient estimators scales quadratically. In this paper, we address this problem through a factor baseline which exploits independence structure encoded in a novel action-target influence network. Factored policy gradients (FPGs), which follow, provide a common framework for analysing key state-of-the-art algorithms, are shown to generalise traditional policy gradients, and yield a principled way of incorporating prior knowledge of a problem domain’s generative processes. We provide an analysis of the proposed estimator and identify the conditions under which variance is reduced. The algorithmic aspects of FPGs are discussed, including optimal policy factorisation, as characterised by minimum biclique coverings, and the implications for the bias-variance trade-off of incorrectly specifying the network. Finally, we demonstrate the performance advantages of our algorithm on large-scale bandit and traffic intersection problems, providing a novel contribution to the latter in the form of a spatial approximation.

heading: 1 Introduction
text: Many sequential decision-making problems in the real-world have objectives that can be naturally decomposed into a set of conditionally independent targets. Control of water reservoirs, energy consumption optimisation, market making, cloud computing allocation, sewage flow systems, and robotics are but a few examples [36]. While many optimisation methods have been proposed [25, 34] — perhaps most prominently using Lagrangian scalarisation [46] — multi-agent learning has emerged as a promising new paradigm for sample-efficient learning [6]. In this class of algorithms, the multi-objective learning problem is cast into a centralised, co-operative stochastic game in which co-ordination is achieved through global coupling terms in each agent’s objective/reward functions. For example, a grocer who must manage their stock could be decomposed into a collection of subagents that each manage a single type of produce, but are subject to a global constraint on inventory. This approach has been shown to be very effective in a number of domains [20, 50, 30, 24, 52], but presents both conceptual and technical issues.
The transformation of a multi-objective Markov decision process (MOMDP) [36] into a stochastic game is a non-trivial design challenge. In many cases there is no clear delineation between agents in the new system, nor an established way of performing the decomposition. What’s more, it’s unclear in many domains that a multi-agent perspective is appropriate, even as a technical trick. For example, the concurrent problems studied by Silver et al. [39] exhibit great levels of homogeneity,
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
lending themselves to the use of a shared policy which conditions on contextual information. The key challenge that we address in this paper is precisely how to scale these single-agent methods — specifically, policy gradients — in a principled way. As we shall see, this study reveals that existing methods in both single- and multi-agent multi-objective optimisation can be formulated as special cases of a wider family of algorithms we entitle factored policy gradients. The contributions of this paper are summarised below:
1. We introduce influence networks as a framework for modelling probabilistic relationships between actions and objectives in an MOMDP, and show how they can be combined with policy factorisation via graph partitioning.
2. We propose a new control variate — the factor baseline — that exploits independence structures within a (factored) influence network, and show how this gives rise to a novel class of algorithms to which we ascribe the name factored policy gradients.
3. We show that FPGs generalise traditional policy gradient estimators and provide a common framework for analysing state-of-the-art algorithms in the literature including actiondependent baselines and counterfactual policy gradients.
4. The variance properties of our family of algorithms are studied, and minimum factorisation is put forward as a principled way of applying FPGs, with theoretical results around the existence and uniqueness of the characterisation.
5. The final contribution is to illustrate the effectiveness of our approach over traditional estimators on two high-dimensional benchmark domains.

heading: 1.1 Related Work
text: Policy gradients. Variance reduction techniques in the context of policy gradient methods have been studied for some time. The seminal work of Konda and Tsitsiklis [19] was one of the earliest works that identified the use of a critic as beneficial for learning. Since then, baselines (or, control variates) have received much attention. In 2001, Weaver and Tao [53] presented the first formal analysis of their properties, and later Greensmith et al. [13] proved several key results around optimality. More recently, these techniques have been extended to include action-dependent baselines [47, 22, 12, 57, 9], though the source of their apparent success has been questioned by some [49] who suggest that subtle implementation details were the true driver. It has also been shown that one can reduce variance by better accounting for the structure of the action-space, such as bounds [4, 10] or more general topological properties [7]. The SVRPG approach of Papini et al. [29] also addresses variance concerns in policy gradients by leveraging advances in supervised learning, and the generalised advantage estimator of Schulman et al. [37] has been proposed as a method for reducing variance in actor-critic methods with fantastic empirical results; both of these can be combined with baselines and the techniques we present in this work. Factorisation. In a related, but distinct line of work, factorisation has been proposed to better leverage the transition structure of MDPs; see e.g. [2, 15, 44]. Indeed, the notion of causality has also been utilised in work by Jonsson and Barto [16]. Most recently, Oliehoek et al. [27] presented an elegant framework for harnessing the influence of other agents (from the perspective of self) in multi-agent systems. This approach is complementary to the work presented in this paper, and more recent extensions have significantly advanced the state-of-theart [41, 5, 28]; we build upon these principles. There is also a long line of research on “influence diagrams” that is pertinent to this work. While the majority of this effort has been focused on dynamic programming, the ideas are very closely related to ours and indeed we see this work as a natural extension of these concepts [45]. Miscellaneous. Causal/graphical modelling has seen past applications in reinforcement learning [11]. Indeed, our proposed influence network is related to, but distinct from, the action influence models introduced by Madumal et al. [23] for explainability. There, the intention was to construct policies that can justify actions with respect to the observation space. Here, the intention was to exploit independence structure in MOMDPs for scalability and efficiency.

heading: 2 Background
text: A regular discrete-time Markov decision process (MDP) is a tupleM .= (S,A,R, p, p0), comprising: a state space S , action space A, and set of rewardsR ⊆ R. The dynamics of the MDP are driven by an initial state distribution such that s0 ∼ p0(·) and a stationary transition kernel where (rt, st+1) ∼
p(·, · | st,at) satisfies the Markov property, p(rt, st+1 |ht) = p(rt, st+1 | st,at), for any history ht . = (s0,a0, r0, s1, . . . , st,at). Given an MDP, a (stochastic) policy, parameterised by θ ∈ Rn, is a mapping πθ : S × Rn → P(A) from states and weights to the set of probability measures on A. The conditional probability density of an action a is denoted by πθ(a | s) . = P(a ∈ da | s,θ) and we assume throughout that πθ is continuously differentiable with respect to θ. For a given policy, the return starting from time t is defined as the discounted sum of future rewards, Gt . = ∑T k=0 γ
krt+k+1, where γ ∈ [0, 1] is the discount rate and T is the terminal time [42]. Value functions express the expected value of returns generated from a given state or state-action pair under the MDP’s transition dynamics and policy π: that is, vπ(s) . = Eπ[Gt | st = s] and qπ(s,a) . = Eπ[Gt | st = s,at = a]. The objective in control is to find a policy that maximises vπ for all states with non-zero measure under p0, denoted by the Lesbesgue integral J(θ) . = Ep0 [vπθ (s0)] = ∫ S vπθ (s0) dp0(s0).

heading: 2.1 Policy Search
text: In this paper, we focus on policy gradient methods which optimise the parameters θ directly. This is achieved, in general, by performing gradient ascent on J(θ), for which Sutton et al. [43] derived
∇θJ(θ) = Eπθ,ρπθ [(qπθ (s,a)− b(s)) z] , (1) where z .= ∇θ lnπθ(a | s) is the policy’s score vector, ρπθ (s) . = ∫ S ∑∞ t=0 γ
tp(st = s | ds0, πθ) denotes the (improper) discounted-ergodic occupancy measure, and b(s) is a state-dependent baseline (or, control variate) [33]. Here, p(st = s | s0, πθ) is the probability of transitioning from s0 → s in t steps under πθ . Equation 1 is convenient for a number of reasons: 1. it is a score-based estimator [26]; and 2. it falls under the class of stochastic approximation algorithms [1]. This is important as it means qπ(s,a) may be replaced by any unbiased quantity, say ψ : S × A → R, such that Eπ,ρπ [ψ(s,a)] = qπ(s,a), while retaining convergence guarantees. It also implies that optimisation can be performed using stochastic gradient estimates, the standard variant of which is defined below.
Definition 2.1 (VPGs). The vanilla policy gradient estimator for target-baseline pair (ψ, b) is denoted
gV(s,a) . = [ψ(s,a)− b(s)] z, (2) where ∇θJ(θ) = Eπθ,ρπθ [ gV(s,a) ] .

heading: 2.2 Factored (Action-Space) MDPs
text: In this paper, we consider the class of MDPs in which the action-space factors into a product, A .= ⊗n i=1Ai = A1 × · · · × An, for some n. This is satisfied trivially when n = 1 and A1 = A,
but also holds in many common settings, such as A .= Rn, which factorises n times as ⊗n
i=1 R. This is equivalent to requiring that actions, a ∈ A, admit a “subscript” operation; without necessarily having A be a vector space. For example, one could have an action-space of the form A .= R× N such that, for any a ∈ A, a1 ∈ R and a2 ∈ N. To this end, we introduce the notion of partition maps which will feature throughout the paper.
Definition 2.2 (Partition Map). Define X .= ⊗n i=1 Xi and J ⊆ [n] with XJ . = ⊗
j∈J Xj such that a partition map (PM) for a pair (X , J) is a function σ : X → XJ with complement σ̄ : X → X[n]\J .
Partition maps are an extension of the canonical projections of the product topology, and are equivalent to the scope operator used by Tian et al. [48]. For example, if (a1, a2, a3) . = a ∈ A .= R3 denotes a three-dimensional real action-space, then one possible PM is given by σ(a) = (a1, a3) with complement σ̄(a) = (a2). Note that there should always exist a unique inverse operation that recovers the original space; in this case, it would be expressed as f((a1, a3) , (a2)) = (a1, a2, a3).

heading: 3 Influence Networks
text: Consider an MOMDP with scalarised objective given by
J(θ) . = Ep0 ψ(s,a) .= m∑ j=1 λjψj(s, σj(a))  , (3)
where λj ∈ R for all 1 ≤ j ≤ m and each ψj(s, σj(a)) denotes some target that depends on a single partition of the action components. Traditional MDPs can be seen as a special case in which m = 1, and ψ = ψ1 . = qπ . The vector ψ(s,a) comprises the concatenation of all m targets and each partitioning is dictated by the non-empty maps σj(a), the form of which is intrinsic to the MOMDP. For convenience, let us denote the collection of targets comprising ψ(s,a) by
Ψ . = {ψj : ψ(s,a) = 〈λ,ψ(s,a)〉} . (4)
The intuition behind FPGs is derived from the observation that each factor of the action-space only influences a subset of the m targets. Take, for example, Figure 1c which depicts an instance of an influence network between a 2-dimensional action vector and a 2-dimensional target. The edges suggest that a1 affects the value of both ψ1 and ψ2, whereas a2 only affects ψ2. This corresponds to an objective of the form λ1ψ1(s, a1) + λ2ψ2(s,a), where each goal’s domain derives from the edges of the graph. This is formalised in Definition 3.1 below. Definition 3.1 (Influence Network). A bipartite graph G(M,Ψ) .= (IA, IΨ, E) is said to be the influence network of an MDPM and target set Ψ if for IA . = [|A|] and IΨ . = [|Ψ|], the presence of an edge, e ∈ E, between nodes i ∈ IA and j ∈ IΨ defines a causal relationship between the ith factor of A and the jth target ψj(s, σj(a)).
An influence network can be seen as a structural equation model [31] in which each vertex in IA has a single, unique parent which is exogenous and drives the randomness in action sampling, and each vertex in IΨ has parents only in the set IA as defined by the set of edges E. The structural equations along each edge (i, j) ∈ E are given by the target functions themselves and the partition maps σj mirror the parents of each node j. Some examples of influence networks are illustrated in Figure 1; see also the appendix. We now define the key concept of influence matrices. Definition 3.2 (Influence Matrix). LetKG denote the biadjacency matrix of an influence network G, defined as the |IA| × |IΨ| boolean matrix with Kij = 1 ⇐⇒ (i, j) ∈ E for i ∈ IA and j ∈ IΨ.
Together, these definitions form a calculus for expressing the relationships between the factors of an action-space and the targets of an objective of the form in Equation 3. We remark that, from an algorithmic perspective, we are free to choose between two representations: graph-based, or partition map-based. The duality between G and K, and the set {σj : j ∈ IΨ}, is intrinsic to our choice of notation and serves as a useful correspondence during analysis.

heading: 3.1 Policy Factorisation
text: Influence networks capture the relationships between A and Ψ, but policies are typically defined over groups of actions rather than the individual axes of A. Consider, for example, a multi-asset trading problem in which an agent must quote buy and sell prices for each of n distinct assets [14, 40]. There is a natural partitioning between each pair of prices and the n sources of profit/loss, and one might therefore define the policy as a product of n bivariate distributions as opposed to a full joint, or fully factored model. This choice over policy factorisation relates to the independence assumptions we make on the distribution πθ for the sake of performance. Indeed, in the majority of the literature, policies are defined using an isotropic distribution [57] since there is no domain knowledge to motivate more complex covariance structure. We formalise this below.
Definition 3.3 (Policy Factorisation). An n-fold policy factorisation, Σ .= {σπi : i ∈ [n]}, is a set of disjoint partition maps that form a complete partitioning over the action space.
The definition above provides a means of expressing any joint policy distribution in terms of PMs,
πθ(a | s) . = n∏ i=1 πi,θ(σ π i (a) | s) , (5)
where σπi ∈ Σ and n = |Σ|. This corresponds to a transformation of the underlying influence network where the action vertices are grouped under the n policy factors and, for any i, j ∈ [n], i 6= j, we have mutual independence: σπi (a) |= σπj (a). This is captured in the following concept. Definition 3.4 (Factored Influence Network). For a given influence network G and policy factorisation Σ, we define a factored influence network, GΣ, by replacing IA with IΣ, the set of partitioned vertices, and merge the corresponding edges to give EΣ. Similarly, denote by KΣ the influence matrix with respect to the Σ-factorisation.
Factored influence networks ascribe links between the policy factors in Equation 5 and the targets ψj ∈ Ψ. They play an important role in Section 4 and provide a refinement of Definition 3.1 which allows us to design more efficient algorithms. As an example, Figure 2 shows how one possible policy factorisation transforms an influence network G into GΣ. Note that while the action nodes and edges have been partitioned into policy factors, the fundamental topology with respect to the attribution of influence remains unchanged; i.e. no dependencies are lost.

heading: 4 Factored Policy Gradients
text: Factored policy gradients exploit factored influence networks by attributing each ψj ∈ Ψ only to the policy factors that were probabilistically responsible for generating it; that is, those with a connecting edge in the given GΣ. The intuition is that the extraneous targets in the objective do not contribute to learning, but do contribute towards variance. For example, it would be counter-intuitive to include ψ2 of Figure 2b in the update for π1 since it played no generative role. Naturally, by removing these terms from the gradient estimator, we can improve the signal to noise ratio and yield more stable algorithms. This idea can be formulated into a set of baselines which are defined and validated below. Definition 4.1 (Factor Baselines). For a given GΣ, the factor baselines (FBs) are defined as
bFi (s, σ̄ π i (a)) . = [(1−KΣ)λ ◦ψ(s,a)]i , (6)
for all i ∈ [|Σ|], where ◦ denotes the Hadamard product and 1 is to be taken as an all-ones matrix. Lemma 4.1. FBs are valid control variates if GΣ is true to the MDP (i.e. unbiased).
Factor baselines are related to the action-dependent baselines studied by Wu et al. [57] and Tucker et al. [49], as well as the methods employed by COMA [9] and DRPGs [3] in multi-agent systems. Note, however, that FBs are distinct in two key ways: 1. they adhere to the structure of the influence network and account not only for policy factorisation, but also the target multiplicity of MOMDPs; and 2. unlike past work, factor baselines were defined using an ansatz based on the structure implied by a given GΣ as opposed to explicitly deriving the arg min of the variance, or approximation thereof; see the appendix. This means that, unlike optimal baselines, FBs can be computed efficiently and
thus yield practical algorithms. Indeed, this very fact is why the state-value function is used so ubiquitously in traditional actor-critic methods as a state-dependent control variate despite being sub-optimal. It follows that we can define an analogous family of methods for MOMDPs with zero computational overhead. Proposition 1 (FPGs). Take a Σ-factored policy πθ(a|s) and |θ| × |Σ| matrix of scores S(s,a). Then, for target vector ψ(s,a) and multipliers λ, the FPG estimator
gF(s,a) . = S(s,a)KΣ λ ◦ψ(s,a) , (7) is an unbiased estimator of the true policy gradient; i.e.∇θJ(θ) = Eπθ,ρπθ [ gF(s,a) ] .
Proposition 1 above shows that the VPG estimator given in Definition 2.1 can be expressed in our calculus as S1λ ◦ψ, where 1 is an all-ones matrix and, traditionally, ψ .= qπ; note that one can still include other baselines in Equation 7 such as vπ. In other words, Proposition 1 strictly generalises the policy gradient theorem [43] and, by virtue of it’s unbiasedness, thus retains all convergence guarantees. We also see that both COMA [9] and DRPGs [3] are special cases in which the influence network reflects the separation of agents withKΣ a square, and often diagonal matrix.

heading: 4.1 Variance Analysis
text: The variance reducing effect of FBs comprises two terms: 1. a quadratic and thus non-negative component which scales with the second moments of bFi ; and 2. a linear term which scales with the expected values of bFi . This is shown in the following result. Proposition 2 (Variance Decomposition). Let gi denote a gradient estimate for the ith factor of a Σ-factored policy πθ (Equation 5). Then, ∆Vi . = V [ gVi ] − V [ gFi ] , satisfies
∆Vi = αi Eσ̄πi (a) [( bFi )2] + 2βiEσ̄πi (a) [ bFi ] , (8)
where zi . = ∇θ lnπi,θ(a | s), αi . = Eσπi (a)[〈zi, zi〉] ≥ 0 and βi . = Eσπi (a) [ 〈zi, zi〉 ( ψ + bFi )] .
The first of these two terms is a “free lunch” which removes the targets that are not probabilistically related to each factor. The linear term, on the other hand, couples the adjusted target with the entries that were removed by the baseline. This suggests that asymmetry and covariance can have a regularising effect in VPGs that is not present in FPGs — a manifestation of the properties of control variates [26]. Now, if we do not assume that the target functions are bounded, then the linear term in Equation 8 can grow arbitrarily in either direction, but we typically require that rewards are restricted to some compact subset R ⊂ R to avoid this. Below, we show that if a similar requirement holds for each target function — namely, that infS,A ψj is well defined for each ψj ∈ Ψ — then we can always construct a set of mappings that constrain (8) to be non-negative without biasing the gradient. Corollary 4.1 (Non-Negative Variance Reduction). Let ψ(s,a) be of the form in Equation 3. If ψj(s,a) ≥ ψj for all (s,a) ∈ S × A and j ∈ [m], with |ψj | < ∞, then there exists a linear translation, ψi → ψi − ∑m j=1 λjψj , which leaves the gradient unbiased but yields ∆Vi ≥ 0.1
Interestingly, numerical experiments on a pair of continuum armed bandits suggest that this transformation is seldom necessary; see Figure 3. As the number of policy factors and targets grow, so too does the potential discrepancy in magnitude between the quadratic and linear terms in Equation 8. The former starts to dominate even for small |Σ|. This is particularly prevalent when the influence matrix KΣ is very sparse and the baselines have wide coverage over Ψ. In other words, applying FBs when the influence network is very dense or even complete will not yield tangible benefits (e.g. in Atari games), but applying them to a problem with a rich structure, such as traffic networks, will almost certainly yield a significant net reduction in variance.
Bias-Variance Trade-Off. It is important to note that, in real-world problems, one does not always know the exact structure of the influence network underling an MDP ex ante. This poses a challenge since incorrectly removing edges can introduce bias and thus constrain the space of solutions that can be found by FPGs. Note, however, that this may not always be a problem, since a small amount
1This inequality can be made strict if either αi > 0 or βi > 0 — where the former equates to having a non-zero trace of the Fisher information matrix — and a small ε > 0 is added to the translation.
(a) Search bandit: Cost(a) .= − ||a− c||1 with fixed centroid vector c ∈ Rn.
(b) ReLU bandit: Cost(a) .= − ∑n
i=1 max (eiai, 0) with fixed sign vector e ∈ {−1, 1}n.
Figure 3: Variance decomposition on a symmetric log scale for two bandit problems as a function of action-space dimensionality. Each term was computed using Monte-Carlo estimation with 105 samples and taking the arithmetic mean across all policy factors.
bias for a large reduction in variance can be desirable. Furthermore, one could leverage curriculum learning to train the policy on (presumed) influence networks with increasing connectedness over time. This trade-off between bias and variance is present in many machine learning settings, and depends strongly on the problem at hand; we explore this empirically in Section 5.2.

heading: 4.2 Minimum Factorisation
text: For many classes of fully-observable MDPs, any policy factorisation is theoretically viable: we can fully factor the policy such that each action dimension is independent of all others; or, at the other extreme, treat the policy as a full joint distribution over A. This holds because, in many classes of (fully-observable) MDPs, there exists at least one deterministic optimal policy [55, 35]. The covariance acts as a driver of exploration, and it’s initial value only affects the rate of convergence.2 As a result, most research uses an isotropic Gaussian with diagonal covariance to avoid the cost of matrix inversion. This poses an interesting question: is there an “optimal” policy factorisation, Σ?G , associated with an influence network G? Below we offer a possible characterisation.
Definition 4.2 (Minimum Factorisation). A minimum factorisation (MF), Σ?G , of an influence network, G, is the minimum biclique vertex cover, disjoint amongst IA.
It follows from Definition 4.2 that for any Σ?G , each σi ∈ Σ?G is a biclique (i.e. complete bipartite subgraph) of the original influence network G, and that the bipartite dimension is equal to the number of policy factors. For example, one can trivially verify that Figure 2b is an MF of the original graph; see also the reductions in Figure 1. In essence, an MF describes a complete partitioning over action vertices — so as to define a proper distribution — where each group is a biclique with the same set of outgoing edges. The “minimum” qualifier then ensures that the maximum number of nodes are included in each of these groups, a property which allows us to prove the following result:
Theorem 4.1. The MF Σ?G always exists and is unique.
Minimum factorisation is a natural construction for the problem domains studied in this paper; see Section 5. It also yields factored policies which, generally, expose the minimum infimum bound on variance for a given influence network. This follows from the fact that an MF yields the greatest freedom to express covariance structure within each of the policy factors whilst also maximising the quadratic term in Equation 8. In fact, when each action corresponds to a single unique target, the MF enjoys a lower bound on variance that is linear in the number of factors. Finally, we remark that, whilst closely related to vertex covering problems (which are known to be NP-complete [18]), we observed experimentally that finding the MF can be done trivially in polynomial time; see e.g. [8].
2Note that this is not true in general: the policy’s covariance structure impacts the set of reachable solutions in partially-observable MDPs and stochastic games, for example.
(a) Convergence of CPGs (with and w/o additional state-dependent baseline) compared with VPGs using different baselines. The error bands denote the standard error on the mean over 10 random seeds. (b) Squared error between an and cn during learning for the search bandit with k = n− 1.
Figure 4: Performance analysis of CPGs on the search bandit domain.

heading: 5 Numerical Experiments

heading: 5.1 Search Bandits
text: Consider an n .= 1000 dimensional continuum armed bandit with action space in Rn, and cost function: Cost(a) .= ||a− c||1 + λ ζ(a), where c ∈ Rn, λ ≥ 0 and ζ : A → R+ is a penalty function. This describes a search problem in which the agent must locate the centroid c subject to an action-regularisation penalty. It abstracts away the prediction aspects of MDP settings, and allows us to focus only on scalability; note that this problem is closely related to the bandit studied by Wu et al. [57] for the same purpose. In our experiments, the centroids were initialised with a uniform distribution, c ∼ U (−5, 5) and were held fixed between episodes. The policy was defined as an isotropic Gaussian with fixed covariance, diag(1), and initial location vector µ .= 0. The influence network was specified such that each policy factor, πi,θ for i ∈ [n], used a reward target ψi(a) . = −∆i(ai)−λζ(ai), with ∆i(a) . = |a− ci|, amounting to a collection of n forks (Figure 1a). The parameter vector, µ, was updated at each time step, and the hyperparameters are provided in the appendix.
We began by examining the case where λ = 0 and the co-ordinate axes were fully decoupled. For VPGs, we note that stability was only possible without a baseline if an extremely low learning rate was used; see the appendix. Including a baseline dramatically improved performance, with the action-dependent case, b(s, a), also leading to better asymptotic behaviour at the expense of a two orders of magnitude longer train-time according to the wall-clock compared with all other algorithms (VPGs and FPGs); see Table 1. In comparison, FPGs, both with and without a learnt state-dependent baseline, yielded significantly reduced variance, leading to faster learning, more consistent convergence and highly robust asymptotic stability.
We then studied the impact of coupling terms in the cost function; i.e. λ > 0. For this, we considered a family of penalties taking the form of partially applied `2 norms: ζk(a) . = √∑k i=1 a 2 i , with 1 ≤ k ≤ n. This set of functions allowed us to vary the penalty attribution across the n factors of A. Further examples demonstrating the performance advantage of FPGs — for k = n and k = n/2
(a) Convergence of joint/factored policies for the global reward, and using FPGs with the spatial baseline. Each curve depicts the mean value across 5 random seeds with standard error bands. (b) Performance degradation as a function of the nlevel spatial baseline approximation in a 2 × 6 grid network. Each point is the average terminal value across 5 seeds and with standard error bands.
Figure 5: Performance analysis of FPGs (with PPO and GAE) on the traffic network domain.
— are given in the appendix. In both cases, the improvement due to FB adjustments was found to be non-negative for every combination of learning rate and action space. This confirms that FPGs can indeed handle coupled targets and retains the variance reduction benefits that were explored in Section 4.1. As an illustrative example, consider the case where k = n− 1 and all but the last action dimension are subject to a penalty. This is a particularly challenging setting for VPGs because the magnitude of the combined cost function is much greater than ∆n(an), leading to an aliasing of the final component of the action vector in the gradient. The result, as exemplified in Figure 4b, was that VPGs favoured reduction of overall error, and was therefore exposed to poor per-dimension performance; hence the increased noise in the an error process. FPGs avoid this effect by attributing gradients directly.

heading: 5.2 Traffic Networks
text: We now consider a classic traffic optimisation domain to demonstrate the scalability of our approach to large real-world problems. In particular, we consider variants of the (3×3) grid network benchmark environment — as originally proposed by Vinitsky et al. [51] — that is provided by the outstanding Flow framework [56, 21]. In this setting, the RL agent is challenged with managing a set of traffic lights with the objective of minimising the delay to vehicles travelling in the network; the configuration should be taken as identical unless explicitly stated. This requires a significant level of co-ordination, and indeed multi-agent approaches have shown exemplary performance in this space [51, 54]. However, much as with the search bandit, the probability of aliasing effects increases substantially with the number of lights/intersections; i.e. the dimensionality of the action-space. This affects both single- and multi-agent approaches when the global reward is used to optimise the policy.
To this end, we propose a “baseline” that removes reward terms derived from streets/edges that are not directly connected to a given traffic light. This is based on the hypothesis that the local problem is sufficiently isolated from the rest of the system that we may still find a (near)-optimal solution; much as with local-form models [27]. Of course, this could introduce bias at the cost of variance if we are incorrect (see Section 4.1), but this turns out to be an effective trade-off as exemplified in Figure 5a.3 In this plot we compare the performance of three policies learnt using PPO [38] and GAE [37] (with an additional state-dependent baseline): (1) a naïve joint policy over the 9-dimensional action-space trained against the global reward; (2) a shared policy trained on the global reward; and (3) a shared policy using the local spatial baseline. In methods 2 and 3, a shared policy refers to the use of a single univariate policy across all nine traffic lights, where only local information and identity variables are provided. The global reward in this case was defined as the negative of the mean delay introduced in the system minus a scaled penalty on vehicles at near standstill; see the appendix for more details.
As expected, we observe that the FBs improve learning efficiency, but, perhaps surprisingly, we also find that the asymptotic behaviour is also superior. We posit that this relates to the fact that, with a fixed learning rate, stochastic gradient descent cannot distinguish between points within a ball of the optimum solution with radius that scales with the variance on the estimator. In other words, significant reductions in variance, even if they introduce a small amount of bias, may increase the likelihood of reaching the true optimal solution by virtue of having much greater precision.
3Note that the standard errors may slightly underestimate the population level due to the low sample size.
To better understand this trade-off, we also explored the impact of “expanding” the local baseline in a larger system of 2 × 6 intersections. With this new baseline we retain reward terms derived from lights up to n edges away in either the east or west directions. The variable n thus provides a dial to directly tweak the bias and variance of the policy gradient estimator (i.e. increasing n reduces bias but increases variance). The result, as shown in Figure 5b, suggest that performance decreases monotonically as a function of n. This corroborates the claim in Section 4.1 that introducing some bias in exchange for a reduction in variance can be a worthwhile trade-off in large problems.

heading: 6 Conclusion
text: Factored policy gradients derive from the observation that many MOMDPs exhibit redundancy in their reward structure. Here, we have characterised this phenomenon using graphical modelling, and demonstrated that conditional independence between factors of the action-space and the optimisation targets can be exploited. The resulting family of algorithms subsume many existing approaches in the literature. Our results in large-scale bandit and concurrent traffic management problems suggest that FPGs are highly suited to real-world problems, and may provide a way of scaling RL to domains that have hitherto remained intractable. What’s more, FPGs are compatible with other techniques that improve policy gradient performance. For example, they can be extended to use natural gradients by pre-multiplying Equation 7 by the inverse Fisher information matrix [17], and can even use additional baselines to reduce variance even further, as in Section 5.2. In future work we intend to address the following interesting questions: (a) Can we infer/adapt the structure of influence networks? (b) Are there canonical structures within S andK? (c) What theoretical insights can be derived from a more detailed analysis of the variance properties of FPGs? We argue that factored approaches such as FPGs — which are complementary to ideas like influence-based abstraction [27] — are a promising direction for practical RL. Addressing some of these questions, we believe, would thus be of great value to the community.
Acknowledgments and Disclosure of Funding
The authors would like to acknowledge our colleagues Joshua Lockhart, Jason Long and Rui Silva for their input and suggestions at various key stages of the research. This work was conducted by JPMorgan’s AI Research group which has no external funding sources; i.e. it was a self-funded project. No sources of financial competing interests (or otherwise) are attributed to this line of research.
Disclaimer This paper was prepared for informational purposes by the Artificial Intelligence Research group of JPMorgan Chase & Co and its affiliates (“J.P. Morgan”), and is not a product of the Research Department of J.P. Morgan. J.P. Morgan makes no representation and warranty whatsoever and disclaims all liability, for the completeness, accuracy or reliability of the information contained herein. This document is not intended as investment research or investment advice, or a recommendation, offer or solicitation for the purchase or sale of any security, financial instrument, financial product or service, or to be used in any way for evaluating the merits of participating in any transaction, and shall not constitute a solicitation under any jurisdiction or to any person, if such solicitation under such jurisdiction or to such person would be unlawful.
© 2021 JPMorgan Chase & Co. All rights reserved.

Review summaries:
This paper studies the problem of policy learning in a factored Markov decision process (MDP). The factorization of the transition probability distribution is graphically described through a set of influence diagrams. The objective function is a linear combination of weighted kernel functions. The authors propose an efficient policy gradient algorithm to learn the optimal policy. Simulation results support the authors’ results.

This paper proposes a policy gradient algorithm for MOMDPs with large factored action spaces, where the background knowledge on the MOMDP structure is given as an influence network. A new baseline for policy gradients is given that makes use of the MOMDP background knowledge and a factorisation of the RL policy. The paper analyses the favourable variance properties of the resulting policy gradient estimator. Empirical results are provided on a pedagogical bandit example, and a traffic control problem. The latter demonstrates applicability of the method even when the graphical assumptions only hold approximately.