Heterogeneous Multi-Agent Self-Attention (HMSA)

• HMSA is a neural attention module that facilitates efficient, scalable, type-adaptive policy learning and communication in heterogeneous multi-agent systems.
• It employs parameterized self-attention layers, featuring Block-gain and HetGAT variants, to manage dynamic topologies and diverse agent capabilities.
• Empirical studies demonstrate improved convergence, parameter efficiency, and effective sim-to-real transfer in distributed control and multi-robot tasks.

The Heterogeneous Multi-Agent Self-Attention module (HMSA) is a class of neural attention mechanisms explicitly designed to facilitate efficient, scalable, and type-adaptive policy learning and communication in teams of interacting agents with heterogeneous capabilities and observation-action spaces. HMSA aligns with emerging trends in distributed multi-agent reinforcement learning (MARL), handling both time-varying communication topologies and agent class diversity through parameterized self-attention architectures. Recent implementations of HMSA have shown strong empirical results, outperforming both model-free and model-based alternatives in complex multi-robot and game-theoretic domains (Sebastián et al., 22 Sep 2025, Seraj et al., 2021).

1. Architectural Principles and Variants

At its core, HMSA comprises several stacked self-attention layers, parameterized either per agent (fully heterogeneous) or per agent class (type-heterogeneous), operating on temporally-varying communication graphs. Each agent maintains an input embedding that aggregates neighbor information according to the current topology, utilizing both linear and nonlinear transformations with type-adaptive projections.

There are two main HMSA instantiations:

• Block-gain HMSA for distributed control: Used in distributed policy gradient frameworks, with attention block gains parameterized for each neighbor to realize time-varying feedback (Sebastián et al., 22 Sep 2025).
• HetGAT-based HMSA for communication learning: Implements edge- and node-type specific projections to learn type-conditioned communication protocols (Seraj et al., 2021).

Both variants maintain efficient scalability by localizing computation to the neighborhood and separating parameter sharing along team/agent-type axes.

2. Mathematical Formulation

Policy Gradient Block-gain HMSA

Let agent $\ell$ in team $i$ at time $k$ collect state observations from its neighbors:

$$x_\ell^{\mathcal{N}}(k) = [x_{n_{\ell_1}}(k), ..., x_{n_{\ell_m}}(k)] \in \mathbb{R}^{x \times |\mathcal{N}_\ell(k)|}$$

A $W$-layer self-attention stack computes at each layer $w$: \begin{align*} Q_i^w &= A_i^w X_i^{(w-1)} \ K_i^w &= B_i^w X_i^{(w-1)} \ V_i^w &= C_i^w X_i^{(w-1)} \ S &= \beta(Q_i^w) \beta(K_i^w)^\top + M \ \alpha &= \text{softmax}(S) \ Y_i^w &= \chi(\alpha \beta(V_i^w)) \ X_i^w &= \psi(D_i^w Y_i^w) \end{align*} Here, mask $M$ zeros out attention to non-neighbors, and nonlinearities $\beta$, $\chi$, $\psi$ are typically identity, ReLU, or tanh. The final block gain matrix is reshaped to form neighbor-specific local feedback parameters $k^{\ell, p}$.

Heterogeneity is enforced by maintaining distinct sets of parameters $$\theta_i^\ell = \{A_i^w, B_i^w, C_i^w, D_i^w\}_{w=1}^W$$ per agent or per team, with no weight sharing across teams/agents.

HetGAT HMSA for Diverse Communication

Let $C$ be agent classes. Each node (agent) $j$ carries $h_j \in \mathbb{R}^d$, with receiver-type $i$ and sender-type $t$ projections $W_i, W_{t \to i} \in \mathbb{R}^{d' \times d}$, and attention vector $a \in \mathbb{R}^{2d'}$.

For neighbor $k$ of type $t$: \begin{align*} e_{jk}^{(t \to i)} &= \text{LeakyReLU}\Big(a^\top [W_i h_j \Vert W_{t \to i} h_k]\Big) \ \alpha_{jk}^{(t \to i)} &= \frac{\exp(e_{jk}^{(t \to i)})}{\sum_{\ell \in \mathcal{N}j} \exp(e{j\ell}^{(**)})} \end{align*} Updates are computed via multi-head attention, combining per-head outputs by concatenation or averaging. The final embedding is passed through a softmax head for stochastic policy output.

Type-conditioned projections $W_{t \to i}$ and $W_i$ enable explicit heterogeneity in message representation and attention calculation (Seraj et al., 2021).

3. Integration into Multi-Agent Policy Optimization

HMSA modules naturally integrate into both decentralized (per-agent) and centralized training paradigms.

• In block-gain HMSA, the output defines neighbor-specific feedback policies:

$$\pi_i^\ell(x_\ell^{\mathcal{N}}(k); \theta_i^\ell) = \sum_{p \in \mathcal{N}_\ell(k)} k^{\ell, p}(x_\ell^{\mathcal{N}}(k); \theta_i^\ell) x_p(k)$$

Stochastic exploration is implemented by wrapping these outputs in Gaussian policies, and optimization is performed using a multi-agent variant of Proximal Policy Optimization (MAPPO).

• In HetGAT HMSA, the output embedding $h_j$ is mapped to per-class action logits, and softmax policies are sampled per agent. During centralized training, an additional State-Summary Node (SSN) may be appended for value estimation by a critic head (Seraj et al., 2021).

Backpropagation through the full HMSA stack is standard for both the actor and critic pathways.

4. Handling of Heterogeneity and Dynamic Topologies

A defining feature of HMSA is explicit support for agent heterogeneity and time-varying communication topologies:

• Parameterization: By maintaining independent projection matrices for each agent, team, or agent-class, HMSA enables learning of specialized attention and feedback patterns. No shared weights are required across agent types, supporting arbitrarily mixed-type teams.
• Masking and Edge Typing: Both block-gain and HetGAT HMSA variants use masking to enforce local communication (only topological neighbors participate in attention), and HetGAT introduces edge-type distinctions for sender-receiver class combinations. These features automatically accommodate variable team sizes and dynamic connectivities.
• Scalability: The per-agent per-layer complexity is $$O(|\mathcal{N}_\ell(k)|^2 d_{\text{model}} + |\mathcal{N}_\ell(k)| d_{\text{model}}^2)$$ for block-gain HMSA and $O(H(|E|d' + |V|d'))$ for HetGAT, both scaling linearly with the number of agents under local communication (Sebastián et al., 22 Sep 2025, Seraj et al., 2021).

5. Empirical Performance and Benchmarking

Evaluation of HMSA has been conducted on a diverse range of environments:

Domain	HMSA Implementation	Observed Outcomes
Distributed LQR	Block-gain / MAPPO	Convergence within 2% of centralized LQR; corrects for time-varying, unknown graphs
Nonlinear MARL	Both variants	Closes within 5% of model-based DP-iLQR cost in <100 iterations, with no dynamics model
Pursuit-evasion	Both variants	Lowest average evader distance/variance, stable catches/rewards, fewer params than GNN/MLP
Mixed-robot teams	HetGAT, MAHAC	∼10–15% improvement in convergence speed over CommNet/IC3Net in both homogeneous and heterogeneous settings

In real-robot experiments (Robotarium platform), HMSA enabled zero-shot sim-to-real policy transfer with adherence to physical safety constraints and the emergence of sophisticated tactics (Sebastián et al., 22 Sep 2025). Ablations show that HMSA is robust to over-specified graphs and prunes spurious edges in attention.

6. Practical Considerations and Hyperparameters

• For block-gain HMSA, most experiments use $d_{\text{model}}=64$ and $W=2$ or $3$ layers. Final projection $D^w$ is sized $(U \cdot X) \times 64$ per layer.
• In HetGAT settings, three HMSA layers ($L=3$), each with four attention heads ($H=4$), and per-head output dimensions of 16 or as required by action space are standard.
• Nonlinearities: ELU or ReLU for feature updates, LeakyReLU for attention score computation.
• Optimizer details include the Adam optimizer with learning rate $\approx 10^{-3}$ and no dropout (Seraj et al., 2021).
• For both approaches, hyperparameter count is independent of agent/team size due to parameter sharing schemes.

7. Relationship to Prior and Contemporary Work

HMSA advances earlier attention-based MARL models by introducing strict heterogeneity both in actor-critic policy parametrization and in communication protocol learning. Unlike traditional homogeneous GAT or communication networks (e.g., CommNet, IC3Net), HMSA supports per-type and per-agent specialization and avoids performance degradation observed in naïve application to mixed-type teams without explicit heterogeneity modeling (Seraj et al., 2021). Empirical results demonstrate clear improvements in convergence rate, final task performance, and parameter efficiency relative to both centralized deep MLP baselines and graph-attention with homogeneous weights (Sebastián et al., 22 Sep 2025).

A plausible implication is that HMSA architectures will become increasingly central in MARL for robotic swarms, autonomous vehicles, and distributed game-theoretic control, especially in settings where agent classes and connectivity are dynamically evolving.