APT: Action Expert Pretraining Improves Instruction Generalization of Vision-Language-Action Policies

Abstract: Vision-Language-Action (VLA) models that couple pretrained Vision-LLMs (VLMs) with continuous action experts have achieved strong manipulation performance, yet generalization to out-of-distribution (OOD) language instructions remains poor. A known challenge is the structural imbalance in VLA data, where language is far less diverse than visual and action content, making policies prone to visual shortcuts. While discrete-action methods mitigate this through vision-language co-training, continuous action experts lack such protection: they start from random initialization and learn entirely from imbalanced data, producing noisy gradients that corrupt the VLM and fail to exploit its language capability. We address this from a Bayesian perspective, factorizing the policy into a language-agnostic Vision-Action (VA) prior and a language-conditioned VLA likelihood, and propose APT, a two-stage training method emphasizing Action expert PreTraining. In Stage 1, the action expert is pretrained as a VA prior on vision-action pairs from a frozen VLM, bypassing the language imbalance. In Stage 2, language tokens are injected through a gated fusion mechanism that integrates VLM features while preserving the learned visuomotor prior. APT applies to mainstream VLA architectures, including the $π$ and GR00T-style architectures. Comprehensive experiments validate that APT achieves consistent gains on unseen instructions and compositional tasks. Project Page: https://xukechun.github.io/papers/APT/

• The paper introduces APT, a two-stage methodology that decouples visuomotor priors from language alignment to improve instruction generalization.
• It leverages Bayesian factorization and gated language token fusion to mitigate dataset imbalances and overcome vision-only shortcut learning.
• Experimental results show up to 62% success in challenging OOD tasks, demonstrating robust improvements in both simulated and real-world robotic manipulation.

Action Expert Pretraining for Improved Instruction Generalization in Vision-Language-Action Policies

Introduction and Motivation

Recent advances in Vision-Language-Action (VLA) models have demonstrated that robotic manipulation policies can benefit from integrating pretrained Vision-LLMs (VLMs) with generative action experts to translate multimodal instructions into executable trajectories. However, these models exhibit persistent deficiencies in generalization to out-of-distribution (OOD) language instructions—especially when language prompts deviate from those seen during training or compose novel demands. This limitation arises from a structural modality imbalance in VLA datasets, where visual and action diversity far outpace language variation, driving continuous-action policies to over-exploit visual shortcuts and ignore language context, ultimately corrupting VLM representations when action experts are trained from random initialization.

The presented work, "APT: Action Expert Pretraining Improves Instruction Generalization of Vision-Language-Action Policies" (2606.12366), directly targets this generalization problem by introducing Action expert PreTraining (APT)—a two-stage training procedure, grounded in a Bayesian factorization of the VLA policy, which decouples the learning of visuomotor priors from language-conditioned task alignment. This methodology enables robust transfer and composition of language instructions across both simulated and real-world manipulation settings.

Figure 1: APT comprises a two-stage protocol in which a vision-action prior is first learned from visual tokens, followed by language-conditioned likelihood alignment via token injection from a frozen VLM backbone.

Bayesian Decomposition and Training Protocol

The core technical insight is the Bayesian decomposition of the policy $\pi(\mathbf{a} | \mathbf{v}, \ell)$ into a language-agnostic Vision-Action prior $\pi^p(\mathbf{a} | \mathbf{v})$ and a VLA likelihood $\mathcal{L}(\ell | \mathbf{v}, \mathbf{a})$, i.e.:

$$\pi(\mathbf{a} | \mathbf{v}, \ell) \propto \pi^p(\mathbf{a} | \mathbf{v}) \cdot \mathcal{L}(\ell | \mathbf{v}, \mathbf{a})$$

Here, $\mathbf{a}$ is the action token, $\mathbf{v}$ is the visual input, and $\ell$ the language prompt. Existing VLA policies, trained end-to-end on triplets $(\mathbf{a}, \mathbf{v}, \ell)$, face an intrinsic incentive to disregard $\ell$ due to visual dominance in the data, leading to degenerate minima where actions are predicted from vision alone. Information-theoretic analysis demonstrates that, under typical dataset imbalance, the conditional mutual information $I(\mathbf{a};\ell|\mathbf{v})$ is negligible, so standard gradient descent is biased toward vision-only solutions.

APT circumvents this by isolating action learning from language during Stage 1: the action expert (a diffusion or flow-based model) is pretrained on balanced vision-action pairs with the language signal masked, ensuring learned priors are not tainted by overfit to visual context alone. Only in Stage 2 are language tokens injected—via a gated fusion mechanism—enabling the network to align the action manifold with instructions without erasing previously acquired visuomotor regularities.

Architectural Mechanisms

APT introduces a specialized action expert architecture featuring:

• Multimodal Self-Attention: Action, vision, and (eventually) language tokens are concatenated and processed with block-wise causal attention.
• Layer-wise VLM Feature Gated Fusion: At each action expert layer, intermediate features from the VLM are injected, gated by learnable scalars, providing a spectrum of abstraction from localized scene details to global semantics.
• Two-Stage Masked Attention: During Stage 1, only $\pi^p(\mathbf{a} | \mathbf{v})$0 layers attend to vision/action, with language masked. In Stage 2, the full $\pi^p(\mathbf{a} | \mathbf{v})$1 layers (including new cross-modal attention) are unlocked, reasoning jointly over all tokens.

This mechanism is agnostic to the specific VLA architecture, as demonstrated by successful applications to both $\pi^p(\mathbf{a} | \mathbf{v})$2-style (interleaved attention) and GR00T-style (final-layer cross-attention) policies. Empirical evidence confirms that the approach generalizes to diverse tokens and embodiment representations.

Figure 2: The two-stage pretraining scheme can be applied to various VLA-hybrid architectures, showing consistent gains in OOD instruction generalization.

Experimental Results

Simulation Benchmarks

APT is evaluated on challenging simulation benchmarks such as LIBERO-PRO and compositional pick-and-place tasks with multiple OOD axes—unseen objects, containers, environmental backgrounds, and paraphrased or novel task instructions. Baselines include OpenVLA, $\pi^p(\mathbf{a} | \mathbf{v})$3, $\pi^p(\mathbf{a} | \mathbf{v})$4 (with gradient stop on VLM), LangForce (MI-based enforcement), and modular programmatic controllers like CaP-X.

Key findings:

• Baseline joint training collapses to 0% success on OOD instruction tasks; both $\pi^p(\mathbf{a} | \mathbf{v})$5 (with VL data co-training and knowledge insulation) and LangForce show trade-offs between visuomotor robustness and language grounding.
• APT consistently surpasses both, achieving up to 62% success rates on LIBERO-PRO OOD tasks, even in the most challenging Goal-Task settings.
• Enabling joint VLM finetuning post-pretraining provides further gains, indicating that, with a well-initialized action prior, allowing gradients to the VLM is not detrimental.

Ablations

Ablative studies demonstrate:

• Two-stage pretraining is effective with or without VL knowledge insulation; the initialization effect dominates.
• Gated fusion for language token injection preserves VA priors better than naive token insertion.
• The benefits of APT persist even with limited pretraining, but are amplified with large-scale diverse demonstrations.

Real-World Transfer

APT is deployed on a real-robot platform and evaluated on single-task generalization (pick-place and cluttered manipulation) and compositional instruction following (task coaching and task chaining). Comparative analysis against $\pi^p(\mathbf{a} | \mathbf{v})$6 highlights:

• Marginal improvement for seen objects (SO), large performance gaps for UO, UOUC, UOUCUE, and long-horizon compositional tasks, with success rates 50–100% higher in OOD settings.
• Robustness to distractors and accurate multi-step execution remain markedly better for APT, especially when evaluating instruction chaining, where baseline policies typically fail to parse and correctly sequence subtasks.

Figure 3: APT achieves strong generalization in both atomic and compositional real-world manipulation tasks, outperforming baselines across OOD settings.

Analysis and Discussion

APT's pretraining regime directly resolves the visual shortcut tendency present in continuous VLA models. By ensuring that the action expert first models the true, balanced distribution over visuomotor behaviors, subsequent language alignment can occur without corrupting either the prior or the VLM backbone, thereby preserving compositionality and instruction fidelity. The approach is fundamentally architectural and methodological; it does not depend on additional VL data or handcrafted auxiliary loss functions.

On the theoretical front, the work justifies the two-stage conditioning via precise information-theoretic arguments, showing that shortcut optima are intrinsic to the standard VLA training objective in the presence of statistical imbalance, and that APT reintroduces a training signal that incentivizes true language grounding.

Limitations and Future Work

The current formulation does not explicitly model persistent long-horizon memory, so transition errors (e.g., continued pushing after grasping, missed sub-task boundaries in compositions) occur primarily when sub-task termination signals are ambiguous. The focus is restricted to manipulation tasks; generalization to mobile manipulation or tasks demanding global scene understanding and extended planning is an area for future extension. Moreover, improvements in hierarchical or memory-augmented architectures could further boost multi-step task compositionality.

Conclusion

This paper demonstrates that action expert pretraining, structured via Bayesian policy factorization, is an effective and architecture-agnostic approach to significantly enhance OOD instruction generalization in continuous VLA policies. This is evidenced by robust gains in both simulated and real-world settings, strong architectural ablation, and resilience to dataset imbalance and OOD shifts. These results indicate a substantial step toward more reliable, compositional robotic policies, and motivate exploration of advanced memory mechanisms and broader task coverage in VLA frameworks.