Permutation Robustness Is Not Enough: Action Collapse in Multi-Agent Transformer Policies

Amit Thakur; Mukesh Singhal

cs.ROOct 2, 2026

Permutation Robustness Is Not Enough: Action Collapse in Multi-Agent Transformer Policies

Authors: Amit Thakur, Mukesh Singhal

Organizations: University of California, Merced

Abstract

Transformer policies are attractive for multi-agent robot learning because self-attention can model interactions among agents. However, multi-agent teams are unordered, while transformers typically process agents as ordered token sequences. We study how this mismatch affects cooperative navigation policies under agent-order permutations. Our results show that low permutation error alone can be misleading: policies may appear robust simply because all agents choose the same action. We therefore evaluate policies using both permutation-consistency metrics and action-collapse diagnostics, including action diversity, same-action fraction, and maximum action frequency. A PPO-ID baseline yields non-collapsed behavior but remains order-sensitive, while strong equivariance regularization can still induce homogeneous behavior. A weak equivariance penalty improves the robustness while preserving more diverse actions for teams with $N=3$ agents, whereas teams with $N=4$ agents require substantially smaller regularization weights. These findings suggest that multi-agent transformer policies should be evaluated not only by return and permutation robustness, but also by whether they maintain non-collapsed, differentiated multi-agent behavior.

Figures & tables

Figure 1: Three-agent PPO results. The ID/position baseline is diverse but order-sensitive. Weak equivariance regularization improves permutation consistency but reduces diversity; stronger regularization moves toward homogeneous behavior.

Method	Return	$E_{\rm eq}^{\rm rand}$	KL rand	$D_{\rm act}^{\rm rand}$	$D_{\rm div}$	$F_{\rm same}$
PPO-ID	$-34.67\pm 4.45$	$0.293\pm 0.355$	$0.117\pm 0.147$	$0.498\pm 0.052$	$0.649\pm 0.047$	$0.187\pm 0.108$
PPO-ID+Eq , $\lambda_{\rm eq}=0.01$	$-41.75\pm 6.24$	$\mathbf{0.009}\pm 0.006$	$\mathbf{0.002}\pm 0.001$	$0.196\pm 0.103$	$0.425\pm 0.053$	$0.739\pm 0.139$
Eq, $\lambda_{\rm eq}=0.05$	$-34.13\pm 17.13$	$\mathbf{0.009}\pm 0.003$	$0.003\pm 0.000$	$\mathbf{0.112}\pm 0.102$	$0.378\pm 0.043$	$0.878\pm 0.117$
Eq+Pair, $\lambda_{\rm eq}=0.05,\lambda_{\rm pair}=0.007$	$-48.16\pm 20.70$	$\mathbf{0.007}\pm 0.004$	$\mathbf{0.002}\pm 0.001$	$0.150\pm 0.135$	$0.396\pm 0.054$	$0.826\pm 0.151$

Table 1: Main PPO results with

N=3

agents, mean

\pm

standard deviation across three seeds. Return is episode return;

E_{\rm eq}^{\rm rand}

, KL, and

D_{\rm act}

are computed under random permutations.

D_{\rm div}

and

F_{\rm same}

are computed under normal ordering.

Figure 2: Four-agent stress test. The collapse floor is

D_{\rm div}=1/4=0.25

. Regularization strengths that are useful for

N=3

can fully collapse

N=4

policies, while much weaker values reveal a return-diversity tradeoff.

Method	Return	$E_{\rm eq}^{\rm rand}$	$D_{\rm act}^{\rm rand}$	$D_{\rm div}$	$F_{\rm same}$	$F_{\max}$
PPO-ID	$-58.43$	$0.063$	$0.331$	$0.384$	$0.582$	$0.719$
Eq, $\lambda_{\rm eq}=2.5\!\times\!10^{-4}$	$-64.99$	$0.022$	$0.369$	$\mathbf{0.427}$	$\mathbf{0.293}$	$\mathbf{0.591}$
Eq, $\lambda_{\rm eq}=5\!\times\!10^{-4}$	$\mathbf{-43.12}$	$0.025$	$0.121$	$0.297$	$0.811$	$0.704$
Eq, $\lambda_{\rm eq}=10^{-3}$	$-66.65$	$\mathbf{0.008}$	$0.310$	$0.367$	$0.575$	$0.751$
Eq, $\lambda_{\rm eq}=10^{-2}$	$-28.04$	$0.016$	$0.000$	$0.250$	$1.000$	$1.000$
Eq, $\lambda_{\rm eq}=5\!\times\!10^{-2}$	$-82.00$	$0.004$	$0.000$	$0.250$	$1.000$	$1.000$

Table 2: Four-agent stress test, seed 0. The coefficient that is useful for

N=3

(

\lambda_{\rm eq}=0.01

) fully collapses for

N=4

. Smaller coefficients expose a tradeoff between return, robustness, and non-collapse.

Figure 3: Robustness-diversity tradeoff. The desirable region is toward the upper-left: low permutation error and high action diversity. Points near the lower-left are robust but collapse-prone. Labels are placed with offsets for readability.

Appendix figures & tables1 asset

Supplementary material from the paper’s appendix.

Appendix

Metric	Good direction	Failure sign
$D_{\rm div}$	high	near $1/N$
$F_{\rm same}$	low	near $1$
$F_{\max}$	low/moderate	near $1$
$E_{\rm eq}$	low	high indicates order sensitivity
$D_{\rm act}$	low	high means selected actions change

Appendix

Table 3: Interpretation of collapse metrics.

Multi-Agent Systems

Multi-Agent Systems studies how autonomous agents coordinate, compete, communicate, and learn in shared or interdependent tasks.

See all 668 papers

Embedding Morphology into Transformers for Cross-Robot Policy Learning
Feb 26, 2026Kei Suzuki, Jing Liu, Ye Wang +4Robot Policy Learning Cross-Embodiment Robot Learning
Mitsubishi Electric Research Laboratories (MERL), Cambridge, Massachusetts, USA.
Is Per-Agent Policy Composition Safe? Rethinking Successor-Feature Transfer in Cooperative Multi-Agent Reinforcement Learning
Aug 12, 2026Zijian Zhao, Sen LiSequential Decision Making Safe RL
1The Hong Kong University of Science and Technology · 2The Hong Kong University of Science and Technology (Guangzhou)
A Transformer-based Approach for Real-time Multi-agent Path Planning with Varying Initial/Final Conditions
Oct 21, 2026Qian Zhang, Shanqing Liu, Jerome Darbon +1Decentralized Multi-Agent Motion Planning Multi-Agent Planning
Division of Applied Mathematics, Brown University, Providence, RI 02912, USA.

Feb 26, 2026J/K move · Enter open · S save

Embedding Morphology into Transformers for Cross-Robot Policy Learning

Kei Suzuki, Jing Liu, Ye Wang, Chiori Hori+3

Mitsubishi Electric Research Laboratories (MERL), Cambridge, Massachusetts, USA.

Robot Policy Learning Cross-Embodiment Robot Learning Robot Kinematics

Open paper AskPDF arXiv

Cross-robot policy learning -- training a single policy to perform well across multiple embodiments -- remains a central challenge in robot learning. Transformer-based policies, such as vision-language-action (VLA) models, are typically embodiment-agnostic and must infer kinematic structure purely from observations, which can reduce robustness across embodiments and even limit performance within a single embodiment. We propose an embodiment-aware transformer policy that injects morphology via three mechanisms: (1) kinematic tokens that factorize actions across joints and compress time through per-joint temporal chunking; (2) a topology-aware attention bias that encodes kinematic topology as an inductive bias in self-attention, encouraging message passing along kinematic edges; and (3) joint-attribute conditioning that augments topology with per-joint descriptors to capture semantics beyond connectivity. Across a range of embodiments, this structured integration consistently improves performance over a vanilla pi0.5 VLA baseline, indicating improved robustness both within an embodiment and across embodiments.

Abstract

Multi-Agent Systems

Explore similar work

Embedding Morphology into Transformers for Cross-Robot Policy Learning

Is Per-Agent Policy Composition Safe? Rethinking Successor-Feature Transfer in Cooperative Multi-Agent Reinforcement Learning

A Transformer-based Approach for Real-time Multi-agent Path Planning with Varying Initial/Final Conditions