Embedding Morphology into Transformers for Cross-Robot Policy Learning

This paper proposes an embodiment-aware transformer policy that enhances cross-robot learning by integrating kinematic morphology through joint-specific tokens, topology-aware attention biases, and per-joint attribute conditioning, thereby improving robustness and performance across diverse robot embodiments compared to standard VLA baselines.

The Gist

Imagine you are trying to teach a group of very different robots how to cook a meal. You have a tiny, nimble robot arm (like a human hand), a bulky, heavy-duty robot arm (like a construction crane), and a robot with a completely different number of fingers.

Currently, most "AI chefs" (robot brains) are trained like this: You show them a video of a human cooking, and the AI has to figure out everything on its own. It has to guess how many joints the robot has, which way they bend, and how they work together just by looking at the video. This is like asking someone to learn how to drive a car, a motorcycle, and a bicycle just by watching a video of a person driving, without ever being told what the steering wheel or handlebars are. It's confusing, slow, and the AI often fails when you swap the vehicle.

This paper proposes a smarter way: Give the AI a blueprint of the robot's body before it starts learning.

Here is how they did it, using three simple tricks:

1. The "Body Map" (Kinematic Tokens)

The Problem: Standard AI treats a robot's movements as a giant, messy list of numbers. It doesn't know that "Joint A" is connected to "Joint B."
The Solution: The authors broke the robot's movement down into individual "body parts." Instead of one big blob of data, they gave the AI a specific token (a little note) for each joint.

• Analogy: Imagine a conductor leading an orchestra. Instead of hearing a wall of noise, the conductor gives a specific sheet of music to the violin section, the drum section, and the trumpet section separately. Now, the AI knows exactly which "musical note" belongs to which "joint."

2. The "Social Network" (Topology-Aware Attention)

The Problem: In a normal AI brain, every part of the robot can talk to every other part instantly. But in real life, a robot's elbow can't directly talk to its shoulder without going through the upper arm. The AI was wasting energy trying to connect things that aren't physically linked.
The Solution: They built a "social network" rule into the AI. They told it: "You can only chat with your immediate neighbors (joints connected by a bone) unless you really need to reach out to the whole group."

• Analogy: Think of a game of "Telephone." If you are in a line of people, you only pass the message to the person standing right next to you. This paper tells the AI to mostly pass messages to neighbors (like a local gossip chain) but occasionally let the message jump to the whole group (global coordination) so the robot doesn't get stuck in a loop. This makes the robot move much more naturally.

3. The "ID Badge" (Joint-Attribute Conditioning)

The Problem: Even if two robots have the same "shape" (topology), their parts might act differently. One robot's joint might be a spinning wheel; another's might be a sliding piston. The AI needs to know what the part is, not just where it is.
The Solution: They gave every joint an "ID Badge" with details like "I am a spinning joint," "I can only turn 90 degrees," or "I am very slippery."

• Analogy: Imagine a sports team. Knowing who is standing next to whom (the topology) is good. But knowing that Player A is a "Goalkeeper" and Player B is a "Striker" (the attributes) is what actually helps the team win. This extra info helps the AI understand the specific rules of each robot's body.

The Results: A Super-Adaptable Robot

When they tested this new "Body-Aware" AI:

• It learned faster: It didn't have to guess how the robot worked.
• It was more robust: If you swapped the robot for a different model (e.g., from a small arm to a big arm), the AI didn't crash. It just looked at the new "blueprint" and adapted.
• It worked better on a single robot too: Even if you only used one type of robot, this method made it perform better than the standard AI.

In a nutshell:
Current robot AI is like a student trying to learn anatomy by staring at a blurry photo. This paper gives the student a clear 3D model, a map of how the bones connect, and a textbook describing what each bone does. The result? The robot learns to move much faster, safer, and can easily switch between different robot bodies without needing to start over.

Technical Summary

1. Problem Statement

Cross-robot policy learning aims to train a single policy capable of performing tasks across multiple robot embodiments (different hardware configurations). While Vision-Language-Action (VLA) models like $\pi_0.5$ have achieved success by scaling on large datasets, they suffer from a critical limitation: they are embodiment-agnostic.

• The Challenge: Standard VLA models must implicitly infer kinematic structures (how joints connect and move) purely from visual observations. This lack of explicit structural knowledge reduces robustness when transferring policies between robots with different morphologies (e.g., different numbers of joints, link lengths, or actuation types) and can even limit performance within a single robot type.
• Existing Gaps: Prior attempts to encode morphology (e.g., using Graph Neural Networks or topology-aware attention) face three specific hurdles in modern VLA architectures:
  1. Token Interface Mismatch: State-of-the-art VLAs compress joint actions into compact tokens, making it difficult to apply per-joint morphology embeddings.
  2. Local-Global Trade-off: Enforcing strict locality (1-hop neighbors) aids kinematic message passing but hinders long-range coordination; conversely, full attention ignores structure.
  3. Missing Semantics: Existing methods encode connectivity (topology) but ignore per-joint semantics (e.g., joint type, axis direction, limits), which are crucial for functional differentiation.

2. Methodology

The authors propose an Embodiment-Aware Transformer Policy that injects robot morphology into the VLA architecture via three distinct mechanisms:

A. Kinematic Tokens (KT)

To address the token interface mismatch, the authors introduce a new token type alongside standard action tokens.

• Mechanism: Instead of compressing the entire action horizon into a single token per timestep, the action sequence is split into temporal chunks. For each joint $j$ and chunk $k$, the actions are concatenated into a vector $b_{j,T_k}$.
• Function: These vectors are projected into embeddings (Kinematic Tokens) that provide a joint-wise representation. This decouples the spatial structure from the temporal compression, allowing the model to explicitly attend to specific joints and their relationships.
• Extension: They also introduce Auxiliary Kinematic Tokens (AKT), generating multiple embeddings per joint to increase representational capacity.

B. Topology-Aware Attention Bias

To address the local-global trade-off, the authors modify the self-attention mechanism using the robot's kinematic graph (nodes = joints, edges = physical connections).

• Hard-Mask Family:
  • Full-Mask: Strictly restricts attention to a joint's 1-hop neighbors (and itself) at every layer.
  • Mix-Mask: Alternates between masked layers (local) and full-attention layers (global). This balances kinematic message passing with long-range coordination.
• Soft-Mask Family: Uses a learnable bias based on the shortest-path distance (SPD) between joints in the kinematic graph. This allows attention to all pairs but biases the model toward kinematically closer joints.
• Finding: The Mix-Mask strategy (Hard-Mask family) proved most effective, avoiding the over-restriction of Full-Mask and the optimization instability of Soft-Mask.

C. Joint-Attribute Conditioning

To address the lack of joint semantics, the authors augment the topology with per-joint descriptors.

• Mechanism: Each joint is described by a feature vector $s_j$ containing:
  • Joint type (prismatic vs. revolute).
  • Axis direction ($x, y, z$).
  • Motion limits (hard upper/lower bounds).
  • Contact properties (damping, friction, stiffness).
• Implementation: These descriptors are processed via Feature-wise Linear Modulation (FiLM) to generate scale ($\gamma$) and shift ($\beta$) parameters. These parameters modulate the kinematic token embeddings, allowing the model to distinguish between joints that may have similar connectivity but different functional roles (e.g., a wrist joint vs. a base joint).

3. Key Contributions

1. Novel Architecture: A unified framework that integrates kinematic tokens, topology-aware attention, and joint-attribute conditioning into a standard VLA backbone ($\pi_0.5$).
2. Token Design: The introduction of Kinematic Tokens solves the interface problem, enabling morphology injection into compressed VLA action spaces.
3. Hybrid Attention Strategy: The proposal of Mix-Mask attention effectively balances the need for local kinematic constraints with global coordination capabilities.
4. Semantic Enrichment: Demonstrating that conditioning on joint attributes (beyond just connectivity) significantly improves policy robustness.

4. Experimental Results

The method was evaluated on three robot platforms: DROID (Franka Panda), Unitree G1 Dex1, and SO101.

• Single-Embodiment Performance (DROID):

  • The baseline $\pi_0.5$ achieved a 19.7% success rate.
  • Adding Kinematic Tokens alone raised this to 36.0%.
  • The full model (KT + Mix-Mask + FiLM) achieved 47.4% success rate.
  • Key Insight: The combination of all three components yielded a 5x improvement on specific difficult tasks compared to the baseline.

• Cross-Robot Generalization (Unitree G1 Dex1):

  • The method maintained effectiveness on a different robot with a 16-DoF action space, achieving 28.0% success rate (vs. 24.7% baseline).

• Multi-Embodiment Training (Panda + SO101):

  • When training a single policy on a mixture of two different robots (8-DoF vs. 6-DoF), the embodiment-aware model consistently outperformed the baseline throughout training.
  • At 50k steps, the proposed method achieved 15.5% Macro Success Rate vs. 5.0% for the baseline.

• Ablation Studies:

  • Temporal Chunking: A single chunk ($G=1$) performed best, suggesting that compressing time too aggressively hurts performance.
  • Auxiliary Tokens: Adding AKT provided consistent gains, especially when combined with Mix-Mask.
  • Soft-Mask vs. Hard-Mask: Soft-Mask variants (even with optimized initialization) consistently underperformed Hard-Mask variants (Mix-Mask), likely due to optimization instability.

5. Significance and Impact

• Robustness: The work demonstrates that explicitly encoding physical structure (morphology) into the policy architecture is superior to relying on implicit learning from observations. This leads to policies that are more robust to hardware variations and failures.
• Scalability: By solving the "embodiment-agnostic" bottleneck, this approach paves the way for true generalist robot policies that can be deployed across diverse hardware without extensive retraining or fine-tuning for each new robot.
• Foundation for Future Research: The paper establishes a clear path for integrating structural priors into large-scale foundation models, moving the field closer to human-like adaptability where a single "brain" can control various "bodies."