Embedding Morphology into Transformers for Cross-Robot Policy Learning

TL;DR

This paper introduces an embodiment-aware transformer policy for cross-robot learning, integrating morphology through kinematic tokens, topology-aware attention, and joint attributes to enhance robustness and performance.

Contribution

It presents a novel transformer architecture that explicitly incorporates robot morphology, improving cross-embodiment policy learning in robotics.

Findings

Improved robustness across multiple robot embodiments.

Enhanced performance over baseline models within single embodiments.

Structured morphology integration benefits in policy generalization.

Abstract

Cross-robot policy learning -- training a single policy to perform well across multiple embodiments -- remains a central challenge in robot learning. Transformer-based policies, such as vision-language-action (VLA) models, are typically embodiment-agnostic and must infer kinematic structure purely from observations, which can reduce robustness across embodiments and even limit performance within a single embodiment. We propose an embodiment-aware transformer policy that injects morphology via three mechanisms: (1) kinematic tokens that factorize actions across joints and compress time through per-joint temporal chunking; (2) a topology-aware attention bias that encodes kinematic topology as an inductive bias in self-attention, encouraging message passing along kinematic edges; and (3) joint-attribute conditioning that augments topology with per-joint descriptors to capture semantics beyond connectivity. Across a range of embodiments, this structured integration consistently improves performance over a vanilla pi0.5 VLA baseline, indicating improved robustness both within an embodiment and across embodiments.