Bayesian Ego-graph Inference for Networked Multi-Agent Reinforcement Learning¶
🕒 Published (v1): 2025-01-01 · Source: NeurIPS · Venue: NeurIPS 2025 · link
Ask a follow-up
Open an assistant pre-loaded with this paper's context.
💬 Ask ChatGPT✦ Ask Claude
TL;DR¶
BayesG introduces a decentralized actor–critic framework for networked MARL that learns sparse, task-adaptive communication graphs via Bayesian variational inference over each agent's local ego-graph. Each agent infers a binary edge mask as a latent variable, approximated by a learned Bernoulli variational posterior, and jointly optimizes interaction topology and policy via an ELBO objective. The method outperforms six baselines on traffic signal control tasks scaling up to 167 agents.
Problem¶
Networked MARL agents operating over fixed physical graphs are constrained to communicate with all immediate neighbors indiscriminately, regardless of contextual relevance. Existing dynamic-graph methods (DGN, variational coordination graphs) require centralized training with global state access, making them inapplicable in fully decentralized, physically constrained networked systems.
Method¶
Each agent \(i\) defines an ego-graph over its physical neighborhood \(\mathcal{V}_i\). A latent binary edge mask \(Z_i \in \{0,1\}^{|\mathcal{V}_i| \times |\mathcal{V}_i|}\) is treated as a posterior variable: \(p(Z_i \mid G^{\text{env}}_{\mathcal{V}_i}, D_i)\), where \(D_i\) is agent-local data (neighbor states, policy fingerprints, RNN trajectory embeddings). The intractable posterior is approximated by a factored Bernoulli variational distribution \(q(Z_i; \phi_i) = \prod_{j \in \mathcal{N}_i} \text{Bern}(z_{ij}; \sigma(\phi_{ij}))\), with Gumbel-Softmax reparameterization for gradient flow through discrete edge samples.
The full training objective is an ELBO that combines: 1. The graph-conditioned A2C actor loss (instantiates the likelihood term) 2. A KL term between \(q(Z_i)\) and a Bernoulli prior (mask sparsity regularization) 3. A mask entropy term preventing premature collapse to deterministic graphs 4. A standard A2C critic loss
The effective subgraph \(G^*_i = Z_i \odot G^{\text{env}}_{\mathcal{V}_i}\) is used for masked GCN message passing over observations, policy fingerprints, and trajectory features to form the agent's input representation.
Key Contributions¶
- Stochastic graph-based policy formulation for networked MARL, conditioning decisions on a sampled subgraph restricted to the physical topology
- Formulation of ego-graph edge selection as Bayesian variational inference with per-edge Bernoulli posteriors
- End-to-end ELBO training that jointly optimizes communication topology and actor–critic policy without global state or centralized infrastructure
- Demonstrated scalability to 167-agent real-world Manhattan traffic networks (NewYork167)
Results¶
- ATSC_Grid (25 agents) and Monaco (28 agents): BayesG achieves higher asymptotic reward and faster convergence than all baselines including NeurComm (the strongest prior communication-based method)
- NewYork33, NewYork51, NewYork167: BayesG shows higher final returns and more stable training curves; baselines plateau or become unstable at scale
- Ablation — masking strategy: Random masking severely degrades performance; learned variational mask significantly outperforms both no-masking and random-masking on Monaco and NewYork51
- Ablation — mask input features: Trajectory-based features > state-only; combining state + trajectory + policy features achieves best performance
- Qualitative: BayesG produces visibly sparser traffic congestion on Grid map at 3500 simulation seconds vs. IA2C, CommNet, NeurComm
(Note: exact numeric reward values are shown in training curves but not reported as discrete tables in the provided text.)
Limitations¶
- Evaluated exclusively on adaptive traffic signal control (ATSC); generalization to other networked MARL domains (power grids, wireless, robotics) is not empirically validated
- The variational distribution \(q(Z_i; \phi_i)\) is conditioned on fixed learnable logits \(\phi_{ij}\), not dynamically input-dependent at inference time (the mask uses neighbor states/trajectories during training but the logit parameters are static per edge pair)
- Communication overhead analysis (message complexity, bandwidth) is not provided, despite the claimed efficiency
- Does not establish theoretical convergence guarantees under the learned stochastic graph, unlike prior sampling-based networked MARL works (Lin et al., Qu et al.)
- Restricted to 1-hop ego-graphs; multi-hop or hierarchical coordination is not addressed
Relevance to Agentic AI / LLM Agents¶
BayesG directly addresses a core challenge in multi-agent systems: how agents should selectively attend to peers rather than processing all available communication indiscriminately. The Bayesian variational approach to learning communication topology is architecturally analogous to learned routing or selective attention in agentic pipelines, where which agents/tools/memory sources to consult is itself a learned decision. The ego-graph formulation—fully decentralized, operating under local observability, constrained to physical connectivity—is a strong model for real-world multi-agent deployments (robot fleets, sensor networks, distributed LLM agent topologies) where centralized orchestration is unavailable. The ELBO-based joint training of structure and behavior provides a principled alternative to heuristic agent-selection rules that dominate current LLM multi-agent frameworks.