Robotic VLA Benefits from Joint Learning with Motion Image Diffusion

TL;DR

The paper addresses the lack of explicit motion reasoning in Vision-Language-Action (VLA) models by introducing a joint learning framework with motion image diffusion. It adds a parallel motion head, implemented as a Diffusion Transformer, that predicts optical-flow-based motion images using a shared VLM backbone, trained via flow-matching losses and a two-stage procedure with a frozen VAE for latent diffusion. This approach preserves the standard inference pathway and improves motion-aware representations, achieving state-of-the-art results on LIBERO (up to 97.5% in π0.5) and RoboTwin (58.0%), with notable real-world gains under limited data. The work demonstrates that dense, pixel-space motion supervision complements action learning, enhancing temporal coherence and generalization in large-scale VLA models.

Abstract

Vision-Language-Action (VLA) models have achieved remarkable progress in robotic manipulation by mapping multimodal observations and instructions directly to actions. However, they typically mimic expert trajectories without predictive motion reasoning, which limits their ability to reason about what actions to take. To address this limitation, we propose joint learning with motion image diffusion, a novel strategy that enhances VLA models with motion reasoning capabilities. Our method extends the VLA architecture with a dual-head design: while the action head predicts action chunks as in vanilla VLAs, an additional motion head, implemented as a Diffusion Transformer (DiT), predicts optical-flow-based motion images that capture future dynamics. The two heads are trained jointly, enabling the shared VLM backbone to learn representations that couple robot control with motion knowledge. This joint learning builds temporally coherent and physically grounded representations without modifying the inference pathway of standard VLAs, thereby maintaining test-time latency. Experiments in both simulation and real-world environments demonstrate that joint learning with motion image diffusion improves the success rate of pi-series VLAs to 97.5% on the LIBERO benchmark and 58.0% on the RoboTwin benchmark, yielding a 23% improvement in real-world performance and validating its effectiveness in enhancing the motion reasoning capability of large-scale VLAs.

Figures (11)

• Figure 1: Overview. Our joint learning strategy seamlessly extends existing large-scale VLAs with motion image diffusion -- learning motion and action jointly through a shared VLM backbone, enhancing their motion reasoning abilities, and maintaining the same inference pipeline as in standard VLA models.
• Figure 2: Overview of joint learning VLA with motion image diffusion.
• Figure 3: Overview of different motion representations. We compare three motion representations for joint learning, where the motion image is most effective.
• Figure 4: Qualitative visualization of predicted action and motion during rollout. We show example rollouts from both LIBERO and RoboTwin. For each example, the first row shows the observed frames with the robot actions predicted by the action head, and the bottom row visualizes the flow images predicted by the motion head. Please see supplementary materials for examples of long-horizon tasks.
• Figure 5: Data efficiency on LIBERO-10. We demonstrate that joint learning with motion image diffusion improves data efficiency over action-only learning.