Robotic VLA Benefits from Joint Learning with Motion Image Diffusion

In plain words

A vision-language-action model looks at camera images and a sentence such as "put the cup on the plate", and outputs the robot's next movements. Models of this kind learn largely by copying demonstrations, and they never learn explicitly how the scene is going to move.

This paper adds a second output to such a model. Next to the head that predicts robot movements, a second head predicts a picture of how the scene moves over the next second: each pixel is coloured by the direction and size of its motion. Both heads share the same backbone and are trained together, so the backbone is pushed to represent motion as well as control. When the robot runs, the motion head is not used, so the speed is the same as the original model. The authors report higher success rates on two simulation benchmarks and a small real-world test. Those numbers come from the authors' own experiments.

TL;DR

• Joint learning adds a diffusion-transformer "motion head" to pi-series VLAs. It predicts optical-flow images §3.2 fig. 2.
• Targets are RAFT flow between frames o_t and o_{t+k}, encoded by a frozen SDXL-VAE §3.3 app. A.1 app. A.2.
• Only the action head runs at inference, so the test-time pathway is unchanged §1 §4.3.
• Author-reported: LIBERO average 97.5% for pi0.5 (96.9% baseline) tab. 1; RoboTwin easy average 58.0% (pi0 45.1%) tab. 2.
• The code is not released (project page says "Code (Coming Soon)"); see Paper vs code.

Problem

VLA models map observations and instructions directly to actions and are usually trained by imitating expert trajectories. The authors argue that this gives no explicit reasoning about future dynamics §1. Video-prediction and unified world-model VLAs predict future appearance, which the authors say "overlook[s] motion learning" §1. Language descriptions of motion and future-frame prediction were the representations the authors compare against §4.4 tab. 3.

Method

The model is a standard pi-series VLA with one extra head §3.2.

1. The VLM encodes the observation and instruction into a prefix z_t §3.1 eq. 1.
2. The action head predicts a chunk of k future actions from z_t §3.1 eq. 2.
3. The motion head μ is a Diffusion Transformer that predicts motion latents m_t from z_t §3.2 eq. 3. A frozen VAE decoder turns m_t into a motion (flow) image §3.2.
4. Both heads use the same flow-matching loss §3.2 eq. 4 eq. 5. The total loss is the unweighted sum of the two §3.2 eq. 6 eq. 7.
5. Training is two-stage §3.3:
   • Warm-up: only the motion head is trained, with all other parameters frozen, on DROID-derived flow §3.3 §4.1.
   • Joint training: all parts are unfrozen except the VAE encoder and decoder, and both heads are optimised §3.3.
6. The motion target is an optical-flow image computed with RAFT between o_t and o_{t+k}. Its time window equals the action chunk's §3.3 app. A.1. The flow is mapped to HSV (angle to hue, magnitude to saturation, value fixed at 1) and then to RGB app. A.1.
7. The motion head is a DiT with N = 6 transformer blocks. Each block has self-attention over latent tokens and cross-attention in which the latents are queries and z_t gives keys and values app. A.2.
8. At inference the motion head is not used, so the action pathway is unchanged §1 §4.3. The motion latent can be decoded into a flow image without optimisation §3.3 app. A.2.

Pseudocode, derived from the text (the code is not released):

# Joint training step  [§3.2] [eq. 4-7] [§3.3]
# Input:  observations, instructions, action chunks A (B, k, d),
#         motion latents m (B, 4, 28, 28)  [app. A.2]
# Output: scalar loss L = L_action + L_motion  [eq. 7]

def training_step(observations, instructions, actions, motion_latents):
    # Shared multimodal prefix z_t from the VLM backbone  [§3.1] [eq. 1]
    prefix = vlm(observations, instructions)

    # Action head: flow-matching loss on the action chunk  [§3.2] [eq. 5]
    action_loss = flow_matching_loss(action_head, actions, prefix)

    # Motion head: same loss on the motion latents  [§3.2] [eq. 3] [eq. 5]
    motion_loss = flow_matching_loss(motion_head, motion_latents, prefix)

    # Unweighted sum of the two losses  [§3.2] [eq. 7]
    return action_loss + motion_loss


def flow_matching_loss(head, targets, prefix):
    # Gaussian noise with the same shape as the target  [§3.2] [eq. 4]
    noise = sample_gaussian_like(targets)
    # Flow time tau in [0, 1]; the Beta parameters are not stated  [§3.2]
    tau = sample_beta()
    # Clean target weighted by tau, noise by (1 - tau)  [§3.2] [eq. 4]
    noisy = tau * targets + (1 - tau) * noise
    # Predicted velocity; u is not given in closed form  [eq. 5]
    velocity = head(noisy, prefix)
    # Squared L2 error; mean or sum reduction not stated  [eq. 5]
    return squared_l2(velocity, targets - noise)  # (see Pitfalls)

# Motion head block: one of N = 6 DiT blocks  [app. A.2]
# Input:  latent tokens (B, n, d), prefix z_t (B, P, d)  [app. A.2]
# Token count n and tokenisation of the latent are not stated  [app. A.2]

def motion_head_block(latent_tokens, prefix):
    # Self-attention over the latent tokens  [app. A.2]
    latent_tokens = self_attention(latent_tokens)
    # Cross-attention: latents as queries, z_t as keys and values  [app. A.2]
    latent_tokens = cross_attention(latent_tokens, prefix)
    return latent_tokens

# Build the motion target from a frame pair  [app. A.1] [app. A.2]
# Output: motion latent (4, 28, 28) for a 224 x 224 flow image  [app. A.2]

def make_motion_latent(frame_t, frame_t_plus_k):
    # Dense optical flow with RAFT, one (dx, dy) per pixel  [app. A.1]
    flow_x, flow_y = raft(frame_t, frame_t_plus_k)
    # Magnitude divided by 64, the dataset-wide maximum  [app. A.1]
    magnitude = sqrt(flow_x ** 2 + flow_y ** 2) / 64
    # Angle sets the hue, magnitude sets the saturation  [app. A.1]
    angle = atan2(flow_y, flow_x)
    # Value is fixed at 1; the result is converted from HSV to RGB  [app. A.1]
    flow_image = hsv_to_rgb(angle, magnitude, 1.0)
    # Frozen SDXL-VAE encoder; latent scaling is not stated  [app. A.2]
    return frozen_vae_encode(flow_image)

# Inference: the motion head is not run  [§1] [§4.3]
# Input:  observation, instruction. Output: action chunk (k, d)

def predict_actions(observation, instruction):
    # Same VLM pass and action sampling as a standard VLA  [§1] [§3.2]
    prefix = vlm(observation, instruction)
    # Number of denoising steps and the solver are not stated  [§3.2]
    return sample_action_flow(action_head, prefix)

Hyperparameters

Results

Author-reported. Setup: LIBERO (4 suites, 10 tasks, 50 demos per task) §4.2; RoboTwin 2.0, 7 tasks, dual Aloha-AgileX, 50 demos per task, easy and hard modes §4.3; a real-world Dobot Nova 5 test §4.5. Baselines in RoboTwin (DP, RDT, ACT) are the official leaderboard numbers, trained per task; this paper's model is one multi-task model §4.3.

Headline comparisons:

LIBERO, success rate per suite tab. 1:

RoboTwin, success rate per task tab. 2:

Motion representation ablation on LIBERO tab. 3 (pi0 and pi0.5 backbones):

Other reported results: data efficiency on LIBERO-Long with pi0 uses 25%, 50% and 75% of the data; the 25% result is 59.0%, which the text reports as a 14.2% gain over the matched pi0 baseline §4.4 fig. 5. Real-world results are shown only as a figure fig. 6.

Paper vs code

No official code. The paper's project page (https://vla-motion.github.io/, linked at the top of the LaTeX source) lists "Code (Coming Soon)" with no link. Searches: the GitHub repository search for "motion image diffusion VLA", "vla-motion" and "joint learning motion image diffusion pi0" returned no matching repository. The vla-motion organisation holds only the project website repository (vla-motion.github.io), which is not a code release. A search restricted to org:SalesforceAIResearch returned no results.

Pitfalls

• Flow-matching convention. Eq. 4 is X^τ = τX + (1 − τ)ε §3.2, so τ = 1 is clean data. The paper does not write the target velocity u in closed form. Differentiating eq. 4 with respect to τ gives X − ε; this is our derivation, not a stated formula. The Beta parameters, the reduction of the squared L2 loss, and the number of inference steps are not stated §3.2 app. A.2.
• Citation key for DROID. §3.3 cites DROID as teed2020raft, which is RAFT's key (the same key is used for RAFT in §3.3 and app. A.1). DROID is cited as khazatsky2024droid in §2. Treat the DROID source as the dataset, not the RAFT paper.
• Action chunk length. k is set to "match the control frequency" so the window covers the next second §4.1. No numeric value of k, and no control frequency per benchmark, is stated.
• Flow image encoding. The hue scale (angle to [0, 1]), the saturation clipping, and whether the divisor 64 applies to every dataset are not stated app. A.1. An implementation will have to pick these.
• Latent handling. The VAE is frozen and never trained §3.2 §3.3. Latent scaling or normalisation before diffusion is not stated app. A.2. The 224 × 224 input to 4 × 28 × 28 latent mapping is given only in the appendix; the main text says H/8 × W/8 §3.2.
• Parameter count. "400M parameters" §4.1 is given together with N = 6 blocks app. A.2, but no width or depth detail is given. Check the count against your own build.
• Joint-stage data. The warm-up uses DROID §3.3. The paper does not say which data the 100k joint-training steps use, or whether the benchmark runs repeat the DROID warm-up §4.1.
• Text vs table, RoboTwin. §4.3 says the gain over pi0 is 13.1%. tab. 2 gives 58.0 − 45.1 = 12.9 points, and the caption says +12.9%.
• Text vs table, LIBERO-Long. §4.2 says the pi0.5 gain on LIBERO-Long is 4.0%. tab. 1 gives 96.2 − 92.4 = 3.8 points.
• "Consistent" hard-mode gains. §4.3 says the method "consistently improves success rates across both easy and hard modes". tab. 2 shows lower hard-mode scores than pi0 on Move Can Pot (2% vs 3%), Move Playingcard (3% vs 16%) and Stack Blocks Two (0% vs 5%). The hard average gain is 0.4 points.
• Largest gains. §4.3 names Stack Blocks Two and Move Playingcard as the largest gains. In tab. 2 the largest easy-mode gain is Move Can Pot (43% to 81%).
• Real-world "23%". The abstract gives "a 23% improvement" and §4.5 says "improving pi0 by 23%". The text does not say whether this is absolute or relative. The pi0 baseline is only in fig. 6. With 43% for the method, an absolute gain implies a baseline near 20%, and a relative gain implies near 35%.
• Small real-world test. Three tasks, 30 teleoperated demos each, LoRA fine-tuning on joint-learned pi0 §4.5. Trial counts, seeds and variance are not reported.
• Single-seed sims. Each simulation policy is tested with 50 or 100 rollouts §4.2 §4.3. No seed or confidence intervals are reported.
• Fair comparison. DP, RDT and ACT are per-task models with leaderboard numbers. pi0 and this method are one multi-task model §4.3. The comparison is not controlled.
• Motion head relies on the third-person camera. Failures come from occlusion in that view, which is a stated limitation §5 app. B.3.
• Cost not reported. The motion head adds about 400M parameters during training, RAFT flow is computed per training pair, and the VAE encodes every target. Training time and memory are not reported §4.1 §3.3. Inference cost is not changed by the design §1 §4.3.
• Qualitative motion claims. The claims that predicted flows are "physically grounded" and "temporally coherent" §4.4 rest on rollout figures fig. 4 and not on a quantitative motion-accuracy metric.
• Language representation details. The language baseline samples an 8 × 8 pixel grid, normalises positions to [0, 1] and describes displacement pairs §4.4. How these are tokenised for the LLM-style objective is not stated.

Lineage

• Builds on: pi0 and pi0.5 (cited as black2024pi_0 and intelligence2025pi_), adapting the PaliGemma architecture with an action expert trained by flow matching §2. No card exists for these papers yet.
• Uses: DiT, a Diffusion Transformer §3.2 (peebles2023scalable); flow matching §3.2 (lipman2022flow); RAFT for optical flow §3.3 (teed2020raft); LoRA for real-world finetuning §4.5 (hu2022lora).
• Compared against: FlowVLA §2 tab. 1, WorldVLA §2 tab. 1, TraceVLA (mentioned in §2), UVA (mentioned in §2), and the language and future-image representations tab. 3.
• Superseded by: not known from content on this wiki.