Bayesian Ego-graph inference for Networked Multi-Agent Reinforcement Learning

Grammatical and formatting issues (e.g., Fig. 2 legend, punctuation in Eq. (12))

We appreciate this comment and will thoroughly revise the manuscript in the final version, including:

• Moving all legends (e.g., Fig. 2) to consistent and appropriate positions.

• Carefully proofread all equations, especially Eq. (12), to ensure correct punctuation and presentation.

• Conducting a full grammatical and typographic pass.

Clarity and completeness of Section 4.2 (Bayesian Graph Inference)

Thank you for pointing this out. We agree that Section 4.2 is central to the paper and could benefit from expanded exposition. In the final version:

• We will revise this section to clarify the motivation for using Bayesian inference, specifically, to enable uncertainty-aware edge masking and regularized graph structure learning.

• We will better explain how the variational distribution over binary masks is amortized using graph neural networks and trained end-to-end with the policy.

• We will also clarify the interaction between the ELBO objective and policy optimization.

We also note that a more detailed explanation of this mechanism and architecture is currently provided in the Appendix, but we agree that bringing some of this detail into the main text will improve accessibility and understanding.

Theoretical guarantees and handling of non-stationarity in DTDE settings

We appreciate the reviewer’s insightful question regarding theoretical guarantees and non-stationarity in decentralized MARL.

1. Non-stationarity in DTDE

In DTDE settings, each agent learns its policy independently while treating the environment as stationary. However, in multi-agent scenarios, the environment includes other agents whose policies are continuously evolving during training. This results in non-stationary learning signals, which often degrade stability and convergence in policy iteration. This issue is well-documented in prior works (e.g., Hernandez-Leal et al., 2019 [survey on MARL]).

2. How CTDE alleviates non-stationarity

Centralized Training with Decentralized Execution (CTDE) mitigates this issue by providing each agent access to global information or centralized critics during training. This stabilizes learning by accounting for co-adaptation among agents. Prior works such as QMIX (ICML 2018) and MADDPG (ICML 2017) have shown that centralized critics can eliminate or significantly reduce non-stationarity during training, leading to more stable and sample-efficient learning.

3. Our relationship to CTDE via Spatiotemporal MDP

Our work lies between DTDE and CTDE in terms of structural assumptions. Specifically, we operate under the Spatiotemporal MDP framework (Definition 1), in which:

• Each agent's transition and reward are conditioned on its physical neighborhood.

• This imposes a partial coupling across agents via a fixed environment graph (e.g., road topology), ensuring that agents only adapt to localized dynamics rather than the entire system.

• Consequently, the variability perceived by each agent is constrained to its stable local neighborhood, substantially reducing global non-stationarity.

Through this structural prior and ego-centric masking, our approach preserves decentralized execution while benefiting from the stability typically associated with CTDE. While we do not offer formal convergence guarantees, our design choices are empirically validated in terms of training stability and sample efficiency (Section 5.4, Fig. 5).

On the lack of formal analysis: We also note that many recent NeurIPS publications proposing novel methods for MARL similarly do not provide theoretical guarantees, focusing instead on empirical performance and structural insights. These include:

• Learning to Share in Multi-Agent Reinforcement Learning (NeurIPS 2022)

• Selectively Sharing Experiences Improves Multi-Agent Reinforcement Learning (NeurIPS 2023)

• Automatic Grouping for Efficient Cooperative Multi-Agent Reinforcement Learning (NeurIPS 2023)

We will expand our discussion in the final version to better articulate the structural mitigation of non-stationarity via spatiotemporal MDPs and position our approach relative to CTDE and DTDE.

Whether the learned graph sampling process impacts the stability of policy iteration

This is an insightful concern. Our response has two parts:

• Learning a variational distribution over edge masks allows us to capture epistemic uncertainty over which neighbors are most informative. This stochasticity can improve robustness in dynamic or partially observable environments, which reflects real-world scenarios (e.g., noisy or fluctuating traffic patterns).

• The sampled communication graph $Z_i$ influences the actor loss via the graph-conditioned policy $\tilde{\pi}_i$, as shown in Definition 4. This actor loss is embedded in the likelihood term of our ELBO formulation: L_ELBO = E_{q(Z_i; φ_i)}[ -L_{θ, φ} + log p(Z_i) - log q(Z_i; φ_i) ].This setup ensures that mask sampling is directly driven by task performance—graphs that yield better policies are more likely under the learned posterior.

Therefore, our sampling-based approach is not only grounded probabilistically but also aligned with policy improvement. Empirically, as shown in our ablation studies (e.g., Fig. 5), this coupling leads to more stable and performant learning than alternatives like deterministic graph construction or random masking.

Thank you for your response! I still have a few concerns regarding this work:

1. For convergence analysis of related methods, I would like to draw the authors’ attention to [1] and [2].
2. Compared to [1] and [2], it might be helpful to include a discussion comparing this work with similar approximation-based approaches. I notice that in the context of networked MARL, several existing methods already employ sampling strategies—[1] and [2] being good examples. What is the key distinction between the proposed method and these prior works? Given time constraints, I understand that additional experiments may not be feasible. However, I encourage the authors to focus on clearly articulating how their approach differs from and advances beyond these closely related methods. Addressing this point would significantly help me better evaluate the paper’s contribution, as I currently have some unresolved concerns.

Reference

[1] Lin Y, Qu G, Huang L, et al. Multi-agent reinforcement learning in stochastic networked systems. In NeurIPS, 2021.

[2] Anand E, Qu G. Efficient reinforcement learning for global decision making in the presence of local agents at scale. arXiv preprint arXiv:2403.00222, 2024.

Conciseness of Section 3 and suggestions to restructure it.

Thank you for the suggestion. Due to space limitations in the submission version, we provided a condensed version of the background. However, we acknowledge the importance of making the preliminaries more self-contained and structured. In the final version, we will:

Restructure Section 3 into two subsections:

• One focusing on the spatiotemporal MDP setting, clearly explaining how state transitions, rewards, and local observability are defined under neighborhood-based interactions.
• The second summarizes key prior approaches and their limitations in the context of dynamic, decentralized interaction modeling.

We will also move some essential formal definitions from the appendix to the main text in the final version to enhance readability and self-containment.

---

Details of adaptive traffic signal control (ATSC) in Section 5.1 and justification for Definition 1

We appreciate this question and the opportunity to clarify. In Section 5.1, we briefly described the mapping from the real-world traffic environment to the MDP components. To elaborate:

• States: Each agent observes its local traffic state via sensors such as induction loop detectors (ILDs), which provide vehicle density, queue length, and waiting times on incoming lanes.

• Actions: The action space of each agent consists of phase control decisions (e.g., switching between green and red for traffic directions).

• Transition Probabilities: These are not analytically available but are implicitly defined via the microscopic traffic simulator (SUMO). The transition of each agent’s local traffic state depends primarily on its own action and the actions of neighboring intersections (e.g., upstream traffic release).

• Rewards: Defined as the (negative) number of halted vehicles at the intersection, normalized by a fixed scale. This encourages reducing traffic congestion locally.

These components are naturally aligned with the Spatiotemporal MDP framework (Definition 1) because:

• The transition dynamics are localized and depend on neighborhood interactions, consistent with the assumption that p_i(s_i'| s_{V_i}, u_i, u_{N_i}).

• The reward function is decomposed per agent, depending only on local state-action pairs.

• The neighborhood structure is explicitly defined by the physical road network graph, which is fixed and sparse.

Moreover, centralized training methods are often impractical or suboptimal in ATSC due to several real-world constraints:

• Scalability: As the size of the traffic network grows (e.g., in city-scale scenarios), centralized models become computationally expensive and suffer from poor generalization due to the curse of dimensionality in joint state-action spaces.
• Communication constraints: Centralized training assumes full observability or global message aggregation, which is unrealistic in large distributed traffic systems with bandwidth and latency limitations.

Thus, our formulation — learning decentralized policies that respect the fixed physical connectivity and local dynamics — not only satisfies the assumptions of the spatiotemporal MDP, but also reflects the operational needs of real-world traffic control systems.

We will revise Section 5.1 to clarify these mappings and add this justification. We also plan to release our code and ATSC simulation environments upon publication.

---

Strengthening the ablation study of graph masking strategies

We agree that understanding when graph masking yields the most benefit is valuable. While our current ablation study includes two maps (Monaco and NewYork51) with different characteristics, we did not include all five due to time constraints during the rebuttal phase. We will extend this study across all ATSC maps in the final version.

Regarding the effectiveness of masking:

• According to our analysis, the benefit of graph masking is more pronounced in denser physical networks, where agents are connected to more neighbors. In such settings, the probability of receiving redundant or irrelevant information increases, which can negatively impact policy learning. Our learned masks help filter out such noisy edges, promoting more focused and effective communication.

• In sparser environments, where each agent has fewer neighbors by design, the amount of incoming noise is naturally limited, and thus the marginal benefit of masking is expected to be smaller.

On the apparent difference between Monaco and NewYork51 in Figure 5:

The seemingly smaller improvement on NewYork51 is partially due to scale compression, caused by the relatively poor performance of the RandMask baseline. In absolute terms: The performance gap between BayesG and NoMask is ~100 in Monaco and ~60–70 in NewYork51, both of which are substantial.

We will clarify this interpretation in the revised version and present both raw and normalized performance gains in the final plots.

On Baseline Selection and Comparison with Graph-Structured Methods like CASEC (Q1 & Q2)

We thank the reviewer for the thoughtful questions regarding our choice of baselines and the omission of recent graph-based methods like CASEC from our experimental comparisons.

Our baseline selection is grounded in the theoretical setting of Spatiotemporal Markov Decision Processes (Spatiotemporal-MDP), as formalized in Definition 1 of our paper. In Spatiotemporal-MDP, each agent’s transition dynamics and reward function are explicitly localized over its physically connected neighbors, meaning that the environment inherently imposes a structured graph topology (e.g., road networks in traffic signal control).

Specifically, the local transition function p_i(s_i' | s_{V_i}, u_i, u_{N_i}) depends only on the state-action tuples of agent $i$ and its neighbors $\mathcal{N}_i$, and the reward R(s_{V_i}, u_{V_i}) is likewise defined over the local subgraph. The resulting learning setting is fully decentralized, grounded in real-world physical constraints, and lacks any centralized critic, global reward, or global coordination mechanism.

In contrast, many recent graph-based MARL methods such as CASEC, DCG [1], DICG [2], SOP-CG [3], and GACG [4] are built on the Decentralized Partially Observable MDP (Dec-POMDP) framework. These methods operate under significantly different assumptions, including:

• Access to or estimation of a global state $s \in S$,

• A shared global reward signal $R(s, \mathbf{a})$ based on the joint action,

• Joint action-based transitions $P(s^{\prime} \mid s, \mathbf{a})$,

• And often, centralized training schemes with global observability or coordination.

These assumptions make Dec-POMDP-based algorithms fundamentally incompatible with real-world physical systems governed by local dynamics, such as traffic networks, power grids, or sensor fields, where each agent only observes and affects its local neighborhood.

Therefore, to ensure a fair and meaningful comparison aligned with our problem setting, we selected baselines that respect the assumptions and structure of the Spatiotemporal-MDP framework. Specifically, methods like CommNet, NeuroComm, and LToS are well-suited because:

• They communicate only with physically connected neighbors.

• They respect the fixed topology imposed by the environment graph.

• They perform local message passing and decentralized policy learning, without relying on a global reward or state.

• These baselines thus share the same structural and semantic assumptions as our proposed method.

Finally, while CASEC and other Dec-POMDP-based methods represent important contributions to multi-agent coordination under different settings, we intentionally excluded them from our main comparisons due to their incompatibility with the physical and decentralized nature of ST-MDP. We will add a clarifying discussion in the final version and cite these works in the related work section to clearly distinguish between the modeling assumptions and ensure transparency in our baseline choices.

References:

[1] Boehmer et al. Deep Coordination Graphs. ICML 2020.

[2] Liu et al. Deep Implicit Coordination Graphs. AAMAS 2021.

[3] Yan et al. Self-Organized Polynomial-Time Coordination Graphs. ICML 2022.

[4] Duan et al. Group-Aware Coordination Graph. IJCAI 2024.

---

The probabilistic interpretation of using the usual policy gradient objective (Q3)

We thank the reviewer for this insightful question.

In our framework, we adopt a probabilistic perspective on reinforcement learning. The trajectory data $D_i$ generated by each agent is treated as being sampled from an underlying distribution governed jointly by the agent's policy and the latent communication structure $Z_i$.

Why Use the Policy Objective as the Likelihood in the ELBO

In variational inference, the log-likelihood term log p(D_i | Z_i, G_env_{V_i}) measures how well the latent structure $Z_i$ explains the data. We instantiate this term using the graph-conditioned actor loss $\mathcal{L}_{\theta,\varphi}$ , defined in Definition 4.

Specifically: $log \tilde{\pi}_i(a_i \mid \tilde{f}_i) \cdot \hat{A}_i$ is treated as the data log-likelihood. This quantity provides task-aligned feedback on how suitable the sampled graph $Z_i$ is for policy optimization: if a sampled structure leads to high advantage (good performance), it is assigned higher likelihood in the ELBO objective. This ties the learning of latent graph structures directly to downstream task performance.

This formulation follows the same probabilistic treatment as entropy-regularized RL, particularly the Soft Actor-Critic (SAC) framework [Haarnoja et al., 2018], which casts RL as probabilistic inference. In SAC, the optimal policy is derived by maximizing the expected log-likelihood of actions under a soft Q-function, plus an entropy bonus to ensure sufficient stochasticity and exploration. This regularization is critical for:

• Improving exploration in high-dimensional action spaces.

• Avoiding premature convergence to suboptimal deterministic policies.

• Encouraging robustness to model misspecification and variance in gradients.

Our actor loss mirrors this design, embedding stochasticity into policy optimization while preserving alignment with the Bayesian structure learning objective.

Why Include Mask Entropy $-\log q(Z_i; \phi_i)$

Analogous to policy entropy in SAC, we include a mask entropy term in our ELBO objective, defined over the variational distribution $q(Z_i; \phi_i)$ of edge masks. This term serves several key purposes:

• Uncertainty Modeling: The entropy of the variational mask distribution captures uncertainty in the learned communication structure. In real-world settings with noisy and partially observable interactions, such uncertainty is inevitable and informative.

• Regularization: Including mask entropy discourages premature collapse of the variational distribution to a deterministic mask. This is particularly important in early training when structural estimates are still evolving.

• Diversity and Exploration: High entropy promotes diversity in sampled graphs, allowing agents to explore multiple plausible communication topologies. Over time, the posterior sharpens around those that consistently lead to higher policy performance.

• Analytical Tractability: Since $q(Z_i)$ is a factorized Bernoulli distribution, its entropy can be computed in closed form using the sigmoid outputs $\sigma(\phi_{ij})$. This allows for efficient optimization.

This design is aligned with the maximum entropy principle in SAC and variational inference theory, where entropy promotes robustness and avoids overfitting to noise or spurious correlations in structure learning.

By treating the graph-conditioned policy objective as a proxy for the log-likelihood and incorporating entropy terms for both policy and structure, our ELBO formulation provides a principled probabilistic mechanism for joint optimization of communication and action strategies. We will clarify this probabilistic interpretation and its relation to Soft Actor-Critic in the final version.

Why not use direct communication? Why use variational latent variables for message masking?

We appreciate this insightful question. There are two layers to unpack here:

(1) Why not direct communication over the full environment graph?

While the environment-provided physical graph (e.g., road topology in traffic networks) provides a grounded prior for interactions, communicating over the entire neighborhood indiscriminately can lead to redundant or noisy information aggregation. Our ablation in Section 5.4 (Figure 5) compares no masking (i.e., direct message passing on the environment graph) with BayesG’s learned mask, and shows that direct communication performs consistently worse. Therefore, selectively filtering communication is essential to avoid distraction and enhance relevance.

(2) Why use a variational latent variable instead of learning a deterministic mask?

Rather than directly learning a fixed mask, we model the mask as a stochastic latent variable sampled from a variational distribution $q(Z_i; \phi)$. This design brings several key advantages:

• Captures uncertainty in graph structure, which is particularly important under partial observability or sparse feedback in MARL.

• Encourages exploration in structural space, enabling agents to discover better interaction topologies during training.

• Incorporates structured priors, such as a Bernoulli sparsity prior, to promote minimal and efficient graphs.

• Supports entropy-based regularization, which smooths optimization and improves training stability.

In summary, direct communication is ineffective due to a lack of relevance filtering; deterministic masking fails to model structural uncertainty. By contrast, our variational latent mask provides a principled, adaptive, and uncertainty-aware mechanism that leads to improved decentralized decision-making, as validated by ablation results.

Mechanism for improved policy performance and empirical evidence (Section 5.4)

The learning signal for $Z_t$ is provided via the graph-conditioned actor loss $\mathcal{L}_{\theta,\phi}$ (see Definition 4), which appears in the ELBO's likelihood term. This loss connects interaction topology learning to downstream task performance: if a sampled graph leads to better policy outcomes, it receives a higher likelihood. This tight coupling ensures that the learned graphs are task-driven and performance-aligned.

In our ablation study on different graph masking strategies (Section 5.4), we compare:

• No masking (i.e., using the full environment graph),

• Random masking, and

• Our learned Bayesian mask (BayesG).

The learned mask consistently outperforms the baselines, showing that task-aware variational latent variables yield more effective communication graphs, ultimately leading to better decentralized policy performance.