DreamFree Distills Future Representations for Imagination-Free Robot Control

KAIST researchers introduce DreamFree, a world-action model that distills predictive capabilities from a video diffusion model into a visual encoder. This enables robots to output actions from a single current frame, bypassing the computationally expensive step of generating future video sequences during deployment.

To learn manipulation, a robot must first anticipate how a scene will evolve. Many world action models achieve this using pretrained video diffusion models: they first imagine future frames and then output actions based on those predictions. However, generating these imagined frames is the most computationally expensive step during deployment, requiring dozens of denoising steps to produce a single action sequence. Any deviation in the imagined frames inevitably leads to erroneous actions.

Video diffusion models have already learned spatiotemporal dynamics—such as falling objects, cloth folding, and opening or closing doors—from massive video datasets, making this information highly valuable for robot control. The challenge lies in extracting these dynamics without having to generate complete future video clips.

On November 3, a KAIST team published the paper "Future Representation Distillation for Imagination-Free World Action Models," introducing DreamFree. This world action model distills the predictive capabilities of a video diffusion model into a visual encoder, enabling the robot to output actions based solely on the current frame.

The paper has eight authors in total. Jaehyeong Jo and Sangwon Jang are co-first authors, with KAIST Professor Sung Ju Hwang and Assistant Professor Kimin Lee serving as co-advisors.

The training pipeline for DreamFree consists of two stages.

• The first stage relies exclusively on video data for training: the video diffusion model is frozen and observes both current and future frames, while an encoder that receives only the current frame as input aligns its intermediate representations with those output by the teacher model.
• In the second stage, the encoder parameters are frozen, and a feature readout module along with an action decoder are trained using action demonstration data.

Built on a 14B foundation model, the system achieves an average success rate of 64% across ten real-robot tasks in DROID; for comparison, DreamZero achieves 56% and π0.5 achieves 16%. Single action chunk inference takes 91 milliseconds, reaching a control frequency of 11 Hz, representing a 17.8x speedup over DreamZero. Across 24 tasks in RoboCasa, the average success rate is 71.7%; after additionally incorporating video data without action labels, the metric improves to 74.4%.

The key innovation of this work lies in where the supervision signal is applied. The team uses structured masks to isolate current and future information, applying the distillation process directly to the feature layer accessed by the action module, thereby seamlessly integrating the learned predictive capabilities into the control pipeline. Action prediction does not necessarily require the pixel-level details of a complete video; it only needs to extract the relevant information accessible to the action module.

The ultimate potential of this approach depends on whether the distilled representations can reliably encode future information, and whether this feature layer can continuously improve robotic control performance by leveraging a larger-scale video foundation model and more extensive video data.

Distilling future representations into current frame features

The DreamFree implementation pipeline consists of four steps: first, modifying the video diffusion backbone to create an encoder that accepts only current-frame inputs; second, applying structured masking to prevent the encoder from directly accessing future information; third, distilling future representations from a frozen teacher model; and finally, decoding these representations to output actions.

Step 1: Modifying the encoder. The team initializes the model on a pre-trained video diffusion backbone, using the hidden state of the current-frame latent tokens as the visual representation. The original backbone employs a bidirectional attention mechanism that allows the current frame to attend to future frames, meaning the distilled features cannot distinguish whether the information was observed or predicted.

Step 2: Blocking this pathway with masking. Current-frame tokens can only attend to one another, while future-frame tokens can attend to both the current frame and themselves, with text cross-attention disabled for the future tokens. After this processing, the hidden state of the current frame is determined solely by the current observation and the task instruction. By restricting the current frame strictly to the present, the mask leaves distillation as the only pathway for the encoder to access future information.

The core of this entire process lies in the target of the supervision signal. With the masking constraints applied, the current-frame features must independently infer the future; the distillation alignment is then applied directly to this exact feature layer, which serves as the input for the downstream action module.

Step 3: Performing representation distillation. The teacher model reuses the same video diffusion architecture with its weights frozen throughout. It takes the full video clip and instruction as input, extracting the j-th layer representation from the noisy latents—a layer that has attended to both current and future frames. The student model, by contrast, takes only a single clean current-frame image as input. A projector expands its N tokens into F×N tokens, which are then augmented with frame embeddings and passed through a two-layer decoder to align with the teacher model's spatiotemporal layout.

Expanding to F×N tokens allows the single current frame to be compared position-by-position against the teacher's entire future sequence. The loss function calculates cosine similarity only at future positions, with current-frame tokens removed from both sides to prevent the model from taking the shortcut of directly copying the input. The team also retains the backbone's native flow-matching generation loss to prevent the encoder from forgetting the video generation dynamics learned during pre-training, with both losses backpropagated jointly to update the parameters.

Step 4: Reading out actions from features. The readout module extracts features from the $k$-th layer of the encoder. It uses 64 learnable queries, passing them through two layers of self-attention and cross-attention to compress the visual tokens into fixed-length action tokens. This design spares the action decoder from having to adapt to a variable number of video tokens. The action decoder is a 1B-parameter DiT, comprising 30 blocks, a hidden dimension of 1024, and 24 attention heads. It performs denoising via flow matching in the action space. The conditional inputs consist of action tokens, task instructions, and the robot's proprioceptive state, with an additional 8-step action history included for the DROID task. During deployment, the encoder runs inference only once, while the decoder iterates for 10 to 16 steps. After generating an action chunk, the system executes the first 10, 24, or 32 steps, depending on the dataset configuration.

Once the network architecture is finalized, the key to the success of the approach lies in the training sequence: the encoder first learns representation distillation, followed by the training of the action module.

Staged training and low-latency deployment optimization

The complete training pipeline consists of three stages: first, adapting the video backbone to robot camera footage; second, training the encoder using video data without action labels; and third, training the readout module and action decoder using action demonstration data. During deployment, only a single forward pass of the encoder is retained.

Stage 1: Backbone adaptation. The team selects the high-noise expert models of Wan2.2-5B and Wan2.2-A14B, fine-tuning them alongside the text encoder and video VAE. The inputs consist of multi-view images processed using a proprietary packing format. The adapted weights are then used to initialize both the teacher and student models simultaneously. The teacher model remains frozen throughout the subsequent process, while the student model continues training with a structured mask.

Stage 2: Using only video data without action labels. Video clips are temporally downsampled by a factor of 4 to yield 9 frames. The current frame remains in its original clean state, while noise is added to the future frames at random timesteps. The teacher model is assigned a separate fixed noise level: 0.7 for simulation benchmark tasks and 0.8 for the DROID task. Representations are then extracted from these noisy videos to serve as distillation targets.

The feature layers selected for the student and teacher models are fixed according to the base model specifications: for the 5B version, distillation uses the 25th layer and feature readout uses the 19th layer; for the 14B version, both distillation and readout use the 31st layer. The simulation experiments employ rank-128 LoRA to fine-tune the encoder, whereas the 14B model on the DROID task undergoes full fine-tuning. All models are optimized using AdamW with a learning rate of 1e-4.

Stage 3: Introducing action labels. The encoder parameters are frozen, and the readout module and action decoder are trained from scratch. The inputs include current observations, task instructions, proprioceptive state, and action chunks. The loss function uses flow matching loss in the action space. The action chunk has a fixed length of 32, and the DROID task additionally includes an 8-step action history. During training, the entire history sequence has a 0.2 probability of being dropped, with missing positions filled using masks. Independent batch sizes are configured for different datasets; for the DROID task, the batch sizes in the two training stages are 64 and 4096, respectively. The entire training process runs on 16 H200 GPUs and takes approximately one week.

Separating the training of the encoder and the action module is another key design of this work. The encoder is trained solely under video signal supervision and its parameters are frozen once training is complete. The action module is then trained independently, preventing action gradients from altering the learned visual representations. In the appendix, the team likens this isolation strategy to Knowledge Insulation, with the key distinction being that this work adapts the paradigm for video diffusion models.

During deployment, the computational overhead is concentrated in a single forward pass of the encoder. The pipeline no longer involves video denoising or classifier-free guidance; on the action side, it requires only 10 or 16 steps of Euler integration.

Engineering optimizations also focus on this stage: instruction embeddings are cached across multiple queries, and the keys and values for the action conditions are computed just once per action chunk, allowing all denoising steps within that chunk to reuse them.

By leveraging fused attention kernels, fused QKV projections, torch.compile, and CUDA graphs, the team chains VAE encoding, feature extraction, and action sampling into a complete rollout pipeline.

For the model built on the 14B foundation, inference for a single action chunk on a single H200 GPU takes 91 milliseconds, achieving a control frequency of up to 11 Hz, with VRAM usage fitting within a single 48 GB A6000. This approach shifts the deployment overhead from video denoising to action denoising: the former typically requires dozens of iterative steps, whereas the latter requires only 10 to 16 steps.

While the speed metrics meet the targets, experiments are still needed to verify exactly what future information the distilled features have captured.

Topping simulation benchmarks and validating real-robot control performance

The validation process follows four progressive stages: first, verifying whether the features capture future information; second, evaluating performance in simulation; third, conducting real-robot experiments; and finally, measuring inference latency and throughput. This step-by-step approach provides evidence spanning from the representation level all the way to physical robot control.

The first set of experiments employs probe evaluations. The team trained two types of probes on frozen features: an action probe to predict action sequences, and a future probe to reconstruct the VAE latent variables of future frames. The comparison includes four setups: a native video model taking only the current frame as input, an encoder with the distillation branch removed, the DreamFree encoder, and a video teacher model that can directly read the ground-truth future frames.

In the action probe evaluation, the prediction error of the DreamFree encoder is lower than that of the video model using only the current frame and the control group with distillation removed. It also outperforms the approach that concatenates the current frame with the VAE latent variables of future frames, achieving performance close to that of the teacher model, which can directly observe future frames. For the future probe task, its performance on the RoboCasa dataset is comparable to the teacher model that has seen future frames, though a slight gap remains on the DROID dataset. Feeding the features output by the readout module back into the action probe further reduces the error. Relying solely on features extracted from the current frame, the model can reconstruct future frame latent variables that approach the quality of those derived from actually observed future frames.

The elegance of this probe design lies in isolating and evaluating two core questions. The future probe reconstructs the VAE latent variables of future frames to answer whether the features contain future information; the action probe predicts actions to verify whether this future information can be effectively utilized by the control module. The two tests are independent, and all control groups reuse the same set of probes, meaning any performance differences can be attributed directly to the quality of the representations themselves.

The second part consists of simulation evaluations. DreamFree achieves an average success rate of 98.6% across the four LIBERO suites, essentially matching the state-of-the-art 98.5% of LingBot-VA. In the 24 kitchen tasks of RoboCasa, the 14B base model achieves an average success rate of 71.7%, surpassing the 67.1% of Cosmos Policy and other VLA baselines; Fast-WAM, which also skips video generation during deployment, reaches only 44.4%. When scaling the base model from 5B to 14B parameters, DreamFree's performance improves by 13.2 percentage points, whereas Fast-WAM improves by only 0.3 percentage points. This demonstrates that generative supervision alone does not equate to a predictive representation suitable for control; distillation is the critical step that translates a larger base model into higher task success rates.

Ablation experiments that remove each of the two losses individually further highlight the critical role of distillation. Removing the distillation loss causes the success rate on RoboCasa to drop from 71.7% to 47.2%, a significantly larger decrease than when the generation loss is removed; additionally, introducing unlabeled video data improves the 14B version's performance from 71.7% to 74.4%.

The third part involves real-robot testing. Across ten DROID tasks, each repeated five times, the system is deployed in a zero-shot manner without any task-specific fine-tuning. DreamFree achieves an average success rate of 64% and a task completion progress of 83%; the corresponding figures for DreamZero are 56% and 71.8%, while π0.5 achieves 16% and 33.5%.

The performance gap is concentrated in tasks requiring the handling of object deformation and the continuous rearrangement of items. For three specific tasks—covering strawberries, hanging a hat on a tripod, and placing three bolts into a bowl—DreamZero fails all five attempts, whereas DreamFree succeeds in at least some trials. Observations from failed rollouts show that once DreamZero's initial action deviates, it tends to continue executing along the erroneous trajectory; DreamFree, however, possesses the ability to correct itself and realign with the target.

The team also conducted a supplementary set of experiments on the YAM dual-arm robotic platform. Training used only a subset of about 100 hours from the ABC-130k dataset, which accounts for roughly one-thirty-fifth of the total 3,500 hours of data. Even in entirely new scenes with novel objects, the model can still perform manipulation tasks such as loading items onto plates, uncapping pens, and folding fabric.

The fourth experiment measures inference latency. DreamFree takes 91 milliseconds per action chunk, compared to 1,632 milliseconds for DreamZero and 81 milliseconds for π0.5. Compared to π0.5, DreamFree incurs only about a 10% increase in latency while achieving four times the average success rate. In a comparative test of placing varying numbers of blocks into a bowl, the completion times for all three models are similar in the two-block and four-block scenarios; however, DreamFree demonstrates a significant speed advantage in the six-block scenario. By eliminating the need to wait for video generation before outputting actions, the time saved in the closed loop translates directly into higher throughput.

Following the validation of the real-robot experiments, one question remains: can the model generalize to entirely new tasks by relying solely on video data?

Generalizing to new tasks using only video

The video transfer approach incorporates three integration methods: adding same-embodiment videos during encoder pre-training, adding videos from different embodiments, or fine-tuning only the encoder after policy training. None of these methods provide action labels for the new tasks.

The first method combines atomic skills into composite tasks. Composite tasks consist of a sequence of multiple atomic skills. The team selected three composite tasks in RoboCasa, adding the corresponding videos to the encoder's pre-training process, while the action module continued to be trained exclusively on demonstration data from atomic tasks.

The baseline model fails completely on these three tasks, achieving an average success rate of 0%. After introducing video and incorporating distillation, the average success rate increases to 21.7%, with successful episodes occurring across all three tasks; however, when video is added but the distillation branch is removed, the average success rate remains at 0%. Taking microwave defrosting as an example, the task requires sequentially opening the door, placing the item inside, closing the door, and pressing the button. Success requires achieving all four subgoals in order, and DreamFree can complete this entire sequence of actions.

The second method is cross-embodiment transfer. The target robot is the ALOHA-AgileX from RoboTwin, and the supplementary transfer videos are collected by a dual-arm Franka Panda, with 50 videos per task, totaling 150 videos. The baseline model only occasionally succeeds in the bowl-stacking task, yielding an average success rate of 6.7% across the three tasks; after introducing cross-embodiment videos, the average success rate increases to 30%. Throughout the process, action labels are derived exclusively from atomic tasks on the ALOHA-AgileX, and no action demonstrations are provided for the new tasks.

The third method adds videos post-training. Based on the RoboTwin atomic task model, the team first removes the target tasks from the training set, then fine-tunes only the encoder while keeping the action module frozen throughout. They provide 10 videos per task, and fine-tuning takes approximately 10 minutes on a single H200 GPU.

All three transfer schemes share a core constraint: videos for new tasks are fed exclusively into the encoder, while the action module is trained solely on action labels from atomic tasks. Task-specific temporal knowledge converges within the representation layer, whereas learned motor skills remain isolated in the action module. Through this isolation design, the team can determine whether the performance gains from the added videos stem from task-level temporal logic or entirely novel action primitives.

Multi-view videos yield significantly higher metric improvements than single-view videos, indicating that the encoder can infer task-relevant state transitions directly from the visual input. However, for tasks requiring entirely novel motor skills—such as opening a cabinet—fine-tuning only the encoder is insufficient. The skills this method can acquire are predominantly novel combinations of existing atomic actions. The added videos primarily supply the execution sequence for a task; mastering completely unseen actions still requires explicit action demonstrations.

DreamFree repurposes a pre-trained video diffusion model into a predictive visual encoder. By leveraging future representation distillation, the features extracted from a single-frame image encapsulate information about future states. During deployment, the system no longer generates future videos; instead, actions are decoded directly from this predictive feature layer. Both simulation and real-robot experiments validate this control performance, while also achieving a substantial reduction in inference latency.

The team's core conclusion centers on predictive representation learning as a viable pathway to translate spatiotemporal knowledge from pre-trained videos into robotic control capabilities. Scaling up the video foundation model and incorporating more video data can further enhance control performance. Meanwhile, video demonstrations lacking explicit action labels can also be utilized to expand the range of novel tasks the model can execute.