LEGO: Local-Canonicalization Equivariant GNNs

Updated 12 July 2026

The paper introduces LEGO as a modular GNN framework that enforces permutation and E(n)-equivariance by converting node observations into locally defined canonical frames.
It combines standard message passing with efficient frame transport and role-specific representations, replacing expensive tensor operations with frame-based computations.
Empirical results show LEGO achieves faster convergence, improved scalability, and enhanced performance in swarm robotics and molecular property prediction compared to baselines.

Local-Canonicalization Equivariant Graph Neural Networks (LEGO) are graph-based architectures that enforce symmetry by mapping node- or agent-centered observations into locally defined canonical coordinates before applying otherwise standard message passing, and by transforming outputs back to the global frame when equivariant outputs are required. In the 2025 swarm-robotics formulation, LEGO is a modular policy architecture for multi-agent reinforcement learning (MARL) that combines graph neural networks for permutation equivariance, canonicalization for $E(n)$ -equivariance, and heterogeneous representations for role-specific inductive biases (Wang et al., 17 Sep 2025). More broadly, recent work on local canonicalization develops the same principle for molecular and geometric learning, especially for $O(3)$ , $SE(3)$ , and Lorentz-equivariant graph models, often with the explicit goal of replacing expensive tensor operations by frame-based canonicalization and transport (Gerhartz et al., 30 Sep 2025, Spinner et al., 26 May 2025). This places LEGO within a wider line of research on learned or computed canonicalization as a route to exact equivariance (Kaba et al., 2022).

1. Conceptual basis and scope

The central premise of LEGO is that many graph-learning problems exhibit symmetries that should be hard-coded into the model. In swarm robotics, the relevant symmetries are permutation symmetry over interchangeable agents and Euclidean symmetry over global translations and rotations of the swarm. The swarm-control LEGO framework is designed to realize both simultaneously: a GNN encoder handles permutation equivariance and varying agent number, while local canonicalization removes global $E(2)$ variation by expressing each agent’s observation in a consistent agent-centric frame (Wang et al., 17 Sep 2025).

In the more general local-canonicalization literature, the same idea is formulated for group actions $G$ . A representation $\rho$ of $G$ on a vector space $V$ is a homomorphism $\rho:G\to GL(V)$ satisfying $\rho(g_1g_2)=\rho(g_1)\rho(g_2)$ , and a map $O(3)$ 0 is $O(3)$ 1-equivariant if $O(3)$ 2 for all $O(3)$ 3 (Gerhartz et al., 30 Sep 2025). Local canonicalization assigns to each node $O(3)$ 4 a local frame $O(3)$ 5, canonicalizes features by $O(3)$ 6, and then performs message passing in canonical space. Because the frame transforms covariantly under the global group action, the canonicalized features become invariant to the global transform (Gerhartz et al., 30 Sep 2025).

A related precursor for geometric deep learning proposed $O(3)$ 7-equivariant graph neural networks with equivariant local complete frames, introducing orthonormal local frames inspired by differential geometry and physics, and emphasizing efficient approximation of geometric quantities through frame construction based only on cross products (Du et al., 2021). That work did not present the later local-canonicalization formalism, but it established the importance of explicit local frames for efficient equivariant graph processing.

2. Canonicalization and equivariance

A general canonicalization formulation factors an equivariant map into a canonicalizer and an unconstrained predictor: $O(3)$ 8 Here $O(3)$ 9 selects a canonicalizing group element and equivariance follows if $SE(3)$ 0 itself is $SE(3)$ 1-equivariant, i.e.

$SE(3)$ 2

Under continuity assumptions, and provided $SE(3)$ 3 is a universal approximator, this factorization yields a universality result for continuous $SE(3)$ 4-equivariant maps (Kaba et al., 2022). This provides a theoretical basis for using learned or computed canonicalizers in place of specialized equivariant layers.

For local-canonicalization GNNs on geometric data, canonicalization is performed per node. At node $SE(3)$ 5, one predicts or computes a local frame $SE(3)$ 6, transforms global features $SE(3)$ 7 into canonical features $SE(3)$ 8, and then applies a standard permutation-equivariant but $SE(3)$ 9-invariant backbone to $E(2)$ 0. The key equivariance mechanism is that under a global action $E(2)$ 1, the frame transforms as $E(2)$ 2, making the canonicalized feature invariant to the global transform (Gerhartz et al., 30 Sep 2025).

The swarm-control LEGO instantiation specializes this construction to $E(2)$ 3. For agent $E(2)$ 4 with velocity $E(2)$ 5 and position $E(2)$ 6, the local $E(2)$ 7-axis is aligned with velocity,

$E(2)$ 8

the center of mass is $E(2)$ 9, $G$ 0, and the local $G$ 1-axis is

$G$ 2

with $G$ 3. The resulting frame is $G$ 4. The canonicalized observation is

$G$ 5

where

$G$ 6

By construction,

$G$ 7

and a global policy $G$ 8 is $G$ 9-equivariant if

$\rho$ 0

LEGO satisfies this by computing a local action $\rho$ 1 from canonicalized inputs and then rotating it back via $\rho$ 2 (Wang et al., 17 Sep 2025).

3. Message passing, transport, and representation choices

Local canonicalization does not remove the need to communicate between nodes with different frames; it changes the form of that communication. In the general LEGO message-passing formulation, information from node $\rho$ 3 must be transported from frame $\rho$ 4 to frame $\rho$ 5. The node-wise update is

$\rho$ 6

where $\rho$ 7 is the frame-to-frame transformation, $\rho$ 8 is the internal feature representation, and $\rho$ 9 expresses relative position in the $G$ 0-th canonical frame (Gerhartz et al., 30 Sep 2025).

This formulation is intended as a computationally lighter alternative to tensor-field convolutions. In a conventional Tensor Field Network, messages couple irreducible features and spherical harmonics through Clebsch–Gordan coefficients. In LEGO, the expensive tensor product is replaced by a composition of frame transport, radial embedding $G$ 1, angular embedding $G$ 2, and pointwise MLP or linear layers. A single EDGE message layer is

$G$ 3

with learned linear maps $G$ 4 and element-wise multiplication $G$ 5 (Gerhartz et al., 30 Sep 2025).

The representation $G$ 6 can be scalar, Cartesian tensor, irreducible, or a learned MLP representation. The trade-off between compactness and transform cost is explicit:

Representation	# components	Total transform cost per frame
Irreducible ( $G$ 7)	$G$ 8	$G$ 9
Cartesian tensor (order $V$ 0)	$V$ 1	$V$ 2

For small $V$ 3, irreducible representations are compact but incur a one-time $V$ 4 cost per rotation to form Wigner- $V$ 5 matrices, whereas Cartesian representations have more components but cheap per-component transforms. The reported conclusion is that, in practice, for $V$ 6 up to $V$ 7 or $V$ 8, Cartesian representations can be faster if the Wigner computation dominates (Gerhartz et al., 30 Sep 2025).

4. LEGO as a MARL policy architecture

In swarm robotics, LEGO models the multi-agent system as a graph $V$ 9, with $\rho:G\to GL(V)$ 0 and edges representing communication. Each node $\rho:G\to GL(V)$ 1 has an initial feature vector $\rho:G\to GL(V)$ 2 encoding its local canonicalized observation. Within each role $\rho:G\to GL(V)$ 3, such as pursuer, evader, or obstacle, a subgraph $\rho:G\to GL(V)$ 4 is formed on node set $\rho:G\to GL(V)$ 5, and a Graphormer-style multi-head attention encoder is applied for $\rho:G\to GL(V)$ 6 layers (Wang et al., 17 Sep 2025).

At layer $\rho:G\to GL(V)$ 7 and head $\rho:G\to GL(V)$ 8, the attention logits are

$\rho:G\to GL(V)$ 9

and the update is

$\rho(g_1g_2)=\rho(g_1)\rho(g_2)$ 0

Because these operations depend on pairwise dot-products and neighborhood sums, the encoder is permutation-equivariant. After $\rho(g_1g_2)=\rho(g_1)\rho(g_2)$ 1 layers, role-wise pooling produces summaries

$\rho(g_1g_2)=\rho(g_1)\rho(g_2)$ 2

whose dimension depends on the number of roles rather than on $\rho(g_1g_2)=\rho(g_1)\rho(g_2)$ 3, which is the stated mechanism by which LEGO scales to varying swarm size (Wang et al., 17 Sep 2025).

The architecture is integrated with the centralized-training decentralized-execution (CTDE) paradigm and is described as plugging into MAPPO with minimal modifications. The actor $\rho(g_1g_2)=\rho(g_1)\rho(g_2)$ 4 outputs a continuous local action $\rho(g_1g_2)=\rho(g_1)\rho(g_2)$ 5, the critic $\rho(g_1g_2)=\rho(g_1)\rho(g_2)$ 6 is a centralized value function on the pooled representation, and the MAPPO loss remains unchanged. The only architectural additions are the canonicalization layer computing $\rho(g_1g_2)=\rho(g_1)\rho(g_2)$ 7 and $\rho(g_1g_2)=\rho(g_1)\rho(g_2)$ 8, and the GNN encoder. No extra equivariance losses are needed because equivariance is built into the architecture (Wang et al., 17 Sep 2025).

A distinctive feature of the swarm formulation is its heterogeneous, role-based representation. Nodes share Graphormer parameters within a role subgraph, while distinct role embeddings or biases may be added to $\rho(g_1g_2)=\rho(g_1)\rho(g_2)$ 9. This is intended to preserve permutation symmetry within each role while allowing role-specific inductive priors (Wang et al., 17 Sep 2025).

5. Empirical behavior across domains

In molecular property prediction, local-canonicalization graph models were evaluated on QM9 using $O(3)$ 00 molecules for training. A LoCaFormer with four internal representations—scalar messages, Cartesian tensor, irreducible, and learned MLP—was compared with Equiformer and MACE. The reported MAE and speed for selected properties were: scalar messages, $O(3)$ 01, $O(3)$ 02, $O(3)$ 03 it/s; Cartesian tensor, $O(3)$ 04, $O(3)$ 05, $O(3)$ 06 it/s; irreducible, $O(3)$ 07, $O(3)$ 08, $O(3)$ 09 it/s; MLP, $O(3)$ 10, $O(3)$ 11, $O(3)$ 12 it/s; Equiformer, $O(3)$ 13, $O(3)$ 14, $O(3)$ 15 it/s. The reported interpretation is that Cartesian and irreducible messages yield better accuracy on geometry-sensitive targets such as $O(3)$ 16 and $O(3)$ 17, and that all LEGO variants run $O(3)$ 18– $O(3)$ 19 faster than Equiformer on the same hardware, an NVIDIA RTX 6000 (Gerhartz et al., 30 Sep 2025).

On tensorial property prediction for the N-methylacetamide dataset with $O(3)$ 20 samples, the targets were the dipole $O(3)$ 21 and polarizability $O(3)$ 22. The RMSE values were reported as follows: scalar, $O(3)$ 23 a.u., $O(3)$ 24 a.u.; Cartesian, $O(3)$ 25, $O(3)$ 26; irreducible, $O(3)$ 27, $O(3)$ 28; MLP, $O(3)$ 29, $O(3)$ 30. The stated conclusion is that tensorial message passing, Cartesian or irreducible, clearly outperforms scalar or learned MLP for tensorial targets. In an explicit comparison against data augmentation, an equivariant LoCaFormer converged faster with more data, showing a steeper error–data curve, while in the very low-data regime the augmented model sometimes had slightly lower error (Gerhartz et al., 30 Sep 2025).

In swarm MARL, the cooperative MPE Spread task was evaluated in configurations with $O(3)$ 31 agents and $O(3)$ 32 landmarks, and $O(3)$ 33 agents and $O(3)$ 34 landmarks. The baselines were MAPPO with an MLP actor and critic on raw observations, MAPPO-local with canonicalization only, and MAPPO-GNN with a GNN only. The reported result is that LEGO-MAPPO converges fastest and achieves highest reward, with MAPPO-local $O(3)$ 35 MAPPO, MAPPO-GNN $O(3)$ 36 MAPPO, and LEGO $O(3)$ 37 both (Wang et al., 17 Sep 2025).

For the competitive Tag-Occlusion task, involving $O(3)$ 38 pursuers versus $O(3)$ 39 evaders plus $O(3)$ 40 obstacles under partial observability via occlusion, LEGO-MAPPO was reported to yield stable pursuit strategies, whereas MAPPO baselines oscillated or failed to converge. In cross-validation, pairing LEGO-trained pursuers against MAPPO evaders, and conversely, showed a clear performance gap. On zero-shot scalability, training on $O(3)$ 41 in Spread and testing on $O(3)$ 42, LEGO-MAPPO generalized strongly and outperformed MAPPO-GNN and MAPPO, even when the latter was retrained on the target $O(3)$ 43. On out-of-distribution geometric generalization, training on left-biased initializations and testing on right-biased and uniform initializations, MAPPO-local and LEGO-MAPPO maintained performance while standard MAPPO degraded, and LEGO achieved the best absolute rewards. In a real-world Crazyflie demonstration with $O(3)$ 44 pursuers, $O(3)$ 45 evader, and $O(3)$ 46 physical obstacles, one pursuer was forced to break down at $O(3)$ 47; LEGO-MAPPO adapted by having the remaining agent pursue while the failed drone served as a blocker (Wang et al., 17 Sep 2025).

A recurrent misconception is that local canonicalization is only a variant of data augmentation. A broader local-canonicalization perspective, developed for Lorentz symmetry, distinguishes three cases: no canonicalization, corresponding to $O(3)$ 48; full data augmentation, corresponding to a single random global frame $O(3)$ 49 per minibatch; and exact equivariance, corresponding to an input-dependent canonicalizer $O(3)$ 50. The same work states that data augmentation emerges as a special choice of reference frame. This suggests that LEGO-style architectures are best understood not as augmentation heuristics but as exact symmetry constructions when the canonicalizer is itself equivariant (Spinner et al., 26 May 2025).

A second misconception is that exact equivariance necessarily requires specialized tensor-product layers. Local-canonicalization work on molecular and geometric data explicitly aims to transfer existing tensor field networks into a more efficient local canonicalization paradigm while preserving equivariance. The accompanying implementation, the tensor_frames package, is a PyPI-installable library built on PyTorchGeometric $O(3)$ 51, with modules tensor_frames.lframes for local-frame predictors, tensor_frames.reps for Irreps, TensorReps, and MLPReps, and tensor_frames.nn.TFMessagePassing as a wrapper over torch_geometric.nn.MessagePassing. The stated API goal is that user-defined message() logic need not change: canonicalization, frame transport, position transformation, and collection of outputs in canonical space are handled automatically (Gerhartz et al., 30 Sep 2025).

A third point, emphasized in the learned-canonicalization literature, is that canonicalization need not be hard-coded by heuristics. Learned canonicalization functions can be implemented directly by small equivariant networks or indirectly through an optimization problem $O(3)$ 52, and the main empirical hypothesis of that line is that learning a small neural network to perform canonicalization is better than predefined heuristics. The same framework reports universality and interpretable “mental rotation”-style behavior, where the prediction network only sees canonicalized inputs (Kaba et al., 2022).

In the swarm-robotics formulation, the stated future extensions are full $O(3)$ 53 canonicalization for $O(3)$ 54D swarms and tensorial message passing for continuous group representations (Wang et al., 17 Sep 2025). More generally, the recent literature suggests a convergence between application-specific LEGO policies, frame-based geometric GNNs, and learned canonicalization theory: all treat equivariance as the composition of local frame selection, canonicalized computation, and, when necessary, de-canonicalization to the original coordinate system.