Embedding Morphology into Transformers for Cross-Robot Policy Learning

Kei Suzuki; Jing Liu; Ye Wang; Chiori Hori; Matthew Brand; Diego Romeres; Toshiaki Koike-Akino

cs.ROFeb 26, 2026

Embedding Morphology into Transformers for Cross-Robot Policy Learning

Authors: Kei Suzuki, Jing Liu, Ye Wang, Chiori Hori, Matthew Brand, Diego Romeres, Toshiaki Koike-Akino

Organizations: Mitsubishi Electric Research Laboratories (MERL), Cambridge, Massachusetts, USA.

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.

Figures & tables

Figure 1: Embedding robot morphology into a Transformer-based VLA policy : We embed kinematic topology and per-joint semantics into the action policy. This design consistently improves cross-robot policy learning.

Panel of Figure 2 — Figure 2: Embodiment-aware Transformer Policy : Morphology is embedded via three mechanisms: (1) kinematic tokens; (2) topology-aware attention bias; and (3) joint-attribute conditioning.

Feature	Description	Example
type_pris	1 if prismatic joint, else 0	0
type_rev	1 if revolute joint, else 0	1
ax	Joint axis X component (unit vector)	0
ay	Joint axis Y component (unit vector)	0
az	Joint axis Z component (unit vector)	1
hard_lower	Hard lower limit [rad or m]	-2.9671

Table 1: Per-joint descriptors for joint-attribute conditioning : Each joint

j

is represented by a descriptor

s_{j}

, which is used for FiLM-based conditioning of kinematic-token embeddings. The descriptor includes joint type indicators, axis direction, motion limits, and contact-related properties (log-transformed where indicated); the Example column shows a representative instance from the DROID (Franka Panda arm) embodiment.

Panel of Figure 3 — Figure 3: Simulation environments for evaluation : All environments are evaluated language-conditioned pick-and-place manipulation.

Setting	Fine-tuning	Batch	Horizon	Steps
DROID (Panda)	AP-FT ( $\sim$ 450M)	32	16	100k
Unitree G1 Dex1	Full-FT ( $\sim$ 3.5B)	8	32	60k
DROID+SO101	AP-FT ( $\sim$ 450M)	32	16	125k

Table 2: Training protocol : Fine-tuning configurations and key hyperparameters for each benchmark. All runs start from pi05-base with a cosine learning-rate schedule. AP-FT updates only the action-policy components, whereas Full-FT updates the entire model.

Kinematic Token	Mask	FiLM	Success Rate (SR%) $\uparrow$ (95% CI)
(Chunk $G$ =1)			Avg	Task 1	Task 2	Task 3
–	–	–	19.7 $\pm$ 4.5	18.3 $\pm$ 4.4	15.7 $\pm$ 4.1	25.0 $\pm$ 4.9
[0.6pt/2pt] ✓	–	–	36.0 $\pm$ 5.4	10.3 $\pm$ 3.6	67.7 $\pm$ 5.3	30.0 $\pm$ 5.2
✓	Full-Mask	–	30.1 $\pm$ 5.2	15.3 $\pm$ 4.1	34.0 $\pm$ 5.3	41.0 $\pm$ 5.5
✓	Mix-Mask	–	36.9 $\pm$ 5.4	27.7 $\pm$ 5.0	56.7 $\pm$ 5.6	26.3 $\pm$ 5.1
✓	Soft-Mask	–	26.1 $\pm$ 4.9	29.7 $\pm$ 5.1	17.3 $\pm$ 4.3	31.3 $\pm$ 5.2

Table 3: Single-embodiment results on DROID : We evaluate whether embedding robot morphology improves performance even under single-embodiment training on DROID (Franka Panda). Starting from the

\pi_{0.5}

baseline, we add kinematic tokens, topology-aware attention, and joint-attribute conditioning (FiLM), individually and in combination. Each component contributes to higher success rates, and the best performance is achieved by combining all three (KT + Mix-Mask + FiLM). Bold : best SR in each column (ties allowed). Underline : second-best. 95% CI is reported with

\pm\Delta

Kinematic Token	Mask	FiLM	Avg SR% $\uparrow$
(Chunk $G$ =1)			(95% CI)
–	–	–	24.7 $\pm$ 4.9
[0.6pt/2pt] ✓	–	–	24.3 $\pm$ 4.9
✓	Full-Mask	–	23.8 $\pm$ 4.8
✓	Mix-Mask	–	27.3 $\pm$ 5.0
✓	–	✓	25.3 $\pm$ 4.9

Table 4: Single-embodiment results on Unitree G1 Dex1 : We evaluate our method on Unitree G1 Dex1 to test its benefit beyond DROID. A key difference is the 16-DoF joint-space action space. Overall, our components remain effective in this setting, and combining KT with Mix-Mask and FiLM achieves the best SR. Bold : best SR. Underline : second-best.

Panel of Figure 4 — Figure 4: Multi-embodiment learning curves on Panda–SO101 : We jointly train a single policy on a mixed Panda (DROID) and SO101 dataset. We report Macro SR, defined as $(\mathrm{SR}_{\text{Panda}}+\mathrm{SR}_{\text{SO101}})/2$ . Our full model outperforms the $\pi_{0.5}$ baseline throughout training. For completeness, per-embodiment success rates are reported in the appendix F . Shaded regions indicate 95% confidence intervals.

Kinematic Tokens	Avg SR% $\uparrow$
Chunk ( $G$ )	(95% CI)
1	36.0 $\pm$ 5.4
2	35.8 $\pm$ 5.4
4	34.4 $\pm$ 5.3
8	30.5 $\pm$ 5.2
16	33.3 $\pm$ 5.3

Table 5: Effect of temporal chunk size on DROID : We study how the temporal granularity of kinematic tokens affects performance. Using a single chunk (

G{=}1

) achieves the best performance, while larger chunk sizes reduce performance. Bold : best SR. Underline : second-best.

Kinematic Tokens		Avg SR% $\uparrow$
Mask	AKT	(95% CI)
–	–	36.0 $\pm$ 5.4
–	✓	37.0 $\pm$ 5.4
Full-Mask	–	30.3 $\pm$ 5.2
Full-Mask	✓	33.0 $\pm$ 5.3
Mix-Mask	–	37.0 $\pm$ 5.4

Table 6: Effect of auxiliary kinematic tokens (AKT) on DROID : We study whether increasing per-joint token capacity improves performance. Adding AKT consistently improves performance, especially with Mix-Mask. Bold : best SR. Underline : second-best.

Bias Init	Avg SR% $\uparrow$ (95% CI)
Zero	26.1 $\pm$ 4.9
Hard	25.1 $\pm$ 4.9
Mix	28.1 $\pm$ 5.1
Linear	20.4 $\pm$ 4.5

Table 7: Effect of bias initialization on DROID : We study the effect of bias initialization in Soft-Mask attention (Zero, Hard, Mix, Linear). Overall, Mix achieves the best performance among Soft-Mask variants, yet all Soft-Mask variants remain below the masked variants in Table 3 . Bold : best. Underline : second-best.

Appendix figures & tables8 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 5: Attention mask : Attention mask used in our VLA action policy with kinematic tokens. Dark cells indicate unmasked attention and light cells indicate masked attention. Tokens are grouped by type for visualization (image/prompt/action/kinematic; each group may contain multiple tokens). We append kinematic tokens and apply a topology-aware mask only in the joint-to-joint block to encode kinematic connectivity, while all other attention patterns follow the base $\pi_{0.5}$ mask.

Figure 6: Verb distribution in the DROID 1/8 subset : Verb frequencies in task instructions.

Figure 7: Object distribution in the DROID 1/8 subset : Object frequencies in task instructions.

Status	Part	Parameter name
Frozen: VLM part (2,923,335,408 parameters)
Frozen	VLM	PaliGemma/img/Transformer/encoder_norm/bias
Frozen	VLM	PaliGemma/img/Transformer/encoder_norm/scale
Frozen	VLM	PaliGemma/img/Transformer/encoderblock/LayerNorm_0/bias
Frozen	VLM	PaliGemma/img/Transformer/encoderblock/LayerNorm_0/scale
Frozen	VLM	PaliGemma/img/Transformer/encoderblock/LayerNorm_1/bias

Appendix

Table 8: Trainable vs. frozen parameters in AP-FT : VLM parameters are frozen, while the action-policy parameters are optimized.

Training	Kinematic Token		Encoder	Mask	Batch	Train Data	Avg		Task 1		Task 2		Task 3
	Chunk	AKT					SR(%) $\uparrow$	Time[s] $\downarrow$	SR(%) $\uparrow$	Time[s] $\downarrow$	SR(%) $\uparrow$	Time[s] $\downarrow$	SR(%) $\uparrow$	Time[s] $\downarrow$
Full-FT	–	–	–	–	256	Full Dataset	66.3	3.215	65.7	2.526	67.7	4.178	65.7	2.939
AP-FT	–	–	–	–	32	1/8 subset	19.7	4.849	18.3	4.731	15.7	6.616	25.0	3.201
AP-FT	1	–	linear	–	32	1/8 subset	26.7	4.606	14.7	5.746	34.3	4.984	31.0	3.088
AP-FT	1	✓	linear	–	32	1/8 subset	35.2	4.571	15.3	4.824	57.0	5.175	33.3	3.714
AP-FT	1	–	lin_swiGlu_lin	–	32	1/8 subset	36.0	4.925	10.3	5.450	67.7	5.254	30.0	4.072

Appendix

Table 9: DROID simulation evaluation summary : Best and second-best are indicated by bold and underline , respectively.

Training	Kinematic Token		Encoder	Mask	FiLM	Batch	Train Data	Avg		Task 1		Task 2		Task 3
	Chunk	AKT						SR(%) $\uparrow$	Time[s] $\downarrow$	SR(%) $\uparrow$	Time[s] $\downarrow$	SR(%) $\uparrow$	Time[s] $\downarrow$	SR(%) $\uparrow$	Time[s] $\downarrow$
AP-FT	–	–	–	–	–	32	1/8 subset	19.7	4.849	18.3	4.731	15.7	6.616	25.0	3.201
AP-FT	–	–	–	–	–	32	1/8 subset	19.7	4.849	18.3	4.731	15.7	6.616	25.0	3.201
AP-FT	1	–	lin_swiGlu_lin	–	–	32	1/8 subset	36.7	4.747	11.3	5.802	67.3	4.597	31.3	3.843
AP-FT	1	–	lin_swiGlu_lin	–	✓	32	1/8 subset	37.7	3.794	6.0	3.727	65.0	4.488	42.0	3.168
AP-FT	1	✓	lin_swiGlu_lin	–	–	32	1/8 subset	37.0	4.361	11.0	4.618	60.3	5.110	39.7	3.355

Appendix

Table 10: DROID simulation evaluation summary : Best and second-best are indicated by bold and underline , respectively.

Base Model	Variant	Init.	Success Rate (SR%) (95% CI) $\uparrow$
			Avg	Task1	Task2	Task3
pi05-base	Soft-Mask (v1.0)	Zero	25.9 $\pm$ 4.9	26.7 $\pm$ 5.0	20.3 $\pm$ 4.5	30.7 $\pm$ 5.2
pi05-base	Soft-Mask (v1.0)	Hard	25.4 $\pm$ 4.9	17.7 $\pm$ 4.3	15.7 $\pm$ 4.1	42.7 $\pm$ 5.6
pi05-base	Soft-Mask (v1.1)	Zero	26.4 $\pm$ 5.0	21.3 $\pm$ 4.6	16.7 $\pm$ 4.2	41.3 $\pm$ 5.5
pi05-base	Soft-Mask (v1.1)	Hard	29.5 $\pm$ 5.1	16.7 $\pm$ 4.2	28.7 $\pm$ 5.1	43.0 $\pm$ 5.6
pi05-base	Soft-Mask (v2.0)	Zero	34.4 $\pm$ 5.3	27.6 $\pm$ 5.0	37.0 $\pm$ 5.4	38.6 $\pm$ 5.5

Appendix

Table 11: Soft-Mask results (SR%) : We report per-task success rates (Task1–Task3) in % along with macro averages, shown as SR%

\pm

95% CI.

Panel of Figure 8 — Figure 8: Multi-embodiment learning curves on Panda–SO101 : Shaded regions indicate 95% CI (per-checkpoint).

Robot Policy Learning

Robot policy learning develops controllers that map sensory observations, task goals, and interaction history to actions for embodied robots.

See all 357 papers

PolyMorph: Learning Locomotion Policy for Cross-Embodiment Adaptation
Oct 14, 2026Hyunsik Oh, Seunghun Jeon, Donghoon Youm +5Robot Policy Learning Cross-Embodiment Robot Learning
Latent Policy Steering: An Efficient and Flexible Framework for Cross-Embodiment Transfer
Sep 18, 2026Yiqi Wang, Mrinal Verghese, Jeff SchneiderCross-Embodiment Transfer Visuomotor Policy Learning
Robotics Institute, School of Computer Science, Carnegie Mellon University, 5000 Forbes Ave, Pittsburgh, United States
The Robot Is Not Its Description: GaugeBench for Representation Robustness in Morphology-Aware Policies
Oct 6, 2026Rahath Malladi, Arshia Sangwan, Rajesh K. Gupta +1Robotic Policy Evaluation Cross-Embodiment Robot Learning
University of California San Diego · New York University

Oct 14, 2026J/K move · Enter open · S save

PolyMorph: Learning Locomotion Policy for Cross-Embodiment Adaptation

Hyunsik Oh, Seunghun Jeon, Donghoon Youm, Jungwoo Hur+4

Robot Policy Learning Cross-Embodiment Robot Learning Robotic Control Robot Policy Generalization RL for Legged Locomotion

Open paper AskPDF arXiv

Reinforcement learning has enabled robust locomotion for legged robots, but learned policies are typically tied to a specific robot morphology and require additional engineering when the number of legs, actuators, or mechanical properties change. We present PolyMorph, a morphology-conditioned framework that learns a single locomotion policy for heterogeneous legged robots. PolyMorph converts each robot's URDF into joint-link modules and combines a recurrent encoder, graph-based Transformer, and per-joint decoder to generate actions across variable leg counts and joint configurations. To improve training stability, PolyMorph uses morphology-normalized rewards, adaptive robot sampling, and URDF-derived PD-gain initialization. We evaluate PolyMorph in simulation across bipedal, tripedal, and quadrupedal robots, including unseen morphology variants and platforms, and deploy the policy on real quadruped platforms operated in multiple configurations. The results show that a single PolyMorph policy generalizes across diverse robot morphologies without per-robot policy retraining.

Abstract

Robot Policy Learning

Explore similar work

PolyMorph: Learning Locomotion Policy for Cross-Embodiment Adaptation

Latent Policy Steering: An Efficient and Flexible Framework for Cross-Embodiment Transfer

The Robot Is Not Its Description: GaugeBench for Representation Robustness in Morphology-Aware Policies