Skip to content

Align-Then-stEer: Adapting the Vision-Language Action Models through Unified Latent Guidance

🕒 Published (v1): 2026-01-01 · Source: ICLR · Venue: ICLR 2026 · link

Why this paper was selected

ICLR 2026; unified latent guidance adapts VLA models to downstream manipulation tasks

Ask a follow-up

Open an assistant pre-loaded with this paper's context.

💬 Ask ChatGPT✦ Ask Claude

TL;DR

ATE (Align-Then-stEer) is a plug-and-play adaptation framework for pre-trained diffusion/flow-based Vision-Language-Action Models (VLAs) that resolves action-distribution mismatch via a two-stage pipeline: a unified latent space built with asymmetric VAEs, followed by classifier guidance that steers fine-tuning toward the target domain. It targets cross-embodiment and cross-task adaptation under limited data, requiring no architectural changes to the base VLA. On real-world dual-arm manipulation it yields a 32% success rate gain over direct fine-tuning.

Problem

Pre-trained VLAs fine-tuned on downstream tasks suffer from severe action-distribution mismatch when (1) the target robot embodiment (e.g., dual-arm or humanoid) differs from the pre-training fleet (mostly single-arm), or (2) tasks and setups diverge significantly from pre-training. Existing parameter-efficient methods (LoRA, dynamic layer activation) address training cost but do not resolve the underlying distributional gap, demanding large adaptation datasets and/or many compute cycles.

Method

ATE operates in two sequential stages:

Stage 1 — Unified Action Latent Space. Two InfoVAEs are trained: a pre-training VAE \(V_\text{pretrain} = \{E_\phi, D_\phi\}\) trained on large cross-embodiment data using the InfoVAE objective (reconstruction + KL toward \(\mathcal{N}(0,I)\) + MMD); and an adaptation VAE \(V_\text{adapt} = \{E_\psi, D_\psi\}\) trained on the small target-domain dataset with a reverse KL divergence constraint toward the pre-training latent distribution \(q_\phi(z)\):

\[\min_\psi \; D_\text{KL}\!\left(q_\psi(z|\tilde{a}_{t:t+L-1})\,\|\,q_\phi(z)\right)\]

Reverse KL is mode-seeking, so \(q_\psi\) concentrates inside one mode of \(q_\phi\), producing a structured, unified latent space \(\mathcal{Z}\) that encompasses both action spaces despite differing chunk lengths (\(H\) vs. \(L\)) and action dimensionalities.

Stage 2 — Latent Classifier Guidance. During fine-tuning, a guidance signal \(g\) is computed as the gradient of an energy-based classifier measuring latent-space discrepancy between the noisy action at denoising step \(k\) and the ground-truth action:

\[g = \nabla_{\hat{a}^k_{t:t+h}} \log p_\psi(y|\hat{a}^k_{t:t+h}) \propto -\nabla_{\hat{a}^k_{t:t+h}} \|E_\psi(\hat{a}^k_{t:t+h}) - E_\psi(a^0_{t:t+h})\|^2\]

For diffusion-based VLAs this modifies the denoising objective by adding \(\sqrt{1-\bar\alpha_k}\cdot\lambda\cdot g\) to the noise prediction target. For flow-based VLAs the analogous correction is \(\frac{1-\tau}{\tau}\cdot\lambda\cdot g\) added to the velocity target. Both VAEs are frozen during VLA fine-tuning; only the base VLA weights are updated.

Key Contributions

  • Unified latent alignment: InfoVAE pair with reverse-KL constraint embeds heterogeneous target actions into a mode of the pre-training latent prior, bridging cross-embodiment and cross-task action-space gaps.
  • Latent classifier guidance: Energy-based guidance function operating in \(\mathcal{Z}\) steers the diffusion/flow denoising trajectory at every noise level toward the target distribution without additional data.
  • Model-agnostic and architecture-preserving: ATE is plug-and-play; it works with any diffusion- or flow-based VLA (demonstrated on Diffusion Policy, RDT-1B, \(\pi_0\)) without modifying the VLA architecture.
  • Data-efficient: Adaptation VAE is trained on limited target-domain data; the guidance module requires no extra data beyond the fine-tuning set.
  • Incidental prior preservation: Constraining outputs to remain in a mode of the pre-training latent manifold implicitly prevents catastrophic forgetting of the pre-trained visuomotor prior.

Results

  • RoboTwin 1.0 (17 bimanual simulation tasks): ATE improves average multi-task success rate by up to 9.8% over direct fine-tuning of representative VLAs (RDT-1B, \(\pi_0\)).
  • ManiSkill3 (contact-rich single-arm simulation): Additional evaluation confirms gains in the fine-tuning regime.
  • Real-world cross-embodiment (dual-arm RealMan 7-DoF robot, 5 tasks): ATE achieves a 32% success rate gain over direct fine-tuning in the cross-embodiment setting.
  • Baselines compared: direct fine-tuning of Diffusion Policy, RDT-1B, and \(\pi_0\) without ATE.
  • Generalization experiments (object placement variation, lighting changes, visual distractors) also show improvements, though specific numbers are not included in the provided excerpt.

Limitations

  • Guidance signal \(g\) relies on ground-truth action chunks and is applied only during training, not inference; the mechanism assumes access to paired action labels in the adaptation data.
  • Pre-training VAE must be trained on the full large-scale cross-embodiment dataset before deployment, adding an upfront compute cost.
  • The adaptation VAE approximates \(q_\phi(z)\) as a single Gaussian \(\mathcal{N}(\mu_\phi, \Sigma_\phi)\), which may not faithfully represent a multimodal pre-training latent distribution.
  • Evaluation covers diffusion- and flow-based VLAs; applicability to autoregressive token-based VLAs (e.g., OpenVLA) is explicitly excluded.
  • Real-world results are limited to one robot platform (dual-arm RealMan); broader cross-platform validation is absent.

Relevance to Vision-Language Models

ATE directly addresses a bottleneck in the VLM-to-VLA transfer pipeline: the semantic and instruction-following capabilities inherited from VLMs are underutilized when the low-level action head cannot adapt to a new embodiment without massive data. By decoupling action-space alignment (VAE layer) from policy steering (guidance layer), ATE lets the VLM backbone remain intact while precisely reshaping the action distribution — a clean separation that is relevant to any VLM extended with a generative action head. The classifier-guidance formulation also generalizes the well-established diffusion-guidance paradigm (Dhariwal & Nichol 2021) from image generation into the robotic action domain, connecting VLA adaptation to a broader line of work on guided generative models. For researchers tracking VLMs, this paper is notable as a case study showing that latent-space geometry (mode-seeking vs. mean-seeking KL) is a practically significant design choice when bridging foundation-model priors to task-specific distributions.