Future Representation Distillation for Imagination-free World Action Models

Future Representation Distillation for Imagination-free World Action Models. By Jaehyeong Jo and seven others, from KAIST. Published on the fourth of November.

World Action Models, or WAMs, draw on the predictive knowledge inside pretrained video diffusion models to control robots. They do it by predicting future video. That route has two weaknesses. Iterative video denoising is expensive, and when the visual prediction is wrong, the action predicted from it can be compromised as well.

The authors want the predictive knowledge without the imagined footage. Their key idea is to expose that knowledge through visual representations computed from the current observation alone, so that actions can be predicted without generating any future video. The model built on this is DreamFree, a WAM that repurposes a pretrained video diffusion model as a predictive visual encoder, by way of representation distillation.

Training has two parts. A visual encoder learns to predict the representation of a frozen video diffusion model, while seeing only the current observation. After that, an action module is trained to predict actions by extracting the action-relevant information from that visual encoder. The diffusion model acts as a teacher whose internal view of the future is distilled, rather than as a generator run at control time.

On real-world and simulated tasks, DreamFree reaches higher success rates than state-of-the-art VLAs and WAMs. It does so while running real-time closed-loop control at eleven hertz with a fourteen billion parameter video diffusion backbone. Further experiments show it can turn stronger video priors and additional video data into better robot control. It can also learn new robot tasks from video demonstrations, including videos from a different embodiment, without action annotation for the new tasks.

Imagination-free is the operative word in the title. The video model's sense of what comes next is kept, but it is carried in a representation read off the present frame instead of a rendered future. That is what lets a fourteen billion parameter backbone sit inside a real-time control loop, and why plain video, even of a different body, can still teach the robot something.