Future Representation Distillation for Imagination-free World Action Models

AI summary

Overview

Research area: Robot learning / vision-language-action policies built on pretrained video diffusion models (cs.RO).

Technical level: Advanced — the paper assumes familiarity with diffusion/flow-matching models, transformer attention masking, and imitation-learning policy training.

Scope: The paper introduces DreamFree, a World Action Model that distills the predictive knowledge of a frozen video diffusion model into a visual encoder so that robot actions can be predicted from the current observation alone, without generating future videos.

What This Paper Is About

World Action Models (WAMs) borrow the predictive knowledge of pretrained video diffusion models for robot control, but they typically do so by generating future video frames at deployment. That iterative video denoising is computationally expensive, and mistakes in the imagined future can propagate into wrong actions. The authors ask whether the same predictive knowledge can instead be baked directly into the visual features computed from the current observation, so that a policy can act well without ever imagining a future video.

Key Contributions

1. Future representation distillation. A method that trains a visual encoder to predict the intermediate representations of a frozen video diffusion model that has access to future frames, using only the current observation and language instruction. The frozen teacher and the alignment projector are discarded at deployment.

2. DreamFree, a two-stage World Action Model. Stage one learns the predictive visual encoder from video (no action labels required); stage two trains a readout module and a flow-matching action decoder on action-labeled robot demonstrations, with the encoder either frozen or jointly fine-tuned.

3. Imagination-free real-time control. With a 14B video diffusion backbone, DreamFree performs real-time closed-loop control at 11Hz, a 17.8× inference speedup over the video-generating WAM DreamZero with the same backbone size, and latency comparable to the VLA π0.5.

4. Evidence that video priors and unlabeled video transfer to control. Scaling the backbone and adding unlabeled video data improve success, and new tasks — including tasks from a different embodiment — can be learned from video demonstrations without action annotations.

Main Findings

• Representation probing: Using normalized mean squared error from frozen-feature probes on RoboCasa-Kitchen and DROID, the DreamFree encoder achieves lower action-prediction error than the pretrained video diffusion model given only the current observation, and lower than a variant trained without distillation. It also beats probing VAE latents of current and future frames, and approaches the video model that is given ground-truth future frames. Probing the readout module's features gives even lower error.

• Future probing: A linear probe reconstructing VAE latent patches of future frames shows the DreamFree encoder captures future visual information better than the no-distillation variant and the video model with current observation only; it approaches the future-frame-conditioned video model on RoboCasa-Kitchen, with a remaining gap on DROID.

• LIBERO: DreamFree reaches 98.6% average success across the four task suites (Spatial, Object, Goal, Long) without cross-embodiment pretraining, versus 94.1 (π0), 96.9 (π0.5), 97.1 (OpenVLA-OFT), 97.7 (Motus), 98.5 (LingBot-VA) and 97.6 (Fast-WAM).

• RoboCasa-Kitchen (24 tasks, 50 demos per task): 71.7% average success with the 14B backbone, the highest reported, versus 67.1 (Cosmos Policy), 66.4 (FLARE), 64.1 (GR00T-N1.5), 62.5 (π0), 44.4 (Fast-WAM) and 44.7 (Fast-WAM 14B). Scaling from 5B (58.5%) to 14B (71.7%) adds 13.2 percentage points, while Fast-WAM barely moves when scaled. Adding unlabeled videos from the atomic-task subset of RoboCasa365 raises 5B to 68.8% and 14B to 74.4%.

• Ablation: Removing the distillation loss drops LIBERO to 94.6 and RoboCasa to 47.2; removing the generative flow-matching loss gives 98.0 and 55.6. Distillation matters more, consistent with the probing results.

• Real-world DROID (10 tasks, 5 rollouts each): 64% average success versus 56% for DreamZero and 16% for π0.5; average task progress 83% versus 71.8% and 33.5%. DreamZero failed all trials on tasks such as cover strawberry with towel, hang hat on tripod, and put bolts in bowl, where inaccurate video predictions were observed. DreamFree often recovered after an initial error, whereas DreamZero tended to persist with unsuccessful actions.

• Efficiency: DreamFree achieves 4× the average success rate of π0.5 at 1.1× its latency. In a block-placing task, completion times are comparable to π0.5 for 2 or 4 blocks and much faster with 6 blocks. The paper reports the 17.8× speedup and 11Hz rate but does not state the raw per-chunk latency values in the text.

• Task transfer from video: Composite tasks absent from action-labeled data go from 0% average success (base model) to 21.7% when their videos are added with distillation; adding the same videos without distillation yields 0%. With RoboTwin and ALOHA-AgileX as the target robot, videos of composite tasks recorded on a dual-arm Franka Panda raise average success from 6.7% to 30%. Fine-tuning only the encoder on 10 video demonstrations per task (about 10 minutes per task on a single H200, action module frozen) improves success on both tested tasks, with larger gains from multi-view than head-camera-only video.

• Bimanual manipulation: Trained on a filtered 100-hour subset of ABC-130k without cross-embodiment training, the policy is shown qualitatively executing pick-and-place, pen uncapping, cloth folding and dish-rack loading on a YAM setup with new scenes and objects; no numeric success rates for these tasks are reported in the main text.

• A stated limitation: Not every task can be learned from video; a policy that could not open cabinets also failed to learn placing an object inside a cabinet through video-only fine-tuning, suggesting transfer is bounded by the atomic skills the policy already has.

Methodology in Plain English

The pipeline has two separate training stages.

Stage one — teach an encoder to "see ahead." Start from a pretrained video diffusion backbone (Wan2.2-5B or Wan2.2-A14B), which has learned how scenes evolve. A frozen copy of this model serves as the teacher: it is fed the whole video clip, including future frames (with noise added at a fixed level), and its internal hidden representation is recorded. The student is the same architecture, modified with a structured attention mask so that current-frame tokens can only attend to each other — meaning the student's current-frame features depend on the current observation and instruction alone, never on future frames. A small transformer projector expands the student's current-frame tokens to match the teacher's spatiotemporal layout. The loss maximizes cosine similarity between the projected student features and the teacher's features on future positions (this is the REPA-inspired distillation objective). A second loss, the backbone's original flow-matching generative objective, keeps the model from losing its native predictive ability. No action labels are needed in this stage, so any video demonstration can be used.

Stage two — read out actions. The predictive encoder's hidden states are passed to a readout module that uses 64 learned queries with self-attention and cross-attention (Perceiver-style) to compress the visual tokens into a fixed set of action-facing tokens. A 1B-parameter diffusion transformer action decoder, conditioned on those tokens, the instruction, proprioceptive state and (for DROID) the previous 8 actions, predicts an action chunk of horizon 32 via a flow-matching loss.

At deployment, one encoder forward pass plus a small number of action-denoising steps produce an action chunk. There is no video denoising loop, no teacher, and no projector — which is the source of the 17.8× speedup over a video-generating WAM.

Why This Matters

Impact on research. The paper reframes a design question in video-based robot learning: instead of asking how to generate better future videos, it asks how the video model's predictive knowledge should be represented for control. The probing experiments give direct evidence that future-frame information is recoverable from a current-frame encoder, and the ablations show the distillation objective is what carries most of the benefit. It also offers a route to using video data — including video from other embodiments — without action labels, which is a practical bottleneck for scaling robot data.

Real-world applications:

• Real-time manipulation on physical robots where reaction latency matters (the 11Hz closed-loop rate is the enabling property).
• Deformable-object and multi-step rearrangement tasks, where imagined video tends to be inaccurate and the video-generating baseline failed entirely.
• Scaling policies by absorbing large unlabeled video corpora, including human or other-robot video, without new teleoperation.
• Quick task adaptation on a deployed robot via short video-only encoder fine-tuning (about 10 minutes per task on a single H200 in the reported setup).

Industry relevance. The result that a 14B video backbone can be deployed at latency comparable to a VLA, while outperforming it by 4× in average success on real tasks, speaks to the cost side of deploying large video models in robotics. The ability to convert stronger video backbones and extra video data directly into better control, without extra action annotation, is relevant to anyone with a video model already in hand.

Future Directions

• Understanding the limits of video-only transfer. The authors report that tasks requiring atomic skills the policy lacks cannot be learned from video alone. What determines which new tasks are learnable this way, and can the atomic-skill repertoire be expanded first?
• Choosing what to distill. The distillation target is the teacher's layer-j full-video representation, and the readout draws from a separately chosen layer k. How sensitive are results to these layer choices and to the fixed teacher noise level t_d?
• Reducing the gap on DROID. Future-latent reconstruction approaches the future-frame-conditioned video model on RoboCasa-Kitchen but not on DROID, which suggests room to improve the encoder on messier real-world data.
• Closing the simulation–real-world and cross-embodiment gaps at scale. The transfer experiments are small (20 episodes per task, 10 tasks with 5 rollouts each), and the bimanual results are qualitative; larger, standardized evaluations across embodiments would clarify how far the approach generalizes.

Target Audience

Researchers and engineers working on robot learning, vision-language-action models and video-diffusion-based policies; practitioners who want the predictive benefits of video models without their inference cost; and students with a background in deep generative models who are interested in how pretrained video representations can be repurposed for control. Readers without familiarity with diffusion/flow matching and transformer attention will find the training details dense, though the core idea is stated plainly.

Authors’ abstract

World Action Models (WAMs) use the predictive knowledge of pretrained video diffusion models for robot control through future video prediction. However, iterative video denoising is costly, and inaccurate visual predictions can compromise action prediction. Our key idea is to make this predictive knowledge available through visual representations computed from the current observation, enabling action prediction without generating future videos. We introduce DreamFree, a WAM that repurposes a pretrained video diffusion model as a predictive visual encoder through representation distillation. The visual encoder learns to predict a frozen video diffusion model's representation using only the current observation. An action module is then trained to predict actions by extracting action-relevant information from the visual encoder. Across real-world and simulated tasks, DreamFree achieves higher success rates than state-of-the-art VLAs and WAMs, while performing real-time closed-loop control at 11Hz with a 14B video diffusion backbone. Experiments show that DreamFree can translate stronger video priors and additional video data into better robot control and enable learning new robot tasks from video demonstrations, including videos from a different embodiment, without action annotation for the new tasks.