Factored Value Functions for Graph-Based Multi-Agent Reinforcement Learning

TL;DR

This work tackles credit assignment in large-scale graph-structured cooperative MARL by introducing the Diffusion Value Function (DVF), a vector-valued, graph-diffused surrogate that decays with both time and graph distance via the diffusion operator $\Gamma = \gamma \mathds{A} D^{-1}$. DVF is proven to have a Bellman fixed point, to decompose the global value through an averaging property, and to align with policy improvements when its components improve. The DVF serves as a drop-in critic within a Diffusion A2C (DA2C) framework, learning via diffusion TD errors, with extensions to decentralised scenarios using graph neural networks. To address communication constraints, the Learned DropEdge GNN (LD-GNN) learns sparse message passing and local outputs, enabling scalable, distributed algorithms. Empirically, DA2C with LD-GNN outperforms local and global critics across firefighting and distributed computation tasks, with up to 11% improvements in average reward and strong generalisation to unseen graph sizes and topologies.

Abstract

Credit assignment is a core challenge in multi-agent reinforcement learning (MARL), especially in large-scale systems with structured, local interactions. Graph-based Markov decision processes (GMDPs) capture such settings via an influence graph, but standard critics are poorly aligned with this structure: global value functions provide weak per-agent learning signals, while existing local constructions can be difficult to estimate and ill-behaved in infinite-horizon settings. We introduce the Diffusion Value Function (DVF), a factored value function for GMDPs that assigns to each agent a value component by diffusing rewards over the influence graph with temporal discounting and spatial attenuation. We show that DVF is well-defined, admits a Bellman fixed point, and decomposes the global discounted value via an averaging property. DVF can be used as a drop-in critic in standard RL algorithms and estimated scalably with graph neural networks. Building on DVF, we propose Diffusion A2C (DA2C) and a sparse message-passing actor, Learned DropEdge GNN (LD-GNN), for learning decentralised algorithms under communication costs. Across the firefighting benchmark and three distributed computation tasks (vector graph colouring and two transmit power optimisation problems), DA2C consistently outperforms local and global critic baselines, improving average reward by up to 11%.

Figures (9)

• Figure 1: The Learned DropEdge GNN (LD-GNN) architecture. Given node and edge memory states $(X,E)$ and observations $O^t$, the policy $\lambda_\theta$ samples neighbour sets $C^t$, inducing a subgraph on which a GNN aggregates messages. GRUs update $(X,E)$, and $\psi_\theta$ samples outputs $Y^t$.
• Figure 2: Test performance as a function of the message-passing penalty $p_m$ for the three tasks (graph colouring, service quality and energy efficiency). Error bars indicate the lower and upper quartiles across runs. We omit IA2C from the service quality task due to consistently low rewards, and from the energy efficiency task because it coincides with the no-pass baseline.
• Figure 3: Training curves for different message-passing penalties $p_m$ on the three tasks (graph colouring, service quality and energy efficiency). Error bars indicate the lower and upper quartiles across runs. Curves are smoothed with a moving average for visual clarity.
• Figure 4: TX-RX communication graph.
• Figure 5: Transformation from a communication graph to the influence graph of its edge GMDP.

Theorems & Definitions (12)

• Proposition 3.1: Existence
• Proposition 3.2: Uniqueness
• Proposition 3.3: Factored value
• Proposition 3.4: Policy alignment
• Lemma 1.1: Limited influence
• proof
• Lemma 2.1
• proof
• proof : Proof of Proposition \ref{['prop:bounded']}
• proof : Proof of Proposition \ref{['prop:uniqueness']}