BayesianVLA: Bayesian Decomposition of Vision Language Action Models via Latent Action Queries

TL;DR

This work identifies a fundamental weakness in Vision-Language-Action models: when training data ties visuals too strongly to instructions, models collapse to vision-only policies and fail in out-of-distribution settings. To combat this, BayesianVLA introduces a Bayesian decomposition of the action policy into a vision-only prior $p(a\mid v)$ and a language-conditioned posterior $\pi(a\mid v, \ell)$, together with Latent Action Queries and a dual-branch training objective that maximizes the conditional PMI between actions and instructions via the Log-Likelihood Ratio $\mathcal{L}_{\text{LLR}} = \log p(\ell \mid a, v) - \log p(\ell \mid v)$. The method yields substantial gains on SimplerEnv (66.5% average, +11.3% absolute over baselines) and RoboCasa (50.4% average), demonstrating stronger grounding of language in action, especially in OOD scenarios. By enforcing instruction following through information-theoretic regularization and preserving general VLM capabilities, BayesianVLA achieves robust generalization without additional data and offers a principled path toward reliable, language-grounded robotic manipulation.

Abstract

Vision-Language-Action (VLA) models have shown promise in robot manipulation but often struggle to generalize to new instructions or complex multi-task scenarios. We identify a critical pathology in current training paradigms where goal-driven data collection creates a dataset bias. In such datasets, language instructions are highly predictable from visual observations alone, causing the conditional mutual information between instructions and actions to vanish, a phenomenon we term Information Collapse. Consequently, models degenerate into vision-only policies that ignore language constraints and fail in out-of-distribution (OOD) settings. To address this, we propose BayesianVLA, a novel framework that enforces instruction following via Bayesian decomposition. By introducing learnable Latent Action Queries, we construct a dual-branch architecture to estimate both a vision-only prior $p(a \mid v)$ and a language-conditioned posterior $π(a \mid v, \ell)$. We then optimize the policy to maximize the conditional Pointwise Mutual Information (PMI) between actions and instructions. This objective effectively penalizes the vision shortcut and rewards actions that explicitly explain the language command. Without requiring new data, BayesianVLA significantly improves generalization. Extensive experiments across on SimplerEnv and RoboCasa demonstrate substantial gains, including an 11.3% improvement on the challenging OOD SimplerEnv benchmark, validating the ability of our approach to robustly ground language in action.

Figures (5)

• Figure 1: Examples of the vision shortcut in RoboCasa RoboCasa_2024_RSS. Training data exhibits visual diversity but limited task diversity. As a result, the model learns to execute tasks directly based on specific visual cues rather than relying on language instructions.
• Figure 2: (a) In LIBERO Goal LIBERO_2023_NeurIPS, the same scene corresponds to multiple tasks, revealing the ambiguity that vision-only models fail to resolve. (b) Action loss curves on BridgeDataV2 Bridgedatav2_2023_CoRL show that the vision-only model achieves comparable training loss to the full vision-language model, indicating the presence of visual shortcuts even in diverse, in-the-wild datasets.
• Figure 3: The framework of BayesianVLA. The framework employs a dual-branch architecture with shared VLM weights. The Priori Branch (left) processes $[v, \mathcal{Q}, \ell]$ with causal masking to learn the vision-only prior $p(a \mid v)$. The Posteriori Branch (right) processes $[v, \ell, \mathcal{Q}]$ to learn the full policy $\pi(a \mid v, \ell)$. Latent Action Queries $\mathcal{Q}$ serve as a bottleneck interface, and the LLR objective (in Eq. \ref{['eq:llr']}) encourages the model to maximize the information between actions and instructions. At inference, only the Posteriori Branch is used, incurring no additional computational overhead.
• Figure 4: Qualitative comparison of general multimodal reasoning. We present a case where the model is asked to solve a mathematical problem. The standard VLA baseline (QwenGR00T) suffers from catastrophic forgetting; while the text before the comma implies "differentiating all terms together", the subsequent output degenerates into repetitive and meaningless gibberish (bottom right). In contrast, BayesianVLA (top right) retains the VLM's original reasoning and language generation capabilities (left), successfully solving the problem.
• Figure 5: Additional qualitative comparison. Demonstrating the preservation of general VLM capabilities on another example.