Per-Link Attention for Cross-Embodiment Graph-Conditioned Diffusion Policies
Abstract
Cross-embodiment robot policies encode a robot's kinematic tree as a graph and run a graph neural network over it, producing per-node (per-link) embeddings that are then pooled into a single fixed-size morphology vector before conditioning the policy. We show that this graph-level pooling step -- standard in GNN pipelines across domains -- collapses morphology-discriminative node-level information into an actuation-magnitude default. A seven-intervention conditioning audit applied to two independently trained pipelines yields sign-unanimous evidence (5/5 seeds, p ~ 0.031 each) that the pooled graph embedding fails to function as a structured morphology code. We propose Per-Link Attention Diffusion (PLA-D), an architecture that routes per-node GNN embeddings directly into a Transformer denoiser's cross-attention as a variable-length key/value sequence with learned node-position embeddings and a key-padding mask. PLA-D preserves |V| * d_h node-level dimensions that pooling discards, strictly subsumes the pooled baseline, and trains to comparable validation loss in less than half the optimisation budget. The finding that graph-level pooling destroys task-relevant node structure generalises beyond robotics to any graph foundation model pipeline that pools node embeddings before a downstream task head.