Paper deep dive
Embedding Morphology into Transformers for Cross-Robot Policy Learning
Kei Suzuki, Jing Liu, Ye Wang, Chiori Hori, Matthew Brand, Diego Romeres, Toshiaki Koike-Akino
Intelligence
Status: succeeded | Model: Gemma-4-26B-A4B | Prompt: intel-v1 | Confidence: 92%
Last extracted: 7/20/2026, 8:31:40 AM
Summary
This paper proposes an embodiment-aware transformer policy for cross-robot policy learning by injecting morphology into Vision-Language-Action (VLA) models. The method introduces three mechanisms: kinematic tokens for joint-wise action representation, topology-aware attention bias for kinematic message passing, and joint-attribute conditioning for per-joint semantics. Experiments on DROID (Franka Panda), Unitree G1 Dex1, and SO101 demonstrate that this structured integration improves success rates and robustness compared to the vanilla pi0.5 VLA baseline.
Entities (13)
Relation Signals (10)
DROID → contains → Franka Panda
confidence 95% · DROID dataset, which contains demonstrations collected on an 8-DoF ... Franka Panda arm.
Embodiment-aware transformer policy → uses → Topology-aware attention bias
confidence 95% · ...(2) a topology-aware attention bias that encodes kinematic topology as an inductive bias in self-attention...
Embodiment-aware transformer policy → uses → Joint-attribute conditioning
confidence 95% · ...(3) joint-attribute conditioning that augments topology with per-joint descriptors...
Embodiment-aware transformer policy → uses → Kinematic Tokens
confidence 95% · We propose an embodiment-aware transformer policy that injects morphology via three mechanisms: (1) kinematic tokens...
Embodiment-aware transformer policy → outperforms → $\\pi_{0.5}$
confidence 92% · this structured integration consistently improves performance over a vanilla pi0.5 VLA baseline
Unitree G1 Dex1 → isevaluatedon → Unitree G1
confidence 90% · Unitree G1 Dex1 (Unitree G1): We also evaluate on the Unitree Dex1 benchmark in simulation
Joint-attribute conditioning → uses → FiLM
confidence 90% · We use Feature-wise Linear Modulation (FiLM) ... to map s j to feature-wise scale and shift parameters
Kinematic Tokens → →
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
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.
Tags
Links
- Source: https://arxiv.org/abs/2603.00182v1
- Canonical: https://arxiv.org/abs/2603.00182v1
Trouble viewing inline? Open PDF directly →
Full Text
53,314 characters extracted from source content.
Expand or collapse full text
Embedding Morphology into Transformers for Cross-Robot Policy Learning Kei Suzuki 1 Jing Liu 1 Ye Wang 1 Chiori Hori 1 Matthew Brand 1 Diego Romeres 1 Toshiaki Koike-Akino 1 Abstract Cross-robot policy learning—training a sin- gle 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 typ- ically embodiment-agnostic and must infer kine- matic structure purely from observations, which can reduce robustness across embodiments and even limit performance within a single embodi- ment. We propose an embodiment-aware trans- former 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 topol- ogy with per-joint descriptors to capture seman- tics beyond connectivity. Across a range of em- bodiments, this structured integration consistently improves performance over a vanilla pi0.5 VLA baseline, indicating improved robustness both within an embodiment and across embodiments. 1. Introduction Transformer-based robot policies, especially vision- language-action (VLA) models (Ahn et al., 2022; Brohan et al., 2022; Driess et al., 2023), have advanced rapidly by scaling to large and diverse datasets, yielding increasingly generalist controllers that can execute a broad range of tasks. However, cross-robot policy learning—training a single policy to perform well across diverse robot embodiments— remains a central challenge. This includes robustness to embodiment changes in the real world, such as hardware variation or failures, design upgrades, or deployment on en- tirely different robot platforms. In practice, achieving strong 1 Mitsubishi Electric Research Laboratories (MERL), Cam- bridge, Massachusetts, USA. Correspondence to: Kei Suzuki<ke- suzuki@merl.com>. Preprint. March 3, 2026. 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. performance across embodiments often requires additional data and training for each robot, e.g., via fine-tuning (Intel- ligence et al., 2025; Kim et al., 2024) and/or replacing the action head (Team et al., 2024; Doshi et al., 2024) to match the target action space. Such approaches reflect a structural limitation: many VLA policies are typically embodiment- agnostic and must learn kinematics and cross-joint coor- dination implicitly from observations, making cross-robot policy learning particularly challenging. This motivates in- jecting embodiment as an explicit inductive bias into the policy architecture. A common formulation represents robot morphology as a kinematic graph, where nodes correspond to actuated joints and edges capture physical connectivity. Prior work incorporates this structure through graph neural networks (Wang et al., 2018; Huang et al., 2020; Whitman et al., 2023) or topology-aware attention in transformers (Veli ˇ ckovi ́ c et al., 2017; Hong et al., 2021; Sferrazza et al., 2024), improving cross-robot policy learning. Nevertheless, existing embedding approaches face three challenges. (i) Lack of a kinematic token interface: VLA models such as π 0.5 (Intelligence et al., 2025) compress joint-space struc- ture into a compact set of action tokens, making it unclear where to apply existing morphology-embedding methods. (i) Local–global trade-off in topology-aware attention: 1 arXiv:2603.00182v1 [cs.RO] 26 Feb 2026 Embedding Morphology into Transformers for Cross-Robot Policy Learning Enforcing strong locality promotes kinematic message pass- ing but can limit long-range coordination. (i) Missing joint semantics: Existing methods do not capture per-joint semantics, even though joints with identical topology can play different functional roles (e.g., actuation type and lim- its). Contributions To address these limitations, we propose an embodiment-aware transformer policy that injects robot morphology into the VLA action policy via three mecha- nisms: (1) kinematic tokens to provide a joint-wise action representation; (2) topology-aware attention bias with a local/global schedule to balance kinematic message pass- ing and global context; and (3) joint-attribute condition- ing to incorporate per-joint semantics beyond connectiv- ity. Across single- and multi-embodiment evaluations on DROID (Franka Panda), Unitree G1 Dex1, and SO101, our structured morphology encoding improves success rates over the vanillaπ 0.5 VLA baseline (Intelligence et al., 2025), indicating improved robustness both within an embodiment and across embodiments. 2. Related Work 2.1. Data-driven generalist robot policies Recent progress in robot policy learning has been driven by scaling vision-language-action (VLA) models—large transformer policies that couple a pretrained VLM back- bone with an action policy head (Ahn et al., 2022; Brohan et al., 2022; Driess et al., 2023). A common recipe is to pretrain on large-scale robot datasets (Walke et al., 2023; O’Neill et al., 2024; Khazatsky et al., 2024), and then adapt to a target robot via fine-tuning or lightweight adaptation modules (Team et al., 2024; Doshi et al., 2024). Despite these advances, cross-robot policy learning—training a sin- gle policy to perform well across multiple embodiments— remains a central challenge, as these policies are typically embodiment-agnostic and must infer kinematic structure from observations alone. 2.2.Architectures for embodiment-aware robot policies A common approach to improve cross-robot policy learning is to explicitly encode robot morphology as a kinematic graph, where nodes correspond to actuated joints and edges represent physical connectivity. This explicit representa- tion provides a structured interface for multi-embodiment learning. GNN policies with kinematic-graph message passing Graph neural network (GNN) policies encode morphology by running message passing on a kinematic graph, promot- ing body-wide coordination (Wang et al., 2018; Huang et al., 2020; Whitman et al., 2023). A key design question in this line of work is how information should propagate over the kinematic graph—balancing local interactions with global coordination. NerveNet (Wang et al., 2018) applies repeated local (one-hop) message passing so that information can gradually propagate across the body. SMP (Huang et al., 2020) further explores structured propagation by separating bottom-up and top-down information flow along the kine- matic tree, aiming to capture hierarchical coordination pat- terns. Despite these advances, designing message-passing schemes that simultaneously support local and global infor- mation propagation remains challenging. Transformer policies with topology-aware attention Transformer-based policies facilitate both local and global information exchange by self-attention, and inject kinematic priors by incorporating the kinematic graph directly into the attention mechanism (Veli ˇ ckovi ́ c et al., 2017; Hong et al., 2021; Sferrazza et al., 2024). Concretely, these methods use the kinematic graph to modulate attention, either by Hard-Mask attention or Soft-Mask attention. Hard-masking approaches (Sferrazza et al., 2024; Buterez et al., 2024) inject topology by enforcing a binary attention constraint: each kinematic token can attend only to itself and its 1-hop kinematic neighbors in the kinematic graph, and all other joint pairs are disallowed (implemented by masking their attention logits before the softmax). To recover global co- ordination, some mixed designs alternate topology-masked (local) layers with fully connected (global) attention lay- ers. Soft-Mask attention approaches (Hong et al., 2021; Luo et al., 2025) inject topology by keeping full attention but adding a learnable, topology-conditioned bias term to the attention logits, so kinematically closer joints receive higher attention on average without strictly blocking any pairs. Although Soft-Mask approaches offer greater flexibil- ity for embedding kinematics, prior works (Sferrazza et al., 2024; Buterez et al., 2024) report optimization instability, whereas Hard-Mask designs tend to be more stable and can outperform Soft-Mask variants. In our method, we follow this taxonomy and consider both families of topology-aware attention within a unified framework (Full-/Mix-Mask for Hard-Mask; SPD-based bias for Soft-Mask). Despite these advances, two limitations remain: •Limited applicability to VLA token interfaces: In state-of-the-art VLA models such asπ 0.5 (Intelligence et al., 2025), action generation often compresses joint- space structure into a small set of action tokens, making it difficult to apply existing kinematic biasing schemes at the level of individual joints. • Topology alone lacks joint semantics: Connectivity specifies how information propagates, but not what each joint represents. Encoding per-joint semantics is important for disambiguating functional roles and 2 Embedding Morphology into Transformers for Cross-Robot Policy Learning (a) Kinematic Token (KT): Kinematic tokens provide a joint-wise interface for the VLA action policy. While the standard action tokens retain temporal structure, kinematic tokens compress the horizon into per-joint summaries, emphasizing cross-joint (spatial) structure and enabling topology/semantics embedding. (b) Topology bias and joint-attribute conditioning: We embed kinematic topology and semantics in two ways: (i) a topology bias encourages kinematic message passing by restricting joint-to-joint self-attention to connected joints, and (i) FiLM conditions kinematic- token embeddings on per-joint descriptors to disambiguate joint roles beyond connectivity. 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. enabling more structured action generation (e.g., actu- ation type, axis, or joint limits). 3. Proposed Method We propose an embodiment-aware transformer policy that injects robot morphology into the VLA action policy via three mechanisms: (1) kinematic tokens for factorized joint- action representation; (2) a topology-aware attention bias for kinematic message passing; and (3) joint-attribute condi- tioning for semantics beyond connectivity. An overview is depicted in Figure 2. 3.1. Kinematic Tokens (KT) To factorize the action sequence across joints while com- pressing temporal information, we introduce kinematic to- kens in addition to the standard action tokens used in VLA policies. The action tokens retain fine-grained temporal structure, while the kinematic tokens provide a compact, per- joint view that highlights cross-joint structure (Figure 2a). Leta t,j denote the scalar action for jointj ∈0,...,J− 1 at timet ∈ 0,...,H − 1, whereJandHare the max- imum number of joints and the maximum horizon length, respectively. The originalπ 0.5 VLA uses action tokens cou- pling all joints into one embedding per horizon,[a t,j ] J−1 j=0 , that prevents the use of topology-aware attention. To embed morphology, we decouple the action into the spatial domain. We then split the horizon intoGnon-overlapping temporal chunks of sizeg = H G . For each jointj, we concatenate the g horizon actions within the chunk into a vector b j,T k := [a t,j ] t∈T k ∈R g ,(1) where for each chunkk ∈ K := 0, 1,...,G− 1, we define the index set T k =kg, kg + 1, ..., (k + 1)g− 1.(2) We refer to the vectorb j,T k as the kinematic token for joint jand chunkk. This yieldsGkinematic tokens per joint (andJGkinematic tokens in total forJjoints). We project each kinematic tokenb j,T k to ad-dimensional embedding z j,T k = Enc 0 (b j,T k ) ∈R d with a lightweight multi-layer perceptron (MLP), and append them to the token sequence as additional context for the action expert of VLA. Follow- ing the VLA backbone, the policy predicts actions from the action tokens, which can attend to these kinematic tokens to leverage joint-wise structure. In our experiments, we find that using a single chunk (G=1) performs best. Auxiliary kinematic tokens (AKT)Building on the kine- matic token interface, we further increase the token capacity per joint by introducing auxiliary kinematic tokens. Specif- ically, for each kinematic tokenb j,T k , in addition to the standard embeddingz j,T k = Enc 0 (b j,T k )∈R d , we gener- ate M auxiliary embeddings z (m) j,T k = Enc m (b j,T k )∈R d , m = 1,...,M,(3) 3 Embedding Morphology into Transformers for Cross-Robot Policy Learning using lightweight encoders with the same input but indepen- dent parameters. We then append these auxiliary tokens to the kinematic token sequence, so the transformer can attend to a richer set of per-joint representations. 3.2. Topology-aware attention Vanilla self-attention treats the token sequence as a fully connected graph, allowing arbitrary interactions between joints. In contrast, robot embodiments admit a natural kine- matic topology represented as a graph. LetG = (V,E )be the kinematic graph, where each vertexv j ∈Vcorresponds to jointj, and each edge(v i ,v j ) ∈ Eindicates physical connectivity. We encode this topology as an inductive bias in self-attention (Figure 2b). Joint-to-joint self-attention modulation In the VLA to- ken sequence, we inject kinematic structure only within the joint-to-joint attention block, while all other attention patterns remain identical to the originalπ 0.5 mask (Ap- pendix 5). Letlogits (ℓ) i,j denote the scaled dot-product atten- tion logits at layerℓbetween query jointiand key jointj. We modulate joint-to-joint attention by adding a topology- dependent term B (ℓ) i,j : α (ℓ) i,j = softmax j logits (ℓ) i,j + B (ℓ) i,j ,(4) where α (ℓ) i,j is the attention weight. We consider three topology-bias variants in two families: Hard-Mask (Full-Mask and Mix-Mask) and Soft-Mask. Hard-Mask uses an adjacency-based hard bias that blocks attention to non-neighbor joints. Soft-Mask uses a shortest- path distance (SPD)-based soft bias that favors nearby joints while retaining full attention. Hard masking family: We first define the 1-hop neigh- borhood indicator M i,j = ( 1, (i = j) or (i,j)∈E, 0, otherwise. (5) Hard masking is implemented by setting B (ℓ) i,j = ( 0, M i,j = 1, −∞, M i,j = 0, (6) so that non-neighbor joints receive zero attention after the softmax. Full-Mask: Full-Mask applies hard masking Eq.(6)at every layer, enforcing strictly local (1-hop) joint-to-joint interactions. Table 1. Per-joint descriptors for joint-attribute conditioning: Each jointjis represented by a descriptors 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 indi- cated); the Example column shows a representative instance from the DROID (Franka Panda arm) embodiment. FeatureDescriptionExample typepris1 if prismatic joint, else 00 typerev1 if revolute joint, else 01 axJoint axis X component (unit vector) 0 ayJoint axis Y component (unit vector) 0 azJoint axis Z component (unit vector) 1 hardlowerHard lower limit [rad or m] -2.9671 hardupperHard upper limit [rad or m] 2.9671 dampingloglog(contact damping)6.90776 friction anchorFriction anchor flag (0/1)1 lateralfrictionLateral friction coeffi- cient 1 spinningfrictionSpinning friction coeffi- cient 0.1 stiffnessloglog(contact stiffness)10.30895 Mix-Mask: Mix-Mask alternates masked and unmasked layers to balance local message passing with periodic global coordination. Even-numbered layers apply hard masking (Eq.(6)), while odd-numbered layers use full attention (B (ℓ) i,j = 0 for all i,j). Soft-Mask: We define the shortest-path distance onG as d(i,j) = min p:i→j |p|,(7) where the minimum is taken over all pathspfrom jointi to jointjand|p|is the number of edges in the path (with d(i,i) = 0). Following a Graphormer-style design (Ying et al., 2021), we parameterize the topology term by a learn- able bias table indexed by distance: B (ℓ) i,j = θ (ℓ) d(i,j) .(8) Unlike hard masking, Soft-Mask biases attention by kine- matic distance while retaining full attention paths. 3.3. Joint-attribute conditioning Topology-aware attention specifies which joints can ex- change information based on kinematic connectivity, but connectivity alone does not capture the semantics of each joint (e.g., joints with similar local topology may play dif- 4 Embedding Morphology into Transformers for Cross-Robot Policy Learning (a) DROID (Franka Panda): Three pick- and-place tasks: (1) “put the cube in the bowl”; (2) “put the can in the mug”; (3) “put banana in the bin” (b) Unitree G1 Dex1 (Unitree G1): “Pick up the red block and place it inside the yel- low box.” (c) SO101: “Grab orange and place into plate” Figure 3. Simulation environments for evaluation: All environments are evaluated language-conditioned pick-and-place manipulation. ferent functional roles). To complement topology, we condi- tion kinematic-token embeddings on per-joint descriptors derived from robot morphology (Figure 2b). For each jointj, we define a per-joint descriptors j using the features in Table 1, capturing the joint type, axis direc- tion, motion limits, and contact-related properties. We use Feature-wise Linear Modulation (FiLM) (Perez et al., 2018) to map s j to feature-wise scale and shift parameters: γ j ,β j = FiLM(s j ),(9) whereγ j ,β j ∈R d . Given the kinematic-token embedding before conditioningz j ∈R d , we apply feature-wise affine modulation: ̃z j = (1 + γ j )⊙ z j + β j ,(10) where⊙denotes element-wise multiplication. The con- ditioned embedding ̃z j augments topology-based message passing with joint-specific semantics beyond connectivity. 4. Experiments We evaluate morphology-aware transformer policies for cross-robot policy learning, where a single policy is trained and evaluated on the same set of robot embodiments. Our experiments isolate the contribution of each component: (i) whether kinematic tokens improve success rates over the vanillaπ 0.5 baseline; (i) whether topology-aware bias im- proves performance over unmasked attention; (i) whether joint-attribute conditioning provides additional gains over topology-only variants; and (iv) whether these effects persist under single- and multi-embodiment training. 4.1. Single-embodiment evaluations We first study single-embodiment imitation learning on two benchmarks to test whether explicit morphology encoding improves performance even outside cross-embodiment set- tings. DROID (Franka Panda): We use the public DROID Table 2. Training protocol: Fine-tuning configurations and key hyperparameters for each benchmark. All runs start from pi05-basewith a cosine learning-rate schedule. AP-FT up- dates only the action-policy components, whereas Full-FT updates the entire model. SettingFine-tuningBatch Horizon Steps DROID (Panda)AP-FT (∼450M)3216100k Unitree G1 Dex1 Full-FT (∼3.5B)83260k DROID+SO101AP-FT (∼450M)3216125k dataset (Khazatsky et al., 2024), which contains demonstra- tions collected on an 8-DoF (including the gripper) Franka Panda arm. We train on a 1/8 subset of DROID and evaluate in a simulation evaluation suite consisting of three language- conditioned pick-and-place tasks 1 (Figure 3a). Unitree G1 Dex1 (Unitree G1): We also evaluate on the Unitree Dex1 benchmark in simulation, which focuses on manipulation with the Unitree G1 platform while the lower body remains stationary. We train on the public Unitree Dex1 dataset with 16-DoF joint-space actions 2 and evaluate in an IsaacLab- based simulator 3 (Figure 3b). 4.2. Multi-embodiment evaluation We evaluate multi-embodiment imitation learning by jointly training a single policy on a mixture of Panda and SO101 demonstrations, where the two robots differ in joint-space action dimensionality (Panda: 8-DoF; SO101: 6-DoF). This setting is challenging because the policy must reconcile embodiment-specific kinematics and coordination patterns under mismatched action spaces, and naive joint training 1 https://github.com/arhanjain/sim-evals 2 https://huggingface.co/datasets/unitreer obotics/G1_Dex1_PickPlaceRedBlock_Dataset_ Sim 3 https://github.com/unitreerobotics/unitr e_sim_isaaclab 5 Embedding Morphology into Transformers for Cross-Robot Policy Learning 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π 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±∆. Kinematic TokenMaskFiLMSuccess Rate (SR%)↑ (95% CI) (Chunk G=1)AvgTask 1Task 2Task 3 –19.7± 4.518.3± 4.415.7± 4.125.0± 4.9 ✓–36.0± 5.410.3± 3.667.7± 5.330.0± 5.2 ✓Full-Mask–30.1± 5.215.3± 4.134.0± 5.341.0± 5.5 ✓Mix-Mask–36.9± 5.427.7± 5.056.7± 5.626.3± 5.1 ✓Soft-Mask– 26.1± 4.9 29.7± 5.117.3± 4.331.3± 5.2 ✓–✓37.7± 5.56.0± 2.765.0± 5.442.0± 5.6 ✓Mix-Mask✓47.4± 5.65.7± 2.7 77.7± 4.7 58.7± 5.5 can lead to interference between embodiments. For the Panda portion, we use the same 1/8 subset of DROID and evaluation suite as in the single-embodiment setting. For SO101, we use the LeRobot dataset 4 and evaluate in the cor- responding LeIsaac simulator 5 (Figure 3c). During training, each mini-batch contains trajectories from a single embod- iment, and we sample Panda and SO101 batches with an 8:2 ratio (i.e., 80% Panda and 20% SO101). We choose this ratio to keep the same amount of DROID (Panda) data as in our single-embodiment setting, while SO101 can be trained adequately with a smaller dataset in this regime. 4.3. Metrics We evaluate policies using task success rate (SR), computed over 300 rollout trials per experimental condition. Success is determined by an axis-aligned bounding box (AABB) criterion: an episode is successful if the object remains inside the target region for a minimum dwell time and is approximately stationary. To quantify statistical uncertainty, we also report a 95% confidence interval (CI) for each SR using the Wilson score interval (Brown et al., 2001) for a binomial proportion (n=300). In tables, we present results in the formatSR ± ∆, where∆denotes the CI half-width, i.e.,∆ = (CI upper − CI lower )/2. Across trials, we vary the random seed to change initial conditions (e.g., object placement). 4.4. Model variants We evaluate variants built on theπ 0.5 backbone by adding kinematic tokens, topology-aware attention, and joint- attribute conditioning, both individually and in combination. For topology-aware attention, we compare Full-Mask, Mix- 4 https://huggingface.co/datasets/Lightwhe elAI/leisaac-pick-orange 5 https://github.com/LightwheelAI/leisaac Mask and Soft-Mask variants. Kinematic tokens are embed- ded with a linear–SwiGLU–linear MLP, and joint-attribute conditioning uses a linear FiLM generator. In addition, we ablate the temporal chunk size (G∈1, 2, 4, 8, 16) and the kinematic token capacity by introducing auxiliary kinematic tokens. 4.5. Training protocol All variants are initialized from theπ 0.5 checkpoint pi05-base 6 . For architectures that add morphology mod- ules, we zero-initialize the final layer of each module so that it starts near an identity mapping, stabilizing optimization. For each benchmark, our baseline is vanillaπ 0.5 , fine-tuned frompi05-baseon that benchmark. Depending on com- putational budget, we use either action policy fine-tuning (AP-FT) which updates only the diffusion action expert, or full fine-tuning (Full-FT) which updates the entire model. Figure 2 summarizes the training configuration for each setting. 5. Results 5.1. Single-embodiment results on DROID Table 3 shows average success rate (Avg SR) forπ 0.5 vari- ants, including vanilla baseline. Overall, structured mor- phology encoding improves performance, and the best re- sult is achieved by combining kinematic Tokens with Mix- Mask and joint-attribute conditioning. Below, we isolate the contribution of (1) kinematic tokens, (2) topology-aware attention bias, and (3) joint-attribute conditioning. For com- pleteness, Appendix D reports success time on DROID as well as an ablation over encoder variants. (1) kinematic tokens: Compared to the baseline (Avg SR 19.7%), adding kinematic tokens improves Avg SR to 36.0%, 6 gs://openpi-assets/checkpoints/pi05_base 6 Embedding Morphology into Transformers for Cross-Robot Policy Learning 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 com- bining KT with Mix-Mask and FiLM achieves the best SR. Bold: best SR.Underline: second-best. Kinematic Token MaskFiLM Avg SR%↑ (Chunk G=1)(95% CI) –24.7± 4.9 ✓–24.3± 4.9 ✓Full-Mask –23.8± 4.8 ✓Mix-Mask –27.3± 5.0 ✓–✓25.3± 4.9 ✓Mix-Mask✓ 28.0± 5.0 showing that a joint-centric action representation is effective even without explicit topology or semantics. (2) Topology-aware attention: Starting from the kinematic token model (Avg SR 36.0%), Mix-Mask improves Avg SR to 36.9%, suggesting that alternating topology-constrained (local) and full (global) attention is beneficial. In contrast, Full-Mask reduces Avg SR to 30.1%, indicating that enforc- ing 1-hop locality at every layer can be overly restrictive. In addition, Soft-Mask achieves Avg SR 26.1%, which is lower than both masked variants. (3) Joint-attribute conditioning: Adding FiLM-based joint-attribute conditioning yields consistent gains. Without topology encoding, FiLM improves Avg SR from 36.0% to 37.7%, and with Mix-Mask Avg SR is improved from 36.9% to 47.4%, achieving the best overall performance. Notably, the full morphology encoding (KT+Mix+FiLM) yields significantly large gains over the baseline on Task 2 and Task 3, improving success rates by 5-fold and 2.3-fold, respectively. 5.2. Single-embodiment results on Unitree G1 Dex1 Table 4 shows Avg SR for each variant on Unitree G1 Dex1. Consistent with DROID, the full morphology encoding (KT, Mix-Mask and FiLM) achieves the best performance (Avg SR 28.0%). Topology-aware attention masking and joint- attribute conditioning provide additional gains on top of kinematic tokens: Mix-Mask improves Avg SR from 24.3% to 27.3%, and adding FiLM further improves to 28.0%. 5.3. Multi-embodiment results on DROID + SO101 Figure 4 shows learning curves for multi-embodiment train- ing on the Panda/SO101 mixture. We report Macro SR, de- fined as(SR DROID + SR SO101 )/2. Overall, our embodiment- aware policy (KT, Mix-Mask, and FiLM) achieves higher Avg SR throughout training. For example, at 50k steps our 050k80k125k 0 0.1 0.2 0.3 Training steps Macro SR π 0.5 Ours (KT+Mix+FiLM) 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(SR Panda + SR SO101 )/2. Our full model outperforms theπ 0.5 baseline through- out training. For completeness, per-embodiment success rates are reported in the appendix F. Shaded regions indicate 95% confi- dence intervals. 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%↑ Chunk (G)(95% CI) 136.0± 5.4 235.8± 5.4 434.4± 5.3 830.5± 5.2 1633.3± 5.3 method reaches Avg SR of 15.5% whileπ 0.5 has a poor SR of 5.0%, and at 125k steps our method remains higher (20.7% vs. 17.5%). 5.4. Ablation study This section examines key design choices in our embodiment-aware Transformer and further probes a strong alternative for topology-aware attention. We study whether performance depends on (i) the kinematic token chunk size G, (i) auxiliary kinematic tokens (AKT), and (i) bias ini- tialization for the Soft-Mask variant, which is potentially more expressive. Effect of temporal chunk size:To study how the tempo- ral resolution of kinematic tokens affects performance, we evaluated five temporal chunk sizes (G ∈ 1, 2, 4, 8, 16), where smallerGcorresponds to more aggressive tempo- ral compression. Table 5 shows that using a single chunk (G=1) achieves the best Avg SR of 36.0%, while perfor- mance generally degrades asGincreases (33.3% atG=16). 7 Embedding Morphology into Transformers for Cross-Robot Policy Learning 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. Kinematic Tokens Avg SR%↑ MaskAKT(95% CI) –36.0± 5.4 –✓37.0± 5.4 Full-Mask–30.3± 5.2 Full-Mask✓33.0± 5.3 Mix-Mask–37.0± 5.4 Mix-Mask✓47.3± 5.6 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. Bias Init Avg SR%↑ (95% CI) Zero26.1± 4.9 Hard25.1± 4.9 Mix28.1± 5.1 Linear 20.4± 4.5 Effect of auxiliary kinematic tokens (AKT): We test whether increasing per-kinematic token capacity improves performance via auxiliary kinematic tokens (AKT). Table 6 shows that AKT consistently improves Avg SR across con- figurations. Notably, under Mix-Mask, AKT increases Avg SR from 37.0% to 47.3%, indicating that scaling kinematic token capacity can substantially boost performance. Effect of bias initialization:Since Soft-Mask attention is potentially more expressive yet optimization-sensitive, we test whether bias initialization improves its performance on DROID. All models use kinematic tokens (G=1) and no FiLM; only the initialization of the learnable bias table is changed. We study four initializations that impose different topology priors at the start of training: • Zero: All bias parameters are initialized to zero, corre- sponding to no topology prior. •Hard: biases are initialized to approximate a full-mask prior by assigning a negative bias (e.g.,−3) to non- neighbor joint pairs. •Mix: To mimic the alternating local/global schedule, even layers use the Hard initialization while odd layers are initialized to Zero. •Linear: Biases are initialized as a distance-dependent prior, linearly interpolating from0for near pairs to−3 for far pairs. Table 7 shows the effect of bias initialization. We found that Mix initialization performs best (Avg SR28.1%). Neverthe- less, across all initializations, the Soft-Mask variant does not surpass the Hard-Masked variants (e.g., Mix-Mask in Table 3). This is consistent to some related reports (Sfer- razza et al., 2024; Buterez et al., 2024). In Appendix E, we further investigate Soft-Mask with more detailed variants, but we do not observe consistent improvements. Conclusion We presented an embodiment-aware transformer policy that injects robot morphology through (1) kinematic tokens with per-joint temporal chunking, (2) a topology-aware attention mask, and (3) joint-attribute conditioning. Across single- and multi-embodiment evaluations on DROID (Franka Panda), Unitree G1 Dex1, and SO101, our structured mor- phology encoding improves success rates over the vanilla π 0.5 VLA baseline, indicating improved robustness both within an embodiment and across embodiments. In the future, we plan to (i) further study kinematic-token repre- sentations, focusing on token design choices and scalability with token capacity, (i) develop reliable Soft-Mask atten- tion methods with more stable optimization, and (i) de- velop more efficient multi-embodiment training strategies— including curriculum schedules that better preserve single- embodiment performance while scaling to additional robots. Impact Statement Our work aims to enable robot policies that perform well across diverse robot embodiments by explicitly encoding robot morphology. This fits within a broader research area working towards more capable and general robotics founda- tion models. The long-term vision of this field is to realize generalist robot policies that can readily adapt to new tasks, environments, and embodiments, in a manner more anal- ogous to the open-ended flexibility of human intelligence, in contrast to traditional methods for robotics control. We believe that our work contributes some modicum of progress towards this goal, which may eventually have profound so- cietal implications in enabling automation to drastically and rapidly replace more human labor. We feel unqualified to fully assess the potential societal and economic impacts, but think that perhaps, depending on social factors and context, these may range across the spectrum of good to bad, such as aging societies that require more automation to maintain productivity versus the challenges of a labor force made redundant. 8 Embedding Morphology into Transformers for Cross-Robot Policy Learning References Ahn, M., Brohan, A., Brown, N., Chebotar, Y., Cortes, O., David, B., Finn, C., Fu, C., Gopalakrishnan, K., Hausman, K., et al. Do as i can, not as i say: Grounding language in robotic affordances. arXiv preprint arXiv:2204.01691, 2022. Brohan, A., Brown, N., Carbajal, J., Chebotar, Y., Dabis, J., Finn, C., Gopalakrishnan, K., Hausman, K., Herzog, A., Hsu, J., et al. Rt-1: Robotics transformer for real-world control at scale. arXiv preprint arXiv:2212.06817, 2022. Brown, L. D., Cai, T. T., and DasGupta, A. Interval estima- tion for a binomial proportion. Statistical science, 16(2): 101–133, 2001. Buterez, D., Janet, J. P., Oglic, D., and Lio, P. Masked attention is all you need for graphs. arXiv preprint arXiv:2402.10793, 2024. Doshi, R., Walke, H., Mees, O., Dasari, S., and Levine, S. Scaling cross-embodied learning: One policy for ma- nipulation, navigation, locomotion and aviation. arXiv preprint arXiv:2408.11812, 2024. Driess, D., Xia, F., Sajjadi, M. S., Lynch, C., Chowdhery, A., Wahid, A., Tompson, J., Vuong, Q., Yu, T., Huang, W., et al. Palm-e: An embodied multimodal language model. arXiv preprint arXiv:2303.03378, 2023. Hong, S., Yoon, D., and Kim, K.-E. Structure-aware trans- former policy for inhomogeneous multi-task reinforce- ment learning. In International Conference on Learning Representations, 2021. Huang, W., Mordatch, I., and Pathak, D. One policy to con- trol them all: Shared modular policies for agent-agnostic control. In International Conference on Machine Learn- ing, p. 4455–4464. PMLR, 2020. Intelligence, P., Black, K., Brown, N., Darpinian, J., Dha- balia, K., Driess, D., Esmail, A., Equi, M., Finn, C., Fusai, N., et al.π0. 5: a vision-language-action model with open-world generalization, 2025. URL https://arxiv. org/abs/2504.16054, 1(2):3, 2025. Khazatsky, A., Pertsch, K., Nair, S., Balakrishna, A., Dasari, S., Karamcheti, S., Nasiriany, S., Srirama, M. K., Chen, L. Y., Ellis, K., et al.Droid: A large-scale in-the-wild robot manipulation dataset. arXiv preprint arXiv:2403.12945, 2024. Kim, M. J., Pertsch, K., Karamcheti, S., Xiao, T., Balakr- ishna, A., Nair, S., Rafailov, R., Foster, E., Lam, G., San- keti, P., et al. Openvla: An open-source vision-language- action model. arXiv preprint arXiv:2406.09246, 2024. Luo, Y., Yao, M., and Xiao, X. Gcnt: Graph-based trans- former policies for morphology-agnostic reinforcement learning. arXiv preprint arXiv:2505.15211, 2025. O’Neill, A., Rehman, A., Maddukuri, A., Gupta, A., Padalkar, A., Lee, A., Pooley, A., Gupta, A., Mandlekar, A., Jain, A., et al. Open x-embodiment: Robotic learning datasets and rt-x models: Open x-embodiment collab- oration 0. In 2024 IEEE International Conference on Robotics and Automation (ICRA), p. 6892–6903. IEEE, 2024. Perez, E., Strub, F., De Vries, H., Dumoulin, V., and Courville, A. Film: Visual reasoning with a general con- ditioning layer. In Proceedings of the AAAI conference on artificial intelligence, volume 32, 2018. Sferrazza, C., Huang, D.-M., Liu, F., Lee, J., and Abbeel, P. Body transformer: Leveraging robot embodiment for policy learning. arXiv preprint arXiv:2408.06316, 2024. Team, O. M., Ghosh, D., Walke, H., Pertsch, K., Black, K., Mees, O., Dasari, S., Hejna, J., Kreiman, T., Xu, C., et al. Octo: An open-source generalist robot policy. arXiv preprint arXiv:2405.12213, 2024. Veli ˇ ckovi ́ c, P., Cucurull, G., Casanova, A., Romero, A., Lio, P., and Bengio, Y. Graph attention networks. arXiv preprint arXiv:1710.10903, 2017. Walke, H. R., Black, K., Zhao, T. Z., Vuong, Q., Zheng, C., Hansen-Estruch, P., He, A. W., Myers, V., Kim, M. J., Du, M., et al. Bridgedata v2: A dataset for robot learning at scale. In Conference on Robot Learning, p. 1723–1736. PMLR, 2023. Wang, T., Liao, R., Ba, J., and Fidler, S. Nervenet: Learning structured policy with graph neural networks. In Interna- tional conference on learning representations, 2018. Whitman, J., Travers, M., and Choset, H. Learning modular robot control policies. IEEE Transactions on Robotics, 39(5):4095–4113, 2023. Ying, C., Cai, T., Luo, S., Zheng, S., Ke, G., He, D., Shen, Y., and Liu, T.-Y. Do transformers really perform badly for graph representation? Advances in neural information processing systems, 34:28877–28888, 2021. 9 Embedding Morphology into Transformers for Cross-Robot Policy Learning A. Attention Mask 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 π 0.5 mask. 10 Embedding Morphology into Transformers for Cross-Robot Policy Learning B. Task Distribution in the DROID 1/8 Subset We summarize the task distribution of the DROID 1/8 training subset used in our experiments. We compute frequency statistics of verbs (Figure 6) and objects (Figure 7) appearing in the task instructions. 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. 11 Embedding Morphology into Transformers for Cross-Robot Policy Learning C. Training Protocol Table 8. Trainable vs. frozen parameters in AP-FT: VLM parameters are frozen, while the action-policy parameters are optimized. StatusPart Parameter name Frozen: VLM part (2,923,335,408 parameters) FrozenVLM PaliGemma/img/Transformer/encoder_norm/bias FrozenVLM PaliGemma/img/Transformer/encoder_norm/scale FrozenVLM PaliGemma/img/Transformer/encoderblock/LayerNorm_0/bias FrozenVLM PaliGemma/img/Transformer/encoderblock/LayerNorm_0/scale FrozenVLM PaliGemma/img/Transformer/encoderblock/LayerNorm_1/bias FrozenVLM PaliGemma/img/Transformer/encoderblock/LayerNorm_1/scale FrozenVLM PaliGemma/img/Transformer/encoderblock/MlpBlock_0/Dense_0/bias FrozenVLM PaliGemma/img/Transformer/encoderblock/MlpBlock_0/Dense_0/kernel FrozenVLM PaliGemma/img/Transformer/encoderblock/MlpBlock_0/Dense_1/bias FrozenVLM PaliGemma/img/Transformer/encoderblock/MlpBlock_0/Dense_1/kernel FrozenVLM PaliGemma/img/Transformer/encoderblock/MultiHeadDotProductAttention_0/key/bias FrozenVLM PaliGemma/img/Transformer/encoderblock/MultiHeadDotProductAttention_0/key/kernel FrozenVLM PaliGemma/img/Transformer/encoderblock/MultiHeadDotProductAttention_0/out/bias FrozenVLM PaliGemma/img/Transformer/encoderblock/MultiHeadDotProductAttention_0/out/kernel FrozenVLM PaliGemma/img/Transformer/encoderblock/MultiHeadDotProductAttention_0/query/bias FrozenVLM PaliGemma/img/Transformer/encoderblock/MultiHeadDotProductAttention_0/query/kernel FrozenVLM PaliGemma/img/Transformer/encoderblock/MultiHeadDotProductAttention_0/value/bias FrozenVLM PaliGemma/img/Transformer/encoderblock/MultiHeadDotProductAttention_0/value/kernel FrozenVLM PaliGemma/img/embedding/bias FrozenVLM PaliGemma/img/embedding/kernel FrozenVLM PaliGemma/img/head/bias FrozenVLM PaliGemma/img/head/kernel FrozenVLM PaliGemma/img/pos_embedding FrozenVLM PaliGemma/llm/embedder/input_embedding FrozenVLM PaliGemma/llm/final_norm/scale FrozenVLM PaliGemma/llm/layers/attn/attn_vec_einsum/w FrozenVLM PaliGemma/llm/layers/attn/kv_einsum/w FrozenVLM PaliGemma/llm/layers/attn/q_einsum/w FrozenVLM PaliGemma/llm/layers/mlp/gating_einsum FrozenVLM PaliGemma/llm/layers/mlp/linear FrozenVLM PaliGemma/llm/layers/pre_attention_norm/scale FrozenVLM PaliGemma/llm/layers/pre_ffw_norm/scale Trainable: Action policy part (430,098,464 parameters) TrainableAction PaliGemma/llm/final_norm_1/Dense_0/bias TrainableAction PaliGemma/llm/final_norm_1/Dense_0/kernel TrainableAction PaliGemma/llm/layers/attn/attn_vec_einsum_1/w TrainableAction PaliGemma/llm/layers/attn/kv_einsum_1/w TrainableAction PaliGemma/llm/layers/attn/q_einsum_1/w TrainableAction PaliGemma/llm/layers/mlp_1/gating_einsum TrainableAction PaliGemma/llm/layers/mlp_1/linear TrainableAction PaliGemma/llm/layers/pre_attention_norm_1/Dense_0/bias TrainableAction PaliGemma/llm/layers/pre_attention_norm_1/Dense_0/kernel TrainableAction PaliGemma/llm/layers/pre_ffw_norm_1/Dense_0/bias TrainableAction PaliGemma/llm/layers/pre_ffw_norm_1/Dense_0/kernel TrainableAction action_in_proj/bias TrainableAction action_in_proj/kernel TrainableAction action_out_proj/bias TrainableAction action_out_proj/kernel TrainableAction time_mlp_in/bias TrainableAction time_mlp_in/kernel TrainableAction time_mlp_out/bias TrainableAction time_mlp_out/kernel 12 Embedding Morphology into Transformers for Cross-Robot Policy Learning D. Detailed DROID experiments In addition to success rate (SR), this table reports time-to-success (task completion time) and includes additional ablation variants. Time-to-success is an auxiliary metric and may trade off with SR. D.1. Full-/Mix-Mask variants without joint-attribute conditioning (FiLM) Table 9. DROID simulation evaluation summary: Best and second-best are indicated by bold and underline, respectively. Training Kinematic Token EncoderMask Batch Train DataAvgTask 1Task 2Task 3 Chunk AKTSR(%)↑ Time[s]↓ SR(%)↑ Time[s]↓ SR(%)↑ Time[s]↓ SR(%)↑ Time[s]↓ Full-FT–256 Full Dataset66.33.21565.72.52667.74.17865.72.939 AP-FT–32 1/8 subset19.74.84918.34.73115.76.61625.03.201 AP-FT1– linear–32 1/8 subset26.74.60614.75.74634.34.98431.03.088 AP-FT1✓ linear–32 1/8 subset35.24.57115.34.82457.05.17533.33.714 AP-FT1– lin swiGlulin –32 1/8 subset36.04.92510.35.45067.75.25430.04.072 AP-FT1✓ linswiGlulin –32 1/8 subset37.04.36111.04.61860.35.11039.73.355 AP-FT16– linear–32 1/8 subset16.95.02614.35.56516.76.09019.73.422 AP-FT16✓ linear–32 1/8 subset18.65.1908.75.27727.35.99919.74.294 AP-FT16– lin swiGlulin –32 1/8 subset33.33.72427.33.48955.74.50017.03.182 AP-FT16✓ lin swiGlulin –32 1/8 subset31.04.87616.75.66453.74.74822.74.216 AP-FT1– linearFull32 1/8 subset30.55.31716.05.37837.06.54138.04.033 AP-FT1✓ linearFull32 1/8 subset28.04.82411.04.55651.06.01822.03.899 AP-FT1– linswiGlulin Full32 1/8 subset30.14.90215.34.57234.06.50141.03.632 AP-FT1✓ linswiGlulin Full32 1/8 subset33.04.52016.04.29661.05.12022.04.144 AP-FT16– linearFull32 1/8 subset11.95.5683.35.81219.76.60312.74.289 AP-FT16✓ linearFull32 1/8 subset20.54.4718.04.8434.75.659 48.72.912 AP-FT16– linswiGlulin Full32 1/8 subset11.85.37213.05.3032.06.95020.33.856 AP-FT16✓ linswiGlulin Full32 1/8 subset16.95.20313.76.0263.06.25834.03.326 AP-FT1– linearMix32 1/8 subset27.74.83513.35.77733.05.37036.73.358 AP-FT1✓ linearMix32 1/8 subset35.74.53822.34.93444.05.50240.73.178 AP-FT1– linswiGlulin Mix32 1/8 subset36.94.97927.74.84856.75.55526.34.534 AP-FT1✓ linswiGlulin Mix32 1/8 subset47.13.11721.05.01079.34.58941.03.548 AP-FT16– linearMix32 1/8 subset21.15.2720.0–21.76.32641.74.218 AP-FT16✓ linearMix32 1/8 subset27.25.1056.35.80751.75.91623.73.591 AP-FT16– linswiGlulin Mix32 1/8 subset19.05.2840.06.24025.05.94930.33.662 AP-FT16✓ linswiGlulin Mix32 1/8 subset29.64.8495.34.51055.35.68928.04.348 D.2. Full-/Mix-Mask variants with joint-attribute conditioning (FiLM) Table 10. DROID simulation evaluation summary: Best and second-best are indicated by bold andunderline, respectively. Training Kinematic Token EncoderMask FiLM Batch Train DataAvgTask 1Task 2Task 3 Chunk AKTSR(%)↑ Time[s]↓ SR(%)↑ Time[s]↓ SR(%)↑ Time[s]↓ SR(%)↑ Time[s]↓ AP-FT–32 1/8 subset19.74.84918.34.73115.76.61625.03.201 AP-FT–32 1/8 subset19.74.84918.34.73115.76.61625.03.201 AP-FT1– linswiGlulin –32 1/8 subset36.74.74711.35.80267.34.59731.33.843 AP-FT1– linswiGlulin –✓32 1/8 subset37.73.7946.03.72765.04.48842.03.168 AP-FT1✓ lin swiGlulin –32 1/8 subset37.04.36111.04.61860.35.11039.73.355 AP-FT1✓ linswiGlulin –✓32 1/8 subset39.44.24514.64.50272.74.60631.03.626 AP-FT1– linswiGlulin Mix–32 1/8 subset36.34.88127.04.88154.05.22028.04.543 AP-FT1– linswiGlulin Mix✓32 1/8 subset47.44.6415.76.01977.74.87758.73.028 AP-FT1✓ linswiGlulin Mix–32 1/8 subset47.13.11721.05.01079.34.58941.03.548 13 Embedding Morphology into Transformers for Cross-Robot Policy Learning E. Soft-Mask variants and stability exploration Soft-Mask attention is potentially more expressive than Hard-Mask designs because it retains fully connected attention while injecting topology as an additive bias on the joint-to-joint attention logits. In practice, however, Soft-Mask can be optimization-sensitive. Motivated by the bias-initialization ablation in the main text, here we provide a more detailed exploration of Soft-Mask designs, focusing on parameterizations and initialization choices that may improve training stability. We organize Soft-Mask variants into two families: (i) Adj-SoftMask, which uses a binary adjacency indicator (connected vs. disconnected) and learns a suppression strength; and (i) SPD-SoftMask, which uses a shortest-path-distance (SPD) index (Graphormer-style) and learns a distance-indexed bias table (the main-text choice). Soft-Mask: general form:We inject topology into the joint-to-joint attention through an additive bias termB (ℓ) ij applied to the attention logits. Following the main text, we write α (ℓ) ij = softmax j logits (ℓ) ij + B (ℓ) ij .(11) Different Soft-Mask variants correspond to different parameterizations of B (ℓ) ij . E.1. Adj-SoftMask (adjacency + learnable strength) Adj-SoftMask parameterizesB (ℓ) ij using the binary adjacency indicatorM ij ∈0, 1defined in the main text (M ij = 1for i = j or (i,j)∈E , and M ij = 0 otherwise). We instantiate the bias as B (ℓ) ij = (M ij − 1)s (ℓ) ij , s (ℓ) ij ∈R,(12) so thatB (ℓ) ij = 0whenM ij = 1andB (ℓ) ij =−s (ℓ) ij whenM ij = 0. Thus, the effect on non-adjacent pairs depends on the sign of s (ℓ) ij . Adj-SoftMask (v1.0): We instantiate Eq. (12) with an edge-wise, layer-wise strength s (ℓ) ij = exp min θ (ℓ) ij ,θ max , θ (ℓ) ij ∈R.(13) Since s (ℓ) ij > 0, the resulting bias on disconnected pairs (M ij = 0) is always negative. Adj-SoftMask (v1.1): We share the strength across edges within each layer by using a single scalar θ (ℓ) ∈R: s (ℓ) ij = exp min θ (ℓ) ,θ max .(14) Again, s (ℓ) ij > 0 implies a strictly negative bias on disconnected pairs. Adj-SoftMask (v2.0): We use a layer-wise scalar without exponential mapping: s (ℓ) ij = θ (ℓ) , θ (ℓ) ∈R.(15) Compared to v1.x, this formulation avoids the sharp scaling introduced byexp(·)and is intended to improve optimization stability. Initialization: We consider two initializations, matching the main-text bias-initialization taxonomy: • Zero: initialize the effect to be weak (close to NO-MASK). • Hard: initialize the effect to be strong (close to HARD-MASK). 14 Embedding Morphology into Transformers for Cross-Robot Policy Learning E.2. SPD-SoftMask Adj-SoftMask uses a binary adjacency signal (M ij ) and cannot distinguish how far two joints are on the kinematic graph. SPD-SoftMask (the main-text SOFT-MASK) instead indexesB (ℓ) ij by the shortest-path distance (SPD), enabling distance-dependent inductive bias while preserving fully connected attention. LetG = (V,E )be the kinematic graph and letd(i,j)denote the SPD between jointsiandjas defined in the main text. We parameterize the topology term using a learnable bias table: B (ℓ) ij = θ (ℓ) d(i,j) , θ (ℓ) d ∈R, d∈0, 1,...,D max .(16) The resulting joint-to-joint attention follows Eq. (11). Bias initialization: To probe sensitivity to initialization, we consider four initializations of the SPD bias table θ (ℓ) d : • Zero: initialize θ (ℓ) d = 0 (no topology prior). • Hard: initialize θ (ℓ) d to strongly suppress larger distances (strong locality prior). • Mix: use Hard on even layers and Zero on odd layers. • Linear: initializeθ (ℓ) d with a distance-dependent prior, interpolating from0to a negative value (e.g.,−3) asdincreases. Warm-start transfer:In addition to enabling SPD-SoftMask frompi05-base, we also study a warm-start setting where SPD-SoftMask is switched on starting from stronger checkpoints: (i) api05-base+JT model (trained with kinematic tokens but without SPD bias), and (i) a pi05-base+JT+Mix model (trained with kinematic tokens and Mix-Mask). We then continue training under the same protocol to evaluate whether SPD-SoftMask can serve as a refinement step on top of strong masked baselines. E.3. Results Across these experiments (Table 11), we did not find a Soft-Mask configuration that surpasses the Hard-Mask Mix- Mask variant under our training protocol. Within Adj-SoftMask, v2.0 achieves the strongest results among the tested parameterizations, suggesting that learning the bias magnitude directly can be preferable to exponential mappings. For SPD-SoftMask, performance is sensitive to the bias initialization and warm-start choice, and we do not observe consistent gains over strong masked baselines. 15 Embedding Morphology into Transformers for Cross-Robot Policy Learning Table 11. Soft-Mask results (SR%): We report per-task success rates (Task1–Task3) in % along with macro averages, shown as SR%± 95% CI. Base ModelVariantInit.Success Rate (SR%) (95% CI)↑ AvgTask1Task2Task3 pi05-baseSoft-Mask (v1.0) Zero25.9± 4.9 26.7± 5.0 20.3± 4.5 30.7± 5.2 pi05-baseSoft-Mask (v1.0) Hard25.4± 4.9 17.7± 4.3 15.7± 4.142.7± 5.6 pi05-baseSoft-Mask (v1.1) Zero26.4± 5.0 21.3± 4.6 16.7± 4.2 41.3± 5.5 pi05-baseSoft-Mask (v1.1) Hard29.5± 5.1 16.7± 4.228.7± 5.1 43.0± 5.6 pi05-baseSoft-Mask (v2.0) Zero 34.4± 5.3 27.6± 5.0 37.0± 5.4 38.6± 5.5 pi05-baseSoft-Mask (v2.0) Hard23.0± 4.7 26.3± 5.0 11.6± 3.6 31.0± 5.2 pi05-baseSoft-Mask (v3.0) Zero26.1± 4.929.7± 5.1 17.3± 4.3 31.3± 5.2 pi05-baseSoft-Mask (v3.0) Hard25.1± 4.9 18.0± 4.3 20.7± 4.6 36.7± 5.4 pi05-baseSoft-Mask (v3.0) Mix28.1± 5.1 17.0± 4.2 24.3± 4.8 43.0± 5.6 pi05-baseSoft-Mask (v3.0) Linear 20.4± 4.5 11.7± 3.6 16.7± 4.2 32.7± 5.3 pretrained π 0.5 w/ KTSoft-Mask (v3.0) –25.8± 4.9 23.0± 4.7 28.3± 5.1 26.0± 4.9 pretrained π 0.5 w/ KT+Mix-Mask Soft-Mask (v3.0) –28.9± 5.1 45.3± 5.6 17.7± 4.3 23.7± 4.8 16 Embedding Morphology into Transformers for Cross-Robot Policy Learning F. Multi-embodiment Results Figure 8 reports learning curves for multi-embodiment joint training on the Panda–SO101 mixture, evaluated separately on DROID (Panda) and SO101 in terms of average success rate (Avg SR). For our method, we use full morphology encoding: kinematic tokens (chunkG=1), Mix-Mask, and FiLM-based joint-attribute conditioning. Overall, our embodiment-aware policy improves performance on DROID throughout training, while achieving comparable performance on SO101. DROID (Panda): Our method improves earlier and remains higher: at 50k steps it reaches Avg SR 0.210 (π 0.5 : 0.000), and at 125k steps it remains higher at 0.213 (π 0.5 : 0.100). SO101: Performance is comparable across training: at 50k steps both reach Avg SR 0.100, and at 125k steps our method achieves 0.200 while π 0.5 achieves 0.250 (difference: 0.050). We hypothesize that the larger gains on DROID are due to the training setup and the mixture ratio being skewed toward Panda, which can favor improvements on DROID relative to SO101. Under AP-FT, single-embodiment SO101 training with π 0.5 reaches only about Avg SR 0.2, suggesting that SO101 performance is generally more constrained in this fine-tuning regime. 050k80k125k 0 0.1 0.2 0.3 Training steps Avg SR π 0.5 (DROID) π 0.5 (SO101) Ours (DROID) Ours (SO101) Figure 8. Multi-embodiment learning curves on Panda–SO101: Shaded regions indicate 95% CI (per-checkpoint). 17