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

Abstract: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 this https URL.

Robotics (cs.RO)

arXiv:2509.14431 [cs.RO]

(or arXiv:2509.14431v2 [cs.RO] for this version)

https://doi.org/10.48550/arXiv.2509.14431 Focus to learn more arXiv-issued DOI via DataCite

arXivLabs is a framework that allows collaborators to develop and share new arXiv features directly on our website.

Both individuals and organizations that work with arXivLabs have embraced and accepted our values of openness, community, excellence, and user data privacy. arXiv is committed to these values and only works with partners that adhere to them.

Have an idea for a project that will add value for arXiv's community? Learn more about arXivLabs.