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

arXiv: https://arxiv.org/abs/2509.14431 · HTML: https://mlsanity.com/paper/2509.14431/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Multi-agent reinforcement learning, one of 15 directions in Deep reinforcement learning methods.
- Builds on 47 of its 67 references that are in the map.
- First posted on arXiv 17 Sep 2025.

## Placed in

- Deep reinforcement learning methods / [Multi-agent reinforcement learning](https://mlsanity.com/deeprl/multi-agent-reinforcement-learning/)

## Builds on (most cited first, 30 of 47)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Semi-Supervised Classification with Graph Convolutional Networks](https://mlsanity.com/paper/1609.02907/) (2016)
- [Graph Attention Networks](https://mlsanity.com/paper/1710.10903/) (2017)
- [A Comprehensive Survey on Graph Neural Networks](https://mlsanity.com/paper/1901.00596/) (2019)
- [Multi-Agent Actor-Critic for Mixed Cooperative-Competitive Environments](https://mlsanity.com/paper/1706.02275/) (2017)
- [Geometric Deep Learning: Grids, Groups, Graphs, Geodesics, and Gauges](https://mlsanity.com/paper/2104.13478/) (2021)
- [Group Equivariant Convolutional Networks](https://mlsanity.com/paper/1602.07576/) (2016)
- [Graph Convolutional Neural Networks for Web-Scale Recommender Systems](https://mlsanity.com/paper/1806.01973/) (2018)
- [E(n) Equivariant Graph Neural Networks](https://mlsanity.com/paper/2102.09844/) (2021)
- [The Surprising Effectiveness of PPO in Cooperative, Multi-Agent Games](https://mlsanity.com/paper/2103.01955/) (2021)
- [Value-Decomposition Networks For Cooperative Multi-Agent Learning](https://mlsanity.com/paper/1706.05296/) (2017)
- [Monotonic Value Function Factorisation for Deep Multi-Agent Reinforcement Learning](https://mlsanity.com/paper/2003.08839/) (2020)
- [Spherical CNNs](https://mlsanity.com/paper/1801.10130/) (2018)
- [Steerable CNNs](https://mlsanity.com/paper/1612.08498/) (2016)
- [Emergence of Grounded Compositional Language in Multi-Agent Populations](https://mlsanity.com/paper/1703.04908/) (2017)
- [Learning SO(3) Equivariant Representations with Spherical CNNs](https://mlsanity.com/paper/1711.06721/) (2017)
- [Geometric and Physical Quantities Improve E(3) Equivariant Message Passing](https://mlsanity.com/paper/2110.02905/) (2021)
- [PettingZoo: Gym for Multi-Agent Reinforcement Learning](https://mlsanity.com/paper/2009.14471/) (2020)
- [Teacher-Student Curriculum Learning](https://mlsanity.com/paper/1707.00183/) (2017)
- [Graph Convolutional Reinforcement Learning](https://mlsanity.com/paper/1810.09202/) (2018)
- [Deep Decentralized Multi-task Multi-Agent Reinforcement Learning under Partial Observability](https://mlsanity.com/paper/1703.06182/) (2017)
- [MDP Homomorphic Networks: Group Symmetries in Reinforcement Learning](https://mlsanity.com/paper/2006.16908/) (2020)
- [ROMA: Multi-Agent Reinforcement Learning with Emergent Roles](https://mlsanity.com/paper/2003.08039/) (2020)
- [SO(2)-Equivariant Reinforcement Learning](https://mlsanity.com/paper/2203.04439/) (2022)
- [Attending to Graph Transformers](https://mlsanity.com/paper/2302.04181/) (2023)
- [SE(3) Equivariant Graph Neural Networks with Complete Local Frames](https://mlsanity.com/paper/2110.14811/) (2021)
- [Learning Safe Multi-Agent Control with Decentralized Neural Barrier Certificates](https://mlsanity.com/paper/2101.05436/) (2021)
- [Heterogeneous-Agent Reinforcement Learning](https://mlsanity.com/paper/2304.09870/) (2023)
- [Benchmarking Graphormer on Large-Scale Molecular Modeling Datasets](https://mlsanity.com/paper/2203.04810/) (2022)
- [PIC: Permutation Invariant Critic for Multi-Agent Deep Reinforcement Learning](https://mlsanity.com/paper/1911.00025/) (2019)

## Abstract (arXiv)

Multi-agent reinforcement learning (MARL) policies for swarm control often learn inefficiently and generalize poorly across coordinate frames, team sizes, and agent roles. We introduce Local-Canonicalization Equivariant Graph Neural Networks (LEGO), a modular policy architecture that combines agent-centric canonicalization with role-aware graph encoding. Canonicalization removes dependence on the global coordinate frame, while transforming predicted local actions back to the world frame produces an E(2)-equivariant policy. Role-wise graph encoders provide intra-role permutation equivariance and fixed-dimensional representations for variable-size teams. LEGO can be paired with standard MARL algorithms; we instantiate it with MAPPO. Across cooperative MPE Spread and competitive Tag-occlusion benchmarks, LEGO-MAPPO improves sample efficiency and task performance relative to MLP-based, graph-only, canonicalization-only, and equivariant baselines. The learned policies transfer without fine-tuning to unseen team sizes, maintain performance under spatial distribution shifts, and benefit from curriculum initialization for larger teams. In Crazyflie experiments, the policy remains operational after one pursuer is disabled. Code is available at https://github.com/CAB-Lab-Princeton/LEGO-MARL.