Venue: ICLR 2026 Affiliation: China Telecom Institute of AI · Tsinghua · CUHK Shenzhen · Northwestern Polytechnical Category: VLA Training — Adaptation / fine-tuning for diffusion- and flow-based VLAs Trend tag: Cross-embodiment / data efficiency / classifier guidance

flowchart LR
  subgraph S1[Stage 1: Unified Action Latent Space]
    PT[Pre-training action chunks<br/>OXE / DROID / ALOHA / Kuka] --> V1[InfoVAE V_phi<br/>Transformer enc/dec<br/>latent dim 512]
    V1 -. "KL to N(0,I)" .-> Prior["Pretrain latent prior q_phi(z)"]
    Adapt[Adaptation action chunks<br/>RoboTwin / ManiSkill / RealMan] --> V2[InfoVAE V_psi<br/>same architecture]
    V2 -. "reverse KL to q_phi(z)" .-> Mode[Mode-seeking embed<br/>into one mode of q_phi]
  end
  subgraph S2[Stage 2: Classifier Guidance Steers Fine-tuning]
    Mode --> E[Frozen E_psi encoder]
    Noisy[a^k_t:t+h noisy chunk] --> E
    Clean[a^0_t:t+h clean chunk] --> E
    E --> Dist["‖E_psi(a^k) − E_psi(a^0)‖²"]
    Dist --> Grad["g = -∇_a ‖z_hat − z‖²"]
    Grad --> Diff["Diffusion-VLA loss + √(1−ᾱ_k)·λ·g"]
    Grad --> Flow["Flow-VLA loss + ((1−τ)/τ)·λ·g"]
  end
  Diff --> POL[Adapted RDT / DP]
  Flow --> POL2[Adapted π₀]

Pre-trained VLAs (π₀, RDT-1B, Diffusion Policy) face a distribution mismatch when fine-tuned on a new embodiment or task: degrees-of-freedom, end-effector representation, joint torques and chunk length all change between pre-training data (typically single-arm 6-DoF from Open X-Embodiment / DROID) and adaptation data (e.g., dual-arm 7-DoF RealMan with single-DoF grippers). Direct fine-tuning either overfits the small target set or drifts away from the pre-training visuomotor prior. Methods that operate at the parameter level (LoRA, dynamic layer activation) or use kinematics retargeting do not directly close the action distribution gap, and diffusion/flow-based VLAs are especially sensitive because their output is a distribution over action chunks, not a point estimate.

The authors instantiate two InfoVAEs (Zhao et al. 2017), not vanilla VAEs, because vanilla VAEs suffer from posterior collapse and inaccurate amortized inference. Each VAE has a Transformer encoder/decoder. Two learnable tokens μ_token and Σ_token are prepended to the encoder input (motion-VAE style, à la Petrovich et al. 2021), and the decoder uses H zero embeddings as queries with the latent z as key/value via cross-attention.

Pre-training VAE V_phi (encoder E_phi, decoder D_phi):

• Trained on the cross-embodiment pre-training corpus with the standard InfoVAE objective:

  L(φ) = E_q[log p_φ(a|z)] − (1−α)·D_KL(q_φ(z|a) ‖ p(z)) − (α−λ−1)·D_KL(q_φ(z) ‖ p(z))

  where p(z) = N(0, I), and the second KL term is replaced by MMD for tractable optimisation (per Zhao et al. 2017). Uses the same chunk length H as the pre-training VLA.

Adaptation VAE V_psi (encoder E_psi, decoder D_psi):

• Same architecture and variant. Trained on the small adaptation set with a reverse KL to the learned pre-training latent distribution q_φ(z):

  L(ψ) = E_q[log p_ψ(ã|z)] − (1−α)·D_KL(q_ψ(z|ã) ‖ q_φ(z)) − (α−λ−1)·D_KL(q_ψ(z) ‖ q_φ(z))

• q_φ(z) is approximated as N(μ_φ, Σ_φ) with statistics aggregated over all pre-training action latents. The mode-seeking property of reverse KL is the key: it pushes q_ψ into one mode of q_φ, producing a unified, hierarchically structured latent space.

• Uses the adaptation chunk length L (which may differ from H).

t-SNE visualisations in Fig. 15 confirm that InfoVAE produces a compact latent space where adaptation embeddings (blue) sit inside pre-training modes (green); vanilla VAE produces poorly aligned, gapped distributions.

A guidance function is constructed in the unified latent Z. Following Carvalho et al. 2023, the classifier is an energy-based model:

p_ψ(y | â^k_t:t+h) = (1/Z_ψ) · exp(−‖E_ψ(â^k_t:t+h) − E_ψ(a^0_t:t+h)‖²)

with guidance gradient:

g = ∇_â log p_ψ(y|â) ∝ −∇_â ‖E_ψ(â^k) − E_ψ(a^0)‖²

This gradient is folded into the training objective of the underlying generative head:

• Diffusion-based VLAs (RDT, DP): L(θ) = E[‖ε − ε_θ(a^k; k, o_t, l) + √(1 − ᾱ_k)·λ·g‖²]
• Flow-based VLAs (π₀): L(θ) = E[‖v_θ(a^τ; τ, o_t, l) + ((1−τ)/τ)·λ·g − (a^0 − ε)‖²]

The encoder E_psi is frozen during this stage; it serves purely as a distance metric. Crucially, since guidance enters at training time only, inference cost is zero — unlike test-time gradient-based steering (DynaGuide) that runs iterative backprop per step. The recommended guidance scale is λ = 2 (Table 13).

Real-robot fine-tuning runs 120k steps on 8× A100 with batch 6/GPU at 20 Hz inference on a single 4090.

Sample-efficiency: RDT + ATE surpasses RDT baseline at 70k steps vs. 90k for the baseline. ATE on π₀ at 25 demos/task averages 29.0% vs. baseline 9.2% (3.15× gain, Table 12).

Four long-horizon tasks at 60k / 90k / 120k steps. Tasks include Cook Bun, Pick Bun, Make Sandwich, Use Toaster. Reported in Fig. 3:

• Cook Bun: π₀ ~15% vs. ATE 100% at 90k steps.
• Pick Bun (single-arm): ATE 50% @ 90k → 70% @ 120k.
• Overall average at 120k: π₀ 16.7% vs. ATE 58.1%.

Make Yogurt Bowl (multi-tool dual-arm, 80 demos): π₀ 15% → ATE 25% (Table 2). Under visual distractors: π₀ 0% → ATE 20% (Table 3).

1. Two-step Info-VAE vs. single-step (Fig. 11): Single-step performs roughly on par with direct fine-tuning; two-step delivers the full gain. Confirms that alignment to the pre-training prior — not just having a latent at all — is the load-bearing piece.
2. Guidance scale λ (Table 13): λ = 1 / 2 / 3 swept on four tasks. λ = 2 is the sweet spot; λ = 3 occasionally hurts because too-strong guidance disrupts chunk smoothness.
3. Reverse-KL weight α (Table 14): α ∈ {0.1, 1, 10}. α = 1 is best; α = 10 over-constrains the adaptation VAE (e.g., Block Handover 92% → 37%); α = 0.1 weakens alignment.
4. InfoVAE vs. Vanilla VAE (Tables 15-16): Vanilla VAE already beats baseline (showing the steering module is modular), but InfoVAE consistently adds another 4-20 points (e.g., DP on Empty Cup: baseline 22 → vanilla 34 → InfoVAE 37; π₀ on Blocks Stack Easy: 30 → 31 → 50).
5. Sample efficiency (Table 12): at 25 demos/task, π₀ averages 9.2%, ATE 29.0% (3.15×). With 25 demos ATE matches/exceeds π₀ with 50 demos on multiple tasks — roughly 2× sample efficiency.
6. DynaGuide comparison (Table 18): ATE matches or exceeds the test-time guidance baseline, without inference overhead (DynaGuide requires iterative gradient steps at every denoising step).

The paper does not have a dedicated "Limitations" section, but the experiments reveal:

• A handful of tasks (Bottle Adjust, Tool Adjust, Pick Apple Messy, Shoe Place) regress with ATE — guidance can occasionally fight the pre-training prior on contact-rich tasks the prior covers well.
• α and λ must be tuned (the sensitivity studies show a 1-order-of-magnitude window).
• Step-1 VAE training itself takes ~12 h on the pre-training corpus per backbone, though this is amortised across all downstream adaptation tasks.
• Currently requires access to the pre-training action distribution — not just the VLA checkpoint — to train V_φ. For closed checkpoints this would block direct application.

ATE regularises adaptation in action-distribution space rather than weight space. Compared to:

• Robust Parameter Merging — weight-space interpolation; complementary.
• UniVLA / GR00T N1 — shared latent action space (VQ-VAE / discrete latents), used as an intermediate representation. ATE instead learns a probability prior over actions and uses it as a classifier-guidance signal.
• Latent Action Diffusion (Bauer et al. 2025) — separate VAEs + contrastive learning; ATE replaces contrastive alignment with the more principled reverse-KL mode-seeking mechanism.
• π₀ FAST / OpenVLA fine-tuning — depend on data scale; ATE is most effective in the low-data, cross-embodiment regime where direct fine-tuning collapses.
• Test-time steering (DynaGuide): ATE bakes guidance into training, paying zero inference cost (vs. DynaGuide's iterative gradient steps), echoing the inference-efficiency arguments in SP-VLA and Action-aware Dynamic Pruning.

Conceptually closest in spirit to classifier-guided diffusion (Dhariwal & Nichol 2021) and motion-planning diffusion (Carvalho et al. 2023), repurposed for the VLA adaptation problem. Plug-and-play across DP, RDT-1B and π₀ is the central practical claim.

• OpenReview: https://openreview.net/forum?id=T3i7Ifeatk
• Code: https://github.com/TeleHuman/Align-Then-Steer

• Robust Parameter Merging — weight-space alternative
• UniVLA — task-centric latent action space
• X-VLA
• π0.6
• Survey: VLA & Manipulation

← Back to ICLR-2026