Diffusion-Trained Multimodal Backbones for Action-Token Representations in Flow-Matching Vision-Language-Action Policies

1Abstract

Vision-Language-Action (VLA) models have emerged as a promising framework for generalist robot control by pairing pretrained vision-language representations with continuous action-generation modules. Recent systems such as π0 and π0.5 demonstrate that pretrained VLM-style backbones can be adapted to robotic manipulation through robotics-specific state/action tokens and flow-matching action experts for continuous action chunks [2,3]. However, many current VLA backbones inherit a causal autoregressive language-model prior, where sequence processing is shaped by next-token prediction.

Robot action chunks are not naturally left-to-right text sequences. A future action trajectory is a structured block in which early approach actions, future contact timing, multi-stage subtask ordering, grasp closure, and stabilization are jointly coupled. This proposal investigates whether a diffusion-trained bidirectional multimodal backbone can produce better action-token hidden states than a causal/autoregressive VLM backbone when both use the same flow-matching action expert.

The proposed study isolates the backbone as the key variable while keeping the action representation, flow objective, action head, dataset, task suite, and control stack fixed. Evaluation is carried out entirely in simulation on the LIBERO benchmark, which enables fair, reproducible, high-n comparison, removes real-robot calibration and latency confounds, provides ground-truth state for representation probing, and is directly comparable to published diffusion-VLA results. The core hypothesis is that diffusion-style block refinement and bidirectional token context will yield hidden states that better encode trajectory phase, subtask ordering, and corrective structure for flow-matching robot control.

2Motivation

Current VLA models are built around a powerful idea: use large-scale vision-language pretraining to obtain semantic and perceptual grounding, then adapt the model to continuous robot control. In the π0 family this is achieved by integrating robot state/action inputs into a VLM and using conditional flow matching to generate continuous action chunks [2]. π0.5 further extends this direction with open-world generalization and a flow-matching action expert that operates on action tokens representing partially denoised robot actions [3,4].

This design raises a deeper architectural question, and it is important to be precise about which part of the model is causal. In the OpenPI implementation of π0.5, the attention mask is block-wise rather than strictly left-to-right: the image/language prefix attends bidirectionally within itself, and the action chunk forms a single bidirectional block whose tokens all attend to one another. Causality survives only across blocks (prefix → state → action), so the backbone behaves as a prefix-LM at inference time, not as a token-level autoregressive decoder. What remains autoregressive is therefore not the action-token attention mask but the pretraining prior: the weights were trained under a next-token objective that shapes how the model organizes and refines sequences. This proposal asks whether that pretraining prior — rather than the inference-time mask, which is already bidirectional over the action block — limits the quality of action-token hidden states, and whether a backbone pretrained around block-level refinement would organize them better.

Diffusion-trained language and multimodal models provide a useful alternative. DiffusionGemma, for example, is a Google DeepMind model based on a Gemma 4 Mixture-of-Experts architecture that generates tokens using discrete diffusion rather than standard autoregressive decoding [5,6]. Rather than emitting text strictly token by token, a diffusion language model refines corrupted token blocks over multiple steps. This motivates a robotics hypothesis: action chunks may benefit from hidden states produced by a backbone trained around bidirectional block refinement, even if the final action decoder remains a continuous flow-matching module.

3Background and Prior Evidence

3.1Flow and diffusion policies for robot action generation

Diffusion Policy introduced a visuomotor policy formulation in which robot behavior is generated through a conditional denoising diffusion process over action space [14]. The method showed strong performance across manipulation benchmarks and emphasized several advantages relevant to this proposal: handling multimodal action distributions, supporting high-dimensional action spaces, using receding-horizon control, and incorporating a time-series diffusion transformer. Subsequent work such as 3D Diffusion Policy extends this formulation to compact 3D scene representations [9], indicating that the diffusion-policy approach generalizes across observation modalities. The important lesson is that robot action generation is often better treated as iterative trajectory refinement than as one-step regression.

3.2VLM backbones with flow-matching action experts

The π0 model adapts a pretrained VLM for robot control by adding robotics-specific inputs and outputs and using conditional flow matching to model continuous actions [2]. π0.5 builds on this architecture and uses co-training, semantic subtask prediction, and multiple data sources to enable broader generalization [3]. The open-source OpenPI implementation supports the flow-matching head for π0.5 training and inference [4]. This validates the architectural pattern used here: multimodal tokens → multimodal backbone → action-token hidden states → flow-matching action expert → continuous action chunk. The proposed research does not replace flow matching; it asks whether the backbone that feeds the flow expert should be diffusion-trained and bidirectional rather than causal/autoregressive.

3.3Diffusion-trained multimodal backbones

DiffusionGemma is an experimental open model built by Google DeepMind using discrete diffusion over tokens, based on the Gemma 4 26B / 4B-active Mixture-of-Experts architecture [5]. The model card reports 25.2B total parameters, 3.8B active parameters, 30 layers, text and image support, a 256-token canvas, and a 256K-token context length [6]. For robotics, the key point is not that DiffusionGemma should directly output actions through its native text head; the relevant point is that diffusion pretraining encourages a different representation prior — corrupted blocks are refined into coherent outputs using bidirectional context. We flag that this is an analogy, not established transfer: whether a block-refinement prior learned over language and image tokens transfers to the structure of robot action chunks is precisely the empirical question this proposal tests, and the design is built so that a null transfer result is itself informative (§12).

3.4Representation quality matters for diffusion policies

Crossway Diffusion improves diffusion-based visuomotor policies by adding self-supervised representation learning objectives, including reconstructing visual and state information from intermediate representations of the reverse diffusion process [7]. This supports the claim that diffusion-policy performance depends not only on the denoising loss but on the quality and structure of intermediate representations. This proposal extends the idea to VLA backbones: instead of only changing the action denoising module, it investigates whether the multimodal backbone can produce action-token hidden states that better organize task-relevant information.

3.5Action-token structure is now central to VLA design

Recent VLA work increasingly treats action tokens as a core interface between perception-language context and executable control. Ordered Action Tokens identifies desiderata for action tokenization and validates action tokens in both autoregressive control and token co-training settings where token losses shape the VLM context consumed by a flow-based action expert [8]. This directly supports the proposed focus on action-token hidden states as a meaningful research object rather than an implementation detail.

3.6Diffusion vision-language-action models (closest prior work)

A cluster of concurrent work already couples diffusion-trained multimodal backbones with robot control. Dream-VLA builds a VLA directly on the diffusion VLM Dream-VL (Qwen2ViT vision encoder + Dream-7B backbone) and reports strong results (e.g. 97.2% average on LIBERO), pretraining on Open X-Embodiment with a discrete-diffusion loss and optionally switching to continuous losses such as flow matching or L1 during downstream fine-tuning [13]. Related efforts include LLaDA-VLA [15] and Discrete Diffusion VLA [16], which bring discrete diffusion to action decoding.

The critical architectural distinction is that these systems fold action generation into the unified diffusion backbone — there is no separate action expert. This proposal makes the opposite choice on purpose: it keeps a separate, fixed flow-matching action expert in the π0.5 style and swaps only the backbone. That separation is what makes the controlled comparison possible. Concretely, this proposal differs in three ways:

1. Controlled isolation. Dream-VLA shows a diffusion VLA works but does not isolate why — it has no scale/architecture-matched AR control, so its gains cannot be attributed to the diffusion prior rather than scale, data, or the action-decoding change. The C−D (Dream-7B vs. Qwen2.5-7B) contrast is designed precisely to answer that.
2. Fixed action mechanism. Holding a single flow-matching expert fixed across all backbones attributes differences to the backbone's representations, not to a change in how actions are decoded.
3. Representation probing. This proposal probes action-token hidden states directly (§9.2), which the outcome-only evaluations of existing diffusion VLAs do not.

In short, the diffusion-VLA cluster establishes feasibility and strong headline numbers; this proposal supplies the controlled, representation-level explanation those results leave open.

4Research Question and Hypothesis

The hypothesis is that a diffusion-trained bidirectional backbone will outperform a causal/autoregressive backbone when paired with the same flow-matching expert, especially on LIBERO tasks requiring spatial reasoning, instruction/goal grounding, and long-horizon multi-stage execution. More specifically, the diffusion-trained backbone is expected to produce action-token hidden states that better encode trajectory phase and subtask ordering, future contact timing, target-object identity and spatial relations, goal/instruction grounding, corrective motion structure, and final action-chunk success. The expected advantage is strongest on LIBERO-Long and in low-data regimes, and weaker on the shorter LIBERO-Spatial/Object/Goal suites, where near-ceiling success rates can compress differences.

5Proposed Method

5.1System overview

The architecture follows a π0.5-style integration pattern. The model receives image, language, proprioceptive, flow-time, and noisy action tokens; the backbone processes the multimodal token sequence, and the hidden states at the action-token positions are passed into a shared flow-matching action expert. For the LIBERO single-arm end-effector policy the future action chunk is

where H is the action horizon and D the action dimension. LIBERO uses a 7-DoF end-effector action space:

The design is action-dimension agnostic; the same pipeline extends to higher-dimensional (e.g. bimanual) action spaces without changing the backbone comparison. Each noisy action vector is projected into the model hidden dimension, and the full token sequence is [image tokens, language tokens, proprio tokens, flow-time token, action₁, …, action_H]. The backbone outputs hidden states H_action ∈ ℝ^H×d at the action-token positions, which the expert maps to a vector field v_pred = f_expert(H_action) ∈ ℝ^H×D.

5.2Reading hidden states from a diffusion-trained backbone

For an autoregressive or prefix-LM backbone, H_action is produced by a single forward pass. A discrete-diffusion backbone does not share this property: its hidden states are a function of the corruption level (mask ratio) and the number of refinement steps, so "the" action-token hidden state is really a family indexed by diffusion time. This choice must be fixed before any comparison and reported explicitly. We adopt the following convention as a controlled hyperparameter: for the diffusion backbone (Model C), the noisy action tokens are placed on the diffusion canvas and the backbone is run for a fixed, pre-registered number of refinement steps k; the hidden states read into the expert are taken from the final refinement step. We ablate k ∈ {1, 2, 4} so that k=1 (a single forward pass) is directly compute-comparable to the AR/prefix-LM backbones, isolating the effect of the prior from extra test-time computation. The flow-time conditioning s is kept separate from the backbone's internal diffusion step.

5.3Flow-matching objective

Given a clean action chunk A_clean from demonstration data, sample Gaussian noise A_noise ∼ N(0, I) and a flow time s ∼ Uniform(0, 1). Construct the interpolated noisy chunk and target field:

The model predicts v_pred = f_θ(images, language, proprio, A_s, s) and is trained with

At inference the policy starts from random noise and repeatedly applies the predicted vector field to refine the action chunk; the first few actions are executed, then the robot re-observes and replans in a receding-horizon loop.

5.4Backbone selection and feasibility

Backbones are chosen so the identifying contrast is clean, the models fit available hardware, and all components are openly licensed with released training code. The primary study uses the Dream-7B / Qwen2.5-7B pair at ~7B parameters; DiffusionGemma is retained only as an optional large-scale confirmation.

Feasibility. At ~7B parameters the primary models fit comfortably on a multi-A100 (80 GB) setup. Full fine-tuning of a 7B backbone plus the action expert is tractable with FSDP; LoRA is the default to preserve the pretraining difference between C and D. At the k=1 hidden-state setting the diffusion backbone costs one forward pass, matching the AR backbone for latency parity. The optional DiffusionGemma run is heavier (26B weights, ~52 GB bf16, plus encoder-decoder MoE), which is why it is scoped as a confirmation run.

6Model Variants

The study keeps the action expert, action space, flow objective, dataset, and control stack fixed while varying the backbone and its attention topology. Dream-7B, a discrete-diffusion LM, is initialized from Qwen2.5-7B and trained with a diffusion objective; using Dream-7B for the diffusion arm and Qwen2.5-7B for the AR arm holds parameters, architecture, and initialization fixed, so the dominant remaining difference is the training objective itself.

The comparisons isolate: C−D — the diffusion pretraining prior at matched scale/architecture; B−D — attention-mask topology (pretraining held fixed); C−A — the total effect of the proposed backbone vs. the deployed π0.5 baseline. The identifying claim rests on C−D. Because Dream-7B and Qwen2.5-7B share initialization and architecture and differ chiefly in objective, this is close to the ideal "same base, two objectives" comparison, obtained without training a diffusion model from scratch. We do not claim any single contrast is perfectly clean; the set brackets the effect and makes the remaining confounds explicit.

7Why This Is Not Redundant

The proposed architecture does not run two independent denoising processes. The diffusion-trained backbone is used as a multimodal representation processor; the continuous flow-matching action expert remains the only component responsible for denoising robot action chunks. The system is a diffusion-trained multimodal backbone + continuous flow-matching action expert — not a discrete text-diffusion action decoder stacked on a separate continuous action-flow decoder. It is also distinct from unified diffusion VLAs such as Dream-VLA (§3.6), which decode actions inside the backbone with no separate expert. Keeping the flow-matching expert fixed and external is the mechanism that makes the backbone the sole controlled variable.

8Experimental Task Suite

Evaluation is performed entirely in simulation on the LIBERO benchmark [17], a standard single-arm (Franka Panda) manipulation suite. LIBERO is chosen for fairness and comparability: fixed task definitions and deterministic resets, large numbers of rollouts at negligible cost, ground-truth simulator state for representation probing, and it is the benchmark on which the closest prior work (Dream-VLA, OpenVLA-OFT, GR00T) reports results, enabling an apples-to-apples comparison. The study uses the four standard suites, each isolating a different generalization axis:

1. LIBERO-Spatial. Same objects, varied spatial arrangements; isolates spatial reasoning and target disambiguation.
2. LIBERO-Object. Same layout, varied objects; isolates object identity grounding.
3. LIBERO-Goal. Same objects and layout, varied goals/instructions; isolates instruction/goal grounding.
4. LIBERO-Long (LIBERO-10). Long-horizon, multi-stage tasks; isolates trajectory-phase and subtask-ordering structure, where the block-refinement hypothesis is expected to matter most.

Because published diffusion-VLA success rates on the shorter suites are near ceiling, LIBERO-Long and low-data (few-shot) variants are the primary discriminating conditions, complemented by continuous and representation-level metrics that do not saturate.

9Evaluation Metrics

9.1Policy-level metrics

Primary rollout metrics include per-suite and average task success rate, LIBERO-Long success specifically, few-shot success under reduced demonstration budgets, final end-effector pose error, target-object placement error, contact-timing error (from simulator contact events), action smoothness and jerk, and inference latency at the chosen k. Because effect sizes may be modest, all headline comparisons report Wilson confidence intervals for success rates and bootstrap intervals for continuous metrics, with significance assessed by paired tests over the shared, deterministic initial-condition set (e.g. 50 episodes × 10 tasks = 500 rollouts per suite, ≥3 training seeds per model). Continuous metrics are emphasized because they are more sample-efficient than binary success and can reveal differences before success rates separate.

9.2Representation-level metrics

The core scientific contribution is representation analysis. For each trained model, hidden states at the action-token positions are extracted and lightweight probes are trained to predict task-relevant quantities against simulator ground truth: final end-effector pose, contact timestep, distance to contact, grasp success/failure, subtask phase/stage index, target-object pose and identity, spatial relation to the goal region, and whether the action chunk will succeed. If the diffusion-trained backbone produces hidden states more predictive of these quantities, it supports the hypothesis that the backbone learns better action-token representations.

Probe fairness. Linear-probe predictivity scales with hidden dimensionality and pretraining scale, so probes are controlled for capacity: hidden states are projected to a common dimensionality (or compared under a matched probe parameter budget), probe hyperparameters and training data are identical across backbones, and performance is reported against a shared held-out split. The most informative probe comparison is again C vs. D; C vs. A gaps are reported only alongside the dimensionality caveat.

10Training Plan

1. Stage 1 — Baseline reproduction. Fine-tune the stock π0.5 prefix-LM baseline (Model A) with the flow expert on LIBERO via OpenPI, reproducing published π0/π0.5-class numbers before any backbone changes.
2. Stage 2 — Backbone swap. Replace the backbone with Dream-7B (Model C) under the shared recipe, keeping the expert, action space, flow objective, and control stack unchanged.
3. Stage 3 — Attention-topology ablation (Model B). Starting from Model D's AR backbone, expand the topology so prefix and action tokens attend mutually. Building B on D holds pretraining fixed so B−D isolates topology alone; a near-null result is a legitimate finding that the mask is not the bottleneck.
4. Stage 4 — Scale/architecture-matched control (Model D). Train Qwen2.5-7B under the same recipe and expert as C. Because Dream-7B is initialized from Qwen2.5-7B, this is close to the ideal "same base, two objectives" comparison. The identifying contrast C−D is computed here.
5. Stage 5 — LoRA or partial fine-tuning. Frozen backbones may not adapt enough; full fine-tuning may erase the pretraining differences. LoRA provides a practical middle ground and is the default.
6. Stage 6 — LIBERO evaluation. Evaluate all models under identical conditions with full standard rollout counts, seeds, and confidence-interval reporting, reporting LIBERO-Long and low-data conditions as the primary discriminators.
7. Stage 7 — Representation probing. Freeze trained policies and analyze action-token hidden states with linear probes, clustering, temporal metrics, and attention visualization.

11Compute and Timeline

The plan assumes a single node of 4× A100 80 GB with FSDP as the reference. At ~7B parameters all primary models (A, B, C, D) plus the flow expert fit comfortably; LoRA runs on a single A100 and full fine-tuning fits across the node. Estimates are anchored to a measured full-fine-tuning throughput of 1.17 s/step: at ~30k steps this is ~9.75 wall-clock hours (~39 A100-hours per ~7B run). On H100/H200 the same run is expected at roughly 1–2 optimization steps/s, dropping the full primary study to roughly 3–5 days of compute.

12Expected Results

The diffusion-trained backbone is expected to improve performance most clearly on LIBERO-Long, low-data/few-shot regimes, contact-timing prediction, action smoothness, and global trajectory consistency across a chunk. A strong positive result would show higher success on LIBERO-Long and low-data conditions, smoother chunks, and more predictive hidden states, with the effect surviving in the matched C−D contrast rather than only in the confounded C−A contrast. A useful negative result would show that flow-matching experts dominate performance and that the backbone prior matters less than action representation and normalization. A third, particularly informative outcome is a dissociation: more predictive probes but no rollout gain (representation quality does not bottleneck control), or a rollout gain with no probe difference (the benefit is not captured by the probed quantities). All three cases are reported explicitly.

13Risks and Mitigations

14Expected Contributions

1. A controlled architecture for combining diffusion-trained multimodal backbones with flow-matching VLA action experts.
2. A systematic, scale/architecture-matched comparison between AR and diffusion-trained backbones under the same action decoder, with the diffusion prior identified through a matched control (C−D) rather than a confounded raw comparison.
3. A precise characterization of the π0.5 backbone as a prefix-LM (bidirectional action block), correcting the common "causal VLA backbone" framing and separating attention-mask topology from the pretraining prior.
4. Evidence on whether backbone generative prior affects action-token hidden representations for continuous robot control.
5. Representation-level analysis of action-token hidden states in VLA policies.
6. Controlled evaluation on LIBERO, directly comparable to published diffusion-VLA results.

15Minimal Viable Study

A minimal viable study can be completed on a single LIBERO suite and should be run first. To stay interpretable at minimal scope, the MVP includes the matched control rather than only a raw diffusion-vs-stock pair: one suite (LIBERO-Long preferred, least susceptible to ceiling effects); Model A (stock π0.5 prefix-LM baseline); Model C (Dream-7B under the shared recipe); Model D (Qwen2.5-7B, Dream's AR base, identical recipe) so that C−D is identifiable; the same flow expert, dataset, and action horizon across all three; multiple seeds with full standard rollout counts and confidence intervals; and representation probing with matched probe capacity using simulator ground-truth labels. Because both models are ~7B and openly available, Model D is affordable within the MVP, so the identifying C−D contrast is available from the outset.

16Final Positioning

The proposed work is not simply "using DiffusionGemma for robotics." The deeper contribution is to test whether diffusion-trained bidirectional multimodal pretraining produces better hidden representations for flow-matching action generation in VLA policies.

The strongest version keeps the flow-matching action expert fixed and isolates the pretraining prior of the backbone through the scale/architecture-matched C−D contrast, rather than relying on a raw diffusion-vs-stock comparison that conflates prior with scale. No single contrast is perfectly clean; the contribution is a design whose contrasts bracket the effect and make the residual confounds explicit, evaluated on a fair, reproducible, high-n benchmark.

 References

1. C. Chi, Z. Xu, S. Feng, E. Cousineau, Y. Du, B. Burchfiel, R. Tedrake, S. Song. Diffusion Policy: Visuomotor Policy Learning via Action Diffusion. RSS 2023; IJRR 2025. arXiv:2303.04137.
2. Physical Intelligence. π0: A Vision-Language-Action Flow Model for General Robot Control. 2024.
3. K. Black et al. π0.5: A Vision-Language-Action Model with Open-World Generalization. CoRL 2025, PMLR 305:17–40.
4. Physical Intelligence. OpenPI: Open-source implementations for π0 and π0.5. GitHub.
5. Google AI for Developers. DiffusionGemma model overview. 2026.
6. Google AI for Developers. DiffusionGemma model card. 2026.
7. X. Li, V. Belagali, J. Shang, M. S. Ryoo. Crossway Diffusion: Improving Diffusion-based Visuomotor Policy via Self-supervised Learning. ICRA 2024. arXiv:2307.01849.
8. C. Liu, Y. Zhao, H. Chen, X. Han, J. Gao, E. Adeli, Y. Du. Ordered Action Tokens for Visuomotor Policy Learning. arXiv:2607.21670, 2026.
9. Y. Ze, G. Zhang, K. Zhang, C. Hu, M. Wang, H. Xu. 3D Diffusion Policy. arXiv:2403.03954, 2024.
10. J. Ye, Z. Xie et al. (HKU NLP & Huawei Noah's Ark Lab). Dream 7B: Diffusion Large Language Models. arXiv:2508.15487, 2025.
11. Qwen Team. Qwen2.5 Technical Report. arXiv:2412.15115, 2024.
12. S. Li, K. Kallidromitis et al. LaViDa: A Large Diffusion Language Model for Multimodal Understanding. arXiv:2505.16839, 2025.
13. DreamLM (HKU NLP). Dream-VL and Dream-VLA: A Diffusion VLM and a Diffusion VLA. 2025–2026.
14. Y. Wen et al. LLaDA-VLA: Vision Language Diffusion Action Models. 2025.
15. Discrete Diffusion VLA: Bringing Discrete Diffusion to Action Decoding in VLA Policies. arXiv:2508.20072, 2025.
16. B. Liu, Y. Zhu, C. Gao, Y. Feng, Q. Liu, Y. Zhu, P. Stone. LIBERO: Benchmarking Knowledge Transfer for Lifelong Robot Learning. NeurIPS D&B 2023. arXiv:2306.03310.