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
Eeqrand
KL rand
Dactrand
Ddiv
Fsame
PPO-ID
−34.67±4.45
0.293±0.355
0.117±0.147
0.498±0.052
0.649±0.047
0.187±0.108
PPO-ID+Eq , λeq=0.01
−41.75±6.24
0.009±0.006
0.002±0.001
0.196±0.103
0.425±0.053
0.739±0.139
Eq, λeq=0.05
−34.13±17.13
0.009±0.003
0.003±0.000
0.112±0.102
0.378±0.043
0.878±0.117
Eq+Pair, λeq=0.05,λpair=0.007
−48.16±20.70
0.007±0.004
0.002±0.001
0.150±0.135
0.396±0.054
0.826±0.151
Table 1: Main PPO results with N=3 agents, mean ± standard deviation across three seeds. Return is episode return; Eeqrand , KL, and Dact are computed under random permutations. Ddiv and Fsame are computed under normal ordering.
Figure 2: Four-agent stress test. The collapse floor is Ddiv=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
Eeqrand
Dactrand
Ddiv
Fsame
Fmax
PPO-ID
−58.43
0.063
0.331
0.384
0.582
0.719
Eq, λeq=2.5×10−4
−64.99
0.022
0.369
0.427
0.293
0.591
Eq, λeq=5×10−4
−43.12
0.025
0.121
0.297
0.811
0.704
Eq, λeq=10−3
−66.65
0.008
0.310
0.367
0.575
0.751
Eq, λeq=10−2
−28.04
0.016
0.000
0.250
1.000
1.000
Eq, λeq=5×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 ( λ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.
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.
Kei Suzuki, Jing Liu, Ye Wang +4
Mitsubishi Electric Research Laboratories (MERL), Cambridge, Massachusetts, USA.
Many reinforcement learning systems, from fleet management to traffic signal control, must serve an objective that changes dynamically after deployment, and retraining a policy for each new objective is prohibitively expensive. For a single agent, this problem is well understood: successor features with generalized policy improvement, together with their universal extension, recombine a library of learned policies into a policy for any new objective, with a guarantee that the result is never worse than any policy in the library. However, multi-agent transfer has received far less attention, and the common practice of letting each agent recombine its own library independently inherits the recipe but not the guarantee. We prove that this independent composition can produce joint behavior strictly worse than every policy in the library, because recombining teammates changes the environment each agent faces and invalidates the values it relies on, a failure with no single-agent counterpart. We further show that the only unconditionally safe fixed rule is synchronized composition, which moves the whole team to one jointly trained policy but cannot serve objectives that assign different goals to different agents. To attain safety and flexibility at once, we propose MA-USFA, a hierarchical method with two layers: a lower layer of universal successor feature approximators that predicts each agent's successor features while conditioned on its teammates' objectives, and an upper composer that selects, across agents, which library entry each agent should follow and supplies the cross-agent correction a per-agent value cannot represent. Trained once over the distribution of objectives, it is applied at deployment with no per-task adaptation.
Zijian Zhao, Sen Li
1The Hong Kong University of Science and Technology · 2The Hong Kong University of Science and Technology (Guangzhou)
We address the problem of real-time multi-agent path planning under prescribed initial and final configurations, formulated as a high-dimensional optimal control problem subject to inter-agent collision avoidance and obstacle constraints. Classical solvers based on the Hamilton--Jacobi--Bellman equation or Pontryagin's maximum principle (PMP) become computationally intractable as the number of agents grows, due to the curse of dimensionality and the non-convexity introduced by interaction constraints. We propose OCFormer, a transformer-based neural network that parameterizes the inverse of a coordinate transformation mapping a simple, analytically tractable latent space, in which agents evolve independently under a decoupled Hamiltonian, to the physical space of the coupled, interacting system. The model is trained by minimizing a physics-informed loss comprising the Hamiltonian ODE residual from the PMP and a boundary matching term, without requiring pre-computed trajectory data. The self-attention mechanism endows OCFormer with permutation equivariance over agents, a structural symmetry that prior symplectic network approaches do not natively enforce. When trained on a distribution of boundary conditions, the amortized OCFormer produces near-optimal trajectories for unseen initial and final configurations in a single forward pass (∼220ms), enabling real-time inference without retraining. We further introduce a lightweight transfer learning framework that adapts a pretrained OCFormer to out-of-distribution perturbations in seconds via L-BFGS fine-tuning. Numerical experiments on two- and three-dimensional multi-agent navigation with up to 128 agents demonstrate the scalability and efficiency of the proposed approach.
Qian Zhang, Shanqing Liu, Jerome Darbon +1
Division of Applied Mathematics, Brown University, Providence, RI 02912, USA.