LEGO: Local-Canonicalization Equivariant Graph Neural Networks for Sample-Efficient and Generalizable Swarm Robot Control

Figure 1. An example of equivariance in the MPE Tag environment where pursuers chase evaders while passing through obstacles. (Middle to left) As the agents (circles) are permuted by swapping their indices ($s\in S_2$), the optimal actions (arrows) are permuted in the same way. (Middle to right) As the agent positions are rotated $90^\circ$ ($g\in SO(2)$), the optimal actions are also rotated.

Overview

Swarm robotics needs control policies that are sample-efficient to train, generalize across different team sizes, and hold up in the real world. Standard MARL policies built on multi-layer perceptrons ignore the combinatorial and geometric structure of these tasks, so they scale poorly as the number of agents grows and rarely transfer to configurations they were not trained on.

LEGO (Local-Canonicalization Equivariant Graph neural networks) bakes the right inductive biases directly into the policy by decoupling the two symmetries of the problem:

• Permutation symmetry is handled by a graph neural network over the agents.
• Euclidean symmetry — $E(n)$-equivariance — is handled by local canonicalization of each agent’s inputs.
• Role structure (e.g., pursuers vs. evaders) is captured with heterogeneous graphs.

Because it is modular, LEGO drops into off-the-shelf MARL algorithms such as MAPPO, yielding LEGO-MAPPO.

Why it’s hard

• Curse of many agents. MLP policies scale badly as the agent count grows; a policy trained with $N$ agents typically fails when deployed with a different number of agents $M$, forcing costly retraining for every swarm configuration.
• Sample inefficiency. MARL often fails to exploit the geometric symmetries inherent in robotics tasks, wasting samples relearning equivalent situations.
• Competitive instability. In pursuer–evader settings, the simultaneous adaptation of both sides can destabilize training.

Method: LEGO

LEGO’s design philosophy is to decouple the system’s symmetries and handle each with the right tool. Each agent first canonicalizes its local frame, which removes $E(2)$ nuisance variation before encoding; a role-wise graph network then encodes the (sub)graphs into a state representation; actions are produced in the local frame and mapped back to the global frame via the agent’s orientation. The whole pipeline is trained end-to-end with MAPPO under centralized-training / decentralized-execution.

Because canonicalization removes $E(2)$ variation before encoding, the resulting policy is $E(2)$-equivariant by construction, and the graph structure makes it permutation-equivariant within each role.

Figure 2. Using LEGO in MARL. For each agent $i$, its raw observation $O_i$ is canonicalized into the local frame $\mathcal{C}(O_i)=\{v_i',\rho(g_i^{-1})O_i\}$. Role-based subgraphs (e.g., self, pursuers, evaders, obstacles) are encoded with Graphormer, pooled by role, and concatenated into $s_i$. The policy outputs a local action $a_i^{\text{loc}} \sim \pi_\theta(s_i)$, which is transformed back via $a_i = R_i a_i^{\text{loc}}$. Under CTDE, the critic uses the same pipeline but operates on the global state $X$ during training to obtain $s_i^{\text{global}}$, from which individual values are estimated as $V_i = V_\phi(s_i^{\text{global}})$.

Results

Cooperative Tasks: MPE Spread^[1]

Qualitative rollouts on the cooperative MPE Spread benchmark, where agents must spread out to cover all landmarks while avoiding collisions. LEGO-MAPPO is compared against MAPPO, PIC, and SEGNN at two team sizes.

Competitive Tasks: MPE Tag-occlusion

Zero-shot Scalability

Because LEGO is built on a graph neural network, it naturally handles a varying number of input nodes, so a policy trained at one team size can be applied directly to another. We train LEGO-MAPPO on MPE Spread with 4 agents and 4 landmarks, then evaluate that same policy — with no retraining or fine-tuning — on systems with 2, 3, 5, and 6 agents.

Curriculum Learning

LEGO’s scalability also enables a curriculum: a policy trained on a smaller swarm serves as a warm start for a larger, harder configuration. Here LEGO-MAPPO is pre-trained on the 4-agent system and then trained on the target systems of 6, 7, and 8 agents.

Real-World Experiments

Robustness to agent failure. To show that zero-shot scalability translates into real-world robustness, we run a scenario with horizon $T=100$ in which one pursuer deliberately breaks down and lands at $t=30$. The remaining pursuer keeps chasing the evader, and the team stays functional despite the failure: although the “broken” pursuer is inactive, it still serves as a roadblock that prevents the evader from approaching its vicinity, while the surviving pursuer switches its strategy from blocking the evader’s path to directly chasing it.

Citation

@article{wang2025lego,
  title   = {Local-Canonicalization Equivariant Graph Neural Networks for Sample-Efficient and Generalizable Swarm Robot Control},
  author  = {Wang, Keqin and Zhong, Tao and Chang, David and Allen-Blanchette, Christine},
  journal = {arXiv preprint arXiv:2509.14431},
  year    = {2025}
}