One-sentence Summary

The authors propose Latent Interface Training (LIT), a two-stage framework that mitigates vision–action shortcuts by first learning a goal-conditioned action prior without images and then constraining visual information through a pose-supervised latent interface, which boosts robot foundation model generalization under visual distribution shifts with 3.87–10.70 percentage-point gains on LIBERO-Plus across four architectures (π0.5, MolmoAct2, FAST-WAM, and ImageWAM) and 13.30–16.70 percentage-point gains in real-world tasks under unseen camera configurations, lighting variations, and distractors.

Key Contributions

• Latent Interface Training (LIT) is a framework-agnostic two-stage strategy that mitigates vision–action shortcuts by constraining the action expert’s visual conditioning to a pose-supervised latent interface.
• In the first stage, the action expert is pretrained without images on language, robot state, and each action chunk’s terminal SE(3) end-effector pose. In the second stage, a latent interface aggregates visual and semantic representations and is supervised to reconstruct that terminal pose, serving as the action expert’s sole visual conditioning pathway.
• When applied to π0.5, MolmoAct2, FAST-WAM, and ImageWAM, LIT preserves LIBERO performance while improving zero-shot LIBERO-Plus success by 3.87–10.70 percentage points, and yields real-world gains of 13.30–16.70 percentage points across three tasks with unseen camera setups, lighting variations, and distractors.

Introduction

Robot foundation models that combine vision–language backbones with action experts achieve strong in-distribution manipulation performance but struggle to generalize when visual conditions shift. Because action experts receive rich visual representations, they can latch onto task-irrelevant visual cues that correlate with training actions—a vision–action shortcut—and this reliance is worsened by limited visual diversity in robot demonstration data. Prior work tries to address this by injecting structured spatial information or by pretraining image-free action priors, yet neither approach directly constrains how the action expert uses visual inputs when they are reintroduced. The authors propose Latent Interface Training (LIT), a model-agnostic two-stage strategy that first learns a spatial-goal-conditioned action prior without images, then introduces visual conditioning exclusively through a pose-supervised latent interface. By forcing all visual information through this interface and requiring it to reconstruct the same end-effector pose used during pretraining, LIT reduces vision–action shortcuts while preserving the spatial information needed for accurate action generation.

Method

The authors propose Latent Interface Training (LIT), a model agnostic two-stage strategy designed to reduce vision-action shortcuts while preserving task-relevant spatial information. The overall framework first establishes an action prior without images and then introduces visual conditioning exclusively through a pose-supervised latent interface.

Refer to the framework diagram:

Problem Setup At each timestep $t$t, the policy receives a visual observation $\mathbf{o}_t$ot​, a language instruction $l$l, and a robot state $\mathbf{s}_t$st​ to predict an action chunk of horizon $H$H:

$$\mathbf{A}_t = (\mathbf{a}_t, \dots, \mathbf{a}_{t+H-1}).$$At​=(at​,…,at+H−1​).

For each demonstrated chunk, the terminal robot state is used as the spatial goal:

$$\mathbf{g}_t = [\mathbf{p}_{t+H}; \mathbf{r}_{t+H}; \mathbf{q}_{t+H}] \in \mathbb{R}^8,$$gt​=[pt+H​;rt+H​;qt+H​]∈R8,

where $\mathbf{p} \in \mathbb{R}^3$p∈R3 and $\mathbf{r} \in \mathbb{R}^3$r∈R3 represent the world-frame end-effector position and axis-angle orientation, and $\mathbf{q} \in \mathbb{R}^2$q∈R2 contains the gripper joint positions. This goal acts as a conditioning signal in the first stage and a reconstruction target in the second stage.

Stage 1: Spatial-Goal-Conditioned Action Pretraining In the first stage, the authors train the action expert from scratch to generate demonstrated action chunks conditioned on the language instruction, robot state, and the terminal pose $\mathbf{g}_t$gt​, effectively learning a spatial-goal-conditioned action prior without any visual input. The frozen backbone processes only the language instruction and robot state, producing semantic representations $\mathbf{H}_{\ell,t}^{\text{sem}}$Hℓ,tsem​ at each of the $L$L coupling layers. A trainable three-layer MLP with GELU activations maps the goal $\mathbf{g}_t$gt​ to goal tokens $\mathbf{G}_t$Gt​, which are concatenated with the backbone representations:

$$\begin{array}{rl} & \mathbf{G}_t = E_\eta(\mathbf{g}_t), \\ & \mathbf{C}_{\ell,t} = [\mathbf{H}_{\ell,t}^{\text{sem}}; \mathbf{G}_t], \qquad \ell = 1, \dots, L. \end{array}$$​Gt​=Eη​(gt​),Cℓ,t​=[Hℓ,tsem​;Gt​],ℓ=1,…,L.​

These concatenated tokens $\mathbf{C}_{\ell,t}$Cℓ,t​ condition the corresponding action expert layer through the native conditioning mechanism. The authors retain the framework's native flow-matching objective. They sample $\tau \sim \mathcal{U}(0,1)$τ∼U(0,1) and $\epsilon \sim \mathcal{N}(\mathbf{0}, \mathbf{I})$ϵ∼N(0,I) to construct the noisy action chunk and target velocity:

$$\begin{array}{rl} & \widetilde{\mathbf{A}}_t^\tau = (1-\tau)\boldsymbol{\epsilon} + \tau \mathbf{A}_t, \\ & \mathbf{v}_t^\star = \mathbf{A}_t - \boldsymbol{\epsilon}. \end{array}$$​Atτ​=(1−τ)ϵ+τAt​,vt⋆​=At​−ϵ.​

Given the conditioning tokens, the action expert predicts the velocity and minimizes the prior loss:

$$\begin{array}{r} \widehat{\mathbf{v}}_t^\tau = v_\theta \Big(\widetilde{\mathbf{A}}_t^\tau, \tau; \mathbf{C}_{1:L,t}\Big), \\ \mathcal{L}_{\text{prior}} = \mathbb{E}_{\mathbf{A}_t, \tau, \boldsymbol{\epsilon}} \left[ \| \widehat{\mathbf{v}}_t^\tau - \mathbf{v}_t^\star \|_2^2 \right]. \end{array}$$vtτ​=vθ​(Atτ​,τ;C1:L,t​),Lprior​=EAt​,τ,ϵ​[∥vtτ​−vt⋆​∥22​].​

Only the action expert parameters $\theta$θ and the SE(3) encoder parameters $\eta$η are updated, while the backbone and its modality encoders remain frozen.

Stage 2: Vision-Action Interface Learning In the second stage, the action expert is initialized from Stage 1, and visual conditioning is introduced exclusively through a pose-supervised latent interface. The authors define learnable latent tokens $\mathbf{Z}^0 \in \mathbb{R}^{K \times d}$Z0∈RK×d shared across inputs, with $K=100$K=100. For each policy input, the interface starts from $\mathbf{Z}_{0,t} = \mathbf{Z}^0$Z0,t​=Z0. At coupling layer $\ell$ℓ, the backbone provides language and state representations $\mathbf{H}_{\ell,t}^{\text{sem}}$Hℓ,tsem​ alongside visual representations $\mathbf{H}_{\ell,t}^{\text{vis}}$Hℓ,tvis​. The latent tokens are updated through self-attention, semantic cross-attention, and visual cross-attention:

$$\begin{array}{l} \overline{\mathbf{Z}}_{\ell,t} = \mathbf{Z}_{\ell-1,t} + \text{SA}_{q(\ell)}(\mathbf{Z}_{\ell-1,t}), \\ \widetilde{\mathbf{Z}}_{\ell,t} = \overline{\mathbf{Z}}_{\ell,t} + \text{CA}_{q(\ell)}^{\text{sem}}(\overline{\mathbf{Z}}_{\ell,t}; \mathbf{H}_{\ell,t}^{\text{sem}}), \\ \mathbf{Z}_{\ell,t} = \widetilde{\mathbf{Z}}_{\ell,t} + \text{CA}_{q(\ell)}^{\text{vis}}(\widetilde{\mathbf{Z}}_{\ell,t}; \mathbf{H}_{\ell,t}^{\text{vis}}), \quad \ell = 1, \dots, L. \end{array}$$Zℓ,t​=Zℓ−1,t​+SAq(ℓ)​(Zℓ−1,t​),Zℓ,t​=Zℓ,t​+CAq(ℓ)sem​(Zℓ,t​;Hℓ,tsem​),Zℓ,t​=Zℓ,t​+CAq(ℓ)vis​(Zℓ,t​;Hℓ,tvis​),ℓ=1,…,L.​

The latent tokens act as queries in each cross-attention operation, while the backbone representations provide the keys and values. For parameter efficiency, every $m$m consecutive coupling layers share interface attention parameters, indexed by $q(\ell) = \lceil \ell/m \rceil$q(ℓ)=⌈ℓ/m⌉. The updated tokens condition the action expert layers, which predict the velocity:

$$\widehat{\mathbf{v}}_t^\tau = v_\theta \Big(\widetilde{\mathbf{A}}_t^\tau, \tau; \mathbf{Z}_{1:L,t}\Big).$$vtτ​=vθ​(Atτ​,τ;Z1:L,t​).

An MLP decoder reconstructs the terminal goal state $\mathbf{g}_t$gt​ from the final latent tokens:

$$\begin{array}{c} \widehat{\mathbf{g}}_t = D_\omega(\mathbf{Z}_{L,t}), \\ \mathcal{L}_{\text{pose}} = \| \widehat{\mathbf{g}}_t - \mathbf{g}_t \|_2^2, \\ \mathcal{L}_{\text{stage2}} = \mathcal{L}_{\text{act}} + \lambda_{\text{pose}} \mathcal{L}_{\text{pose}}. \end{array}$$g​t​=Dω​(ZL,t​),Lpose​=∥g​t​−gt​∥22​,Lstage2​=Lact​+λpose​Lpose​.​

The authors set $\lambda_{\text{pose}} = 0.3$λpose​=0.3. This reconstruction loss encourages the interface to retain goal-relevant information. Stage 2 jointly optimizes the backbone, the action expert, the latent tokens, the interface attention modules, and the MLP decoder.

Framework Integration and Inference LIT introduces a latent interface between a pretrained backbone and an action expert, making it applicable to various architectures. The interface provides layer-wise conditioning through each architecture's native mechanism while retaining the original architectures and objectives. At inference, the latent interface remains active, while the pose encoder and pose decoder are omitted. The policy requires only visual observations, language, and robot state, following the native action sampling and execution procedures.

Experiment

LIT was evaluated across multiple VLA and WAM architectures, where it preserved or slightly improved in-distribution task performance and consistently enhanced zero-shot generalization under diverse visual and task perturbations. Real-world robot experiments confirmed increased robustness to lighting, camera, and distractor changes, and analyses showed that LIT reduces reliance on spurious visual correlations by routing visual information through a pose-supervised latent interface. Ablation studies further verified that the combination of spatial-goal-conditioned action pretraining, pose supervision, and restricted visual access is essential for these gains.

LIT consistently improves in-distribution success rates over the baseline across diverse architectures on LIBERO. The benefit is more pronounced for lower-performing baselines, while already strong models such as FAST-WAM see only marginal gains. ImageWAM attains the highest overall average, reaching perfect or near-perfect scores on object and goal subtasks. LIT raises the average success of the base model from 87.75% to 91.80%, a gain of over 4 percentage points, whereas for FAST-WAM the improvement shrinks to 0.5 points (from 97.60% to 98.10%). ImageWAM achieves a perfect 100% on object and 97.60% on goal subtasks, resulting in the highest average in-distribution success rate (98.10%) among all evaluated variants.

LIT delivers consistent zero-shot improvements across all evaluated manipulation architectures on LIBERO-Plus, with the largest absolute gains under camera viewpoint changes and sensor noise. Architectures starting from low baselines gain the most, while already-robust models like ImageWAM see meaningful improvement mainly under robot initial state variability. Ablation studies show that both the restricted visual pathway and pose supervision are essential for this generalization, especially against task-irrelevant visual perturbations. For the π₀.₅ VLA, LIT raises success on camera viewpoint changes by 22.0 percentage points and on sensor noise by 12.1 points. MolmoAct2 gains 21.7 points under sensor noise and 15.7 points on object layout changes with LIT. FAST-WAM, with a baseline of only 16.4%, improves by 27.4 points on camera viewpoint perturbations after LIT. ImageWAM already exceeds 96% on sensor noise and lighting; its largest LIT gain is 12.4 points on robot initial states. Removing pose supervision drops sensor noise performance by 10.8 points and background texture performance by 6.0 points. Giving the action expert direct visual access reduces overall OOD success from 71.9% to 67.7%, supporting the restricted latent interface.

LIT substantially improves over the MolmoAct2 baseline in overall out-of-distribution success, with each of its three design elements—Stage 1 action pretraining, pose supervision, and restricted visual pathway—contributing to the gain. Ablating any single component reduces generalization, and alternative simpler designs fail to match LIT, showing that the combination of a spatial-goal-conditioned prior, explicit pose reconstruction, and a latent-only visual bottleneck is essential for the observed robustness. Removing Stage 1 action pretraining lowers overall OOD success, especially under robot initial states and object layout variations, confirming that spatial-goal-conditioned pretraining aids visual policy adaptation. Ablating pose supervision causes the largest drops in sensor noise and background texture scenarios, aligning with the idea that spatial-goal supervision improves robustness to task-irrelevant visual perturbations. Allowing the action expert direct visual access rather than routing through the latent interface reduces OOD success, supporting the importance of the restricted visual pathway for generalization. Training only the latent-token aggregation without the action prior and pose supervision yields only minor improvement over the baseline, indicating the latent interface alone is far from sufficient. Staged training without LIT’s specific components or simply adding pose supervision to the baseline both reach only about 65.5% OOD success, well below LIT, ruling out these simpler explanations for the gains.

The experiments evaluate LIT's in-distribution performance on LIBERO and zero-shot generalization under diverse perturbations on LIBERO-Plus. LIT consistently improves success rates, with the largest gains for lower-performing baselines, while already strong models see only marginal benefits. Across camera changes, sensor noise, and visual variations, the restricted visual pathway and explicit pose supervision prove essential for robustness, and ablations confirm that all three design components—spatial-goal-conditioned action pretraining, pose reconstruction, and the latent-only bottleneck—jointly underpin the observed out-of-distribution improvements, ruling out simpler alternatives.