Dual-Stream Diffusion for World-Model Augmented Vision-Language-Action Model

TL;DR

This work tackles modality conflict in world-model augmented vision-language-action models by proposing DUST, a dual-stream diffusion framework that preserves separate action and vision streams while enabling cross-modal knowledge exchange through shared attention. It introduces a decoupled diffusion training regime with independent per-modality noise schedules and a joint flow-matching objective, plus an asynchronous inference strategy that allocates more diffusion steps to high-dimensional vision tokens. Across RoboCasa, GR-1, and Franka Real tasks, DUST consistently outperforms baselines, with additional gains from test-time scaling and action-free pretraining on BridgeV2, demonstrating strong data efficiency and transfer capability. The results suggest significant practical impact for scalable, robust robotic policy learning in diverse environments and tasks.

Abstract

Recently, augmenting vision-language-action models (VLAs) with world-models has shown promise in robotic policy learning. However, it remains challenging to jointly predict next-state observations and action sequences because of the inherent difference between the two modalities. To address this, we propose DUal-STream diffusion (DUST), a world-model augmented VLA framework that handles the modality conflict and enhances the performance of VLAs across diverse tasks. Specifically, we propose a multimodal diffusion transformer architecture that explicitly maintains separate modality streams while enabling cross-modal knowledge sharing. In addition, we propose training techniques such as independent noise perturbations for each modality and a decoupled flow matching loss, which enables the model to learn the joint distribution in a bidirectional manner while avoiding the need for a unified latent space. Furthermore, based on the decoupled training framework, we introduce a sampling method where we sample action and vision tokens asynchronously at different rates, which shows improvement through inference-time scaling. Through experiments on simulated benchmarks such as RoboCasa and GR-1, DUST achieves up to 6% gains over a standard VLA baseline and implicit world-modeling methods, with our inference-time scaling approach providing an additional 2-5% gain on success rate. On real-world tasks with the Franka Research 3, DUST outperforms baselines in success rate by 13%, confirming its effectiveness beyond simulation. Lastly, we demonstrate the effectiveness of DUST in large-scale pretraining with action-free videos from BridgeV2, where DUST leads to significant gain when transferred to the RoboCasa benchmark.

Figures (11)

• Figure 1: Architectures of world-model augmented VLAs. (a) Unified Joint Diffusion concatenates action and vision tokens and generates both with a single model. (b) Causal Diffusion uses separate models with one-way conditioning. (c) Dual-Stream Diffusion (ours) maintains separate streams for each modality while enabling cross-modal knowledge transfer through shared attention.
• Figure 2: Dual-stream diffusion (DUST) architecture. Our architecture has a (1) VLM model $\text{VLM}\phi(\cdot)$ that processes current observation and task instruction to produce semantic representations, and a (2) diffusion model $\pi_\theta$ which conditions on these representations to generate actions and future observation embeddings.
• Figure 2: Evaluation on GR-1. Success rates (%) on GR-1 benchmark for 16 pick-and-place (PnP) and 8 articulated (Art.) tasks. 300 and 1,000 demos per task are used. $^\dagger$: reproduced results.
• Figure 3: Overview of vision-action joint sampling. During inference, we sample over $N_A$ steps for action tokens and $N_o = q\times N_A$ steps for vision tokens. The global timestep advances by $\Delta\tau_o = 1/N_o$, where vision tokens are updated every step and action tokens are updated only every $q$ steps in $\Delta\tau_A = 1/N_A$ strides. The default $q$ value is 1, and increasing it allows test-time scaling.
• Figure 4: Real-world task instructions. For the real-world experiments, we utilize 4 pick-and-place tasks with the Franka Research 3 robot. The tasks are categorized by their distinct source-target pairs (box, bowl, plate, etc.) and each contains 4 different objects (cup, doll, cube, sponge).