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

2025/09/17 by Keqin Wang, Wang, Keqin, Tao Zhong +6
Computer Science · #Adaptive Dynamic Programming Control #Adversarial Robustness in Machine Learning #FOS: Computer and information sciences #Reinforcement Learning in Robotics #Robotics (cs.RO) #cs.RO

paper · pdf · doi:10.48550/arxiv.2509.14431

Accepted at IROS 2026

openalex publication_date 2025/09/17 · openalex created_date 2025/10/10 · openalex updated_date 2026/07/28 · arxiv created 2026/07/31 · arxiv updated 2026/08/04

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 https://github.com/CAB-Lab-Princeton/LEGO-MARL.

Citations

• PEnGUiN: Partially Equivariant Graph NeUral Networks for Sample Efficient MARL
• SE(3)-Equivariant Robot Learning and Control: A Tutorial Survey
• GAGrasp: Geometric Algebra Diffusion for Dexterous Grasping
• Leveraging Graph Neural Networks and Multi-Agent Reinforcement Learning for Inventory Control in Supply Chains
• Boosting Sample Efficiency and Generalization in Multi-agent Reinforcement Learning via Equivariance
• Equivariant Reinforcement Learning under Partial Observability
• Behavior-Inspired Neural Networks for Relational Inference
• Learning Color Equivariant Representations
• Beyond Canonicalization: How Tensorial Messages Improve Equivariant Message Passing
• Graph Neural Network-based Multi-agent Reinforcement Learning for Resilient Distributed Coordination of Multi-Robot Systems
• Multi-agent Reinforcement Learning: A Comprehensive Survey
• \rm E(3)-Equivariant Actor-Critic Methods for Cooperative Multi-Agent Reinforcement Learning
• Heterogeneous-Agent Reinforcement Learning
• Attending to Graph Transformers
• A Survey on Large-Population Systems and Scalable Multi-Agent Reinforcement Learning
• Benchmarking Graphormer on Large-Scale Molecular Modeling Datasets
• SO(2)-Equivariant Reinforcement Learning
• Coach-Player Multi-Agent Reinforcement Learning for Dynamic Team Composition
• Geometric Deep Learning: Grids, Groups, Graphs, Geodesics, and Gauges
• The Surprising Effectiveness of PPO in Cooperative, Multi-Agent Games
• PettingZoo: Gym for Multi-Agent Reinforcement Learning
• MDP Homomorphic Networks: Group Symmetries in Reinforcement Learning
• ROMA: Multi-Agent Reinforcement Learning with Emergent Roles
• Learning Transferable Cooperative Behavior in Multi-Agent Teams
• Equivariant Multi-View Networks
• Graph Convolutional Reinforcement Learning
• Spherical CNNs
• Learning SO(3) Equivariant Representations with Spherical CNNs
• Graph Attention Networks
• Teacher-Student Curriculum Learning
• Value-Decomposition Networks For Cooperative Multi-Agent Learning
• Attention Is All You Need
• Multi-Agent Actor-Critic for Mixed Cooperative-Competitive Environments
• Emergence of Grounded Compositional Language in Multi-Agent Populations
• Semi-Supervised Classification with Graph Convolutional Networks
• A Comprehensive Survey on Graph Neural Networks

Related

• Heterogeneous Multi-Robot Reinforcement Learning
• PEnGUiN: Partially Equivariant Graph NeUral Networks for Sample Efficient MARL
• Boosting Sample Efficiency and Generalization in Multi-agent Reinforcement Learning via Equivariance
• Graph Neural Network-based Multi-agent Reinforcement Learning for Resilient Distributed Coordination of Multi-Robot Systems
• Using Reinforcement Learning to Herd a Robotic Swarm to a Target\n Distribution