Local-Canonicalization Equivariant Graph Neural Networks for Sample-Efficient and Generalizable Swarm Robot Control
This work addresses the problem of inefficient learning and poor generalization in multi-agent reinforcement learning for swarm control, offering a method that improves sample efficiency and transferability, which is significant for robotics applications.
The paper introduces LEGO, a modular policy architecture for swarm robot control that combines agent-centric canonicalization with role-aware graph encoding to achieve E(2)-equivariance and intra-role permutation equivariance. LEGO-MAPPO improves sample efficiency and task performance over baselines in MPE benchmarks, transfers to unseen team sizes and distribution shifts, and remains operational in Crazyflie experiments after a pursuer is disabled.
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.