PaperlayerView exact revision on arXiv

Source-linked AI summary

World Model Self-Distillation: Training World Models to Solve General Tasks

Sebastian Stapf, Pablo Acuaviva Huertos, Aram Davtyan, Paolo Favaro

arXiv:2606.12072v1cs.CV

TL;DR

Pretrained video world models rely on detailed textual execution descriptions, limiting direct task solving without costly paired task-execution videos. WMSD distills caption-guided behavior into an instruction-conditioned Executor and improves it with VLM-feedback reinforcement learning, consistently improving task-solving ability and surpassing its Demonstrator under VLM-based evaluation.

  • Problem

    Pretrained video world models require detailed textual descriptions and external reasoning, while supervised task-execution videos are costly to collect at scale.

  • Method

    WMSD self-distills a detailed-description-conditioned Demonstrator into an instruction-conditioned Executor and adds VLM task-level feedback through reinforcement learning.

  • Results

    WMSD consistently improves task-solving ability across settings, surpasses its Demonstrator under VLM-based evaluation, and remains competitive with supervised fine-tuning.

  • Takeaways & Limitations

    The framework converts detailed execution knowledge into a compact instruction-following interface without curated task-execution videos while maintaining efficient inference.

  • Takeaways & Limitations

    In the data-free setting, WMSD cannot recover accurate robot-specific dynamics without corresponding robot data beyond the initial frame.

Abstract

from arXiv · show

Pretrained video generators are promising visual world models that exhibit emergent task-solving abilities; however, their reliance on detailed textual descriptions limits their direct use for planning and decision-making. Existing approaches either outsource this reasoning to language or vision-language models, or rely on supervised fine-tuning with paired task-execution videos, which are costly to collect and difficult to scale. We propose a scalable framework that elicits task-solving ability in such models by combining self-distillation with reinforcement learning. Given an unlabeled scene image, a vision-language model generates a candidate task and a detailed step-by-step solution. The solution conditions a pretrained video diffusion model, the Demonstrator; we distill its behavior into an Executor conditioned only on the image and a short task prompt. This transfers execution knowledge from caption-guided generation to instruction-conditioned task solving without curated task-video supervision. We further improve the Executor with reinforcement learning from VLM feedback, exploiting the asymmetry between judging whether a sampled video satisfies a task and generating the solution. Experiments on our proposed WorldTasks-Benchmark and the DreamGen robotics benchmark show that the Executor surpasses the Demonstrator under our VLM-based evaluation protocol and transfers competitively to robotic tasks.

1 Introduction

Pretrained video world models show emergent capabilities but rely on detailed textual solutions and external reasoning, limiting direct task solving. WMSD addresses this through self-distillation from a caption-guided Demonstrator to an instruction-conditioned Executor, stabilized with VLM feedback and teacher-based regularization.

  • Motivation: Pretrained text- or image-to-video generators require detailed scene or action descriptions and external language or vision-language reasoning instead of autonomously inferring task execution.This textual-conditioning dependency limits their direct applicability to task solving.
  • Motivation: Supervised fine-tuning requires large, diverse collections of successful task-execution videos spanning environments, objects, and task abstractions.Such demonstrations are costly to acquire, particularly for long-horizon tasks and fine-grained object interactions.
  • Method: Reinforcement learning lets the model sample candidate solutions and improve generations that satisfy tasks, but video-domain computation makes this strategy difficult to apply directly.Few-step distillation and student-sample objectives motivate iterative improvement without paired real videos at every update.
  • Method: WMSD trains an instruction-conditioned Executor to match a detailed-description-conditioned Demonstrator, transferring execution knowledge without paired task-execution videos.The Executor receives a high-level task instruction and initial observation, while the Demonstrator receives a detailed execution description.
  • Method: WMSD combines VLM feedback as a weak verification signal with distributional regularization from the teacher, stabilizing self-distillation for task-solving video models.The approach exploits the asymmetry between generating a valid solution and checking a proposed one, while addressing noisy and inconsistent raw VLM rewards.
  • Contributions: The proposed method turns pretrained caption-conditioned video diffusion models into instruction-conditioned task solvers and enables the Executor to surpass its teacher under VLM-based evaluation.The contribution also reports competitiveness with methods trained using curated task-specific supervision.

2 Related Work

Related work spans task-conditioned world models, self-distillation and distribution matching, and reinforcement learning for diffusion or flow generators. This work distinguishes itself by combining VLM task rewards with Demonstrator-derived distillation rewards and an anchor loss.

  • Task-Conditioned World Models: Task-conditioned world models use language, actions, task specifications, or demonstrations, while other planners combine video generation with language or vision-language models.
  • Self-Distillation and Distribution Matching: Knowledge distillation and self-distillation train compact or improved students from teacher predictions, with on-policy refinement and distribution matching addressing shift, efficiency, and stability.
  • Reinforcement Learning for Flow Models: Recent reinforcement-learning methods adapt policy optimization to diffusion and flow generators, using flow-specific refinements to improve training stability.
  • Reinforcement Learning for Flow Models: Unlike reward-only alignment, the proposed approach combines VLM task rewards with Demonstrator-derived distillation rewards and an anchor loss to improve task success while regularizing visual dynamics.

3 Method

The method distills detailed, caption-guided execution from a fixed Demonstrator into an instruction-conditioned Executor, then combines on-policy distillation with VLM-based reinforcement learning and teacher anchoring. This transfers execution knowledge while allowing task feedback to improve success beyond imitation.

  • Setup: The Executor conditions on an initial observation and short task instruction, while the fixed Demonstrator receives a richer detailed execution description.The student uses cE = (I, T), whereas the teacher uses cD = (I, D).
  • Distillation: Off-policy distillation is stable but matches the student only on teacher trajectories, allowing errors to compound during student rollouts.On-policy distillation instead evaluates teacher–student discrepancy on trajectories sampled from the student.
  • Distillation: On-policy distillation combines a policy-gradient term that favors low-discrepancy trajectories with direct velocity-field regression on student states.Under shared initial noise, matching teacher velocity on student states bounds student–teacher trajectory drift.
  • Reinforcement learning for task solving: VLM feedback supplies task-level reward indicating whether a generated video solves the instruction, while distillation reward controls agreement with the Demonstrator.The two objectives are weighted by λtask > 0 and λdistill > 0.
  • Optimization objective: The full objective combines policy-gradient reward optimization with a teacher anchor that keeps selected student velocities close to the teacher without backpropagating through every sampler step.The anchor uses sampled states treated as fixed, and the reward term selects trajectories.
  • Optimization objective: Self-distillation transfers execution knowledge, reinforcement learning improves task success, and anchoring prevents uncontrolled drift while allowing the Executor to surpass the Demonstrator.This claim is stated under the chosen reward and evaluation protocol.

4 Experiments

Experiments show that on-policy world-model self-distillation combined with reinforcement learning improves task solving beyond off-policy distillation, standard RL, and baseline models. WMSD generalizes to robotic and out-of-distribution tasks, while remaining limited by data-free modeling of robot-specific dynamics and unstable demonstrator failures.

  • Self-distillation variants: On-policy self-distillation continues improving after approximately 60 training steps, surpassing off-policy self-distillation, with the full objective achieving the best overall performance.Off-policy self-distillation saturates, whereas both on-policy variants improve on both reported metrics.
  • Self-distillation and reinforcement learning: Combining on-policy self-distillation with RL substantially improves task solving and enables the student to surpass the demonstrator under VLM-based evaluation.Standard RL alone performs comparably only through approximately 50 steps before plateauing.
  • Baseline comparisons: After 25 training steps, WMSD improves all reported metrics on HY1.5 and produces larger task-completion, agent-correctness, and physical-consistency gains on distilled LTX-2 than SFT.SFT provides little to no improvement and sometimes degrades performance.
  • WorldTasks-Bench breakdown: Navigation completion rises from 31.1% to 75.6%, while object-interaction completion rises from 17.6% to 55.9% under World-Model Self-Distillation.Perception, positioning, and combat also improve, alongside stronger grounding for first-person prompts.
  • Generalizability: WMSD improves WorldTasks-Bench and robotics-related performance competitively with supervised fine-tuning, while leveraging few-step video models for more efficient RL optimization.In a data-free regime, WMSD is comparable to SFT-trained Cosmos and substantially improves over LTX-2.
  • Limitations: Data-free WMSD cannot recover accurate robot-specific dynamics, and out-of-distribution gains diminish when the Demonstrator produces incoherent solutions, introducing additional instability.The limitation concerns platform-specific appearance and motion beyond the initial frame.

5 Conclusion

WMSD consistently improves task-solving ability, agent correctness, and physical consistency across diverse settings. Its key contribution is transferring detailed caption-guided execution knowledge into a compact instruction-following interface without curated task-execution videos, with on-policy self-distillation and reinforcement learning proving especially effective.

  • WMSD consistently improves task-solving ability, agent correctness, and physical consistency across a wide range of settings.
  • The framework converts detailed execution knowledge from caption-guided video generation into a compact instruction-following interface.
  • WMSD performs this transfer without requiring curated task-execution videos.
  • Combining on-policy self-distillation with reinforcement learning proves especially effective.

A Technical appendices and supplementary material · A.1 Further Implementation Details

The supplementary implementation details report the hyperparameters used to self-distill LTX-2 and HunyuanVideo-1.5 in the Sec. 4 experiments.

  • A.1 Further Implementation Details: Tab. 3 reports the hyperparameters used for self-distilling LTX-2 and HunyuanVideo-1.5 in the Sec. 4 experiments.The passage provides no further hyperparameter values or configurations.

A.2 Compute Resources

The primary results used a 128-GH200-GPU cluster, while ablations used 16 GH200 GPUs for 12 hours.

  • Compute Resources: 128 GH200 GPUs supported the primary results in Sec. 4.4.The primary experiments ran on a large-scale cluster.
  • Compute Resources: 16 GH200 GPUs over 12 hours supported the ablation studies in Sec. 4.6.The ablations used a smaller computational setup.

A.2.1 Distribution-Matching Self-Distillation

This section investigates Distribution Matching Distillation as an alternative on-policy objective for demonstrator–executor self-distillation. Although adapted by minimizing KL divergence between executor and demonstrator distributions, the objective consistently diverged and was excluded from the final method.

  • Motivation: DMD was evaluated as an alternative on-policy distillation objective combined with reinforcement learning.The investigation followed prior work [33].
  • Distribution Matching Distillation: DMD trains a one-step generator to match a pretrained diffusion teacher’s sample distribution rather than reproduce its full denoising trajectory.The approach is motivated by the computational expense of multi-step diffusion sampling.
  • Adaptation: The objective was adapted by minimizing KL divergence between executor distribution pθ(xt, t | cE) and demonstrator distribution pθ′(xt, t | cD).The executor parameters were updated using an approximate distribution-matching gradient.
  • Results: DMD-based self-distillation consistently diverged across experiments and was not used in the final method.The authors describe the objective as conceptually appealing but unstable in practice.

A.2.2 Reward Prompts

The reward prompts used during training provide binary supervision for task success and visual-temporal consistency, enabling stable reward computation from generated videos. They are shown in Fig. 10 and Fig. 11.

  • Reward prompts provide binary supervision for task success during training.
  • Reward prompts also supervise visual-temporal consistency in generated videos.
  • These prompts enable stable reward computation from generated videos and are shown in Fig. 10 and Fig. 11.

A.3 Further Experiments

Further experiments found negligible differences between per-step and trajectory-level distillation rewards, while sharing Executor and Demonstrator weights destabilized training. Consequently, the main experiments use trajectory-level rewards and a fixed Demonstrator.

  • Reward design: Per-step and trajectory-level distillation rewards produced only minor differences in final performance, so the trajectory-level form was used for simplicity.The trajectory-level reward is defined in Eq. 9.
  • Weight sharing: Sharing Executor weights with the Demonstrator caused unstable training across hyperparameter settings and EMA schedules, so the main results use a fixed Demonstrator.The instability was observed across the investigated training configurations.

A.4 Evaluation Denominators … A.10 Method Derivations and Proofs

The appendix documents evaluation validity, out-of-distribution VBVR testing, dataset preparation, benchmark prompts, alternating optimization, and the theoretical foundations and proofs of WMSD. It emphasizes explicit failure accounting and reinforcement-learning formulations designed to improve beyond teacher behavior while stabilizing flow-matching training.

  • A.4 Evaluation Denominators: 3.0% is the largest observed failure rate in WorldTasks-Bench evaluation, with main-text scores computed only over valid VLM judgments.Tab. 4 reports success counts, valid denominators, and failure rates, making discarded data explicit.
  • A.5 VBVR Evaluation; A.6 Further Details on WorldTasks: VBVR evaluation compares the vanilla model, WMSD-trained model, and a prompt-rewrite variant that converts longer benchmark queries into shorter task instructions.VBVR tasks are longer and more abstract than WorldTasks instructions, making them substantially out of distribution.
  • A.6.1 Dataset Filtering: Dataset filtering removes incomplete, processed, blurry, underexposed, overexposed, nearly empty, and semantically unsuitable images before downstream task-conditioned video generation.The pipeline combines Laplacian variance, luminance, near-black and near-white pixel fractions, CLIP-based aesthetic ranking, and VLM screening.
  • A.6.2 Example Task and Solution Prompts from WorldTasks: WorldTasks examples pair an initial frame with two task prompts and their corresponding descriptive solution prompts.Representative samples are presented in Figures 13 and 14.
  • A.7 Evaluation Prompts for WorldTasks-Bench: WorldTasks-Bench evaluates generated videos for end-state task completion, correct agent attribution, and physical realism with temporal consistency.The three prompts separately assess whether the task succeeds, whether the intended actor performs it, and whether execution is physically plausible.
  • A.8 Alternating Training Algorithm; A.9 Theoretical Background: The training procedure alternates between Executor and optional Demonstrator rounds, with the Demonstrator fixed when the alternation period N equals 0.Executor updates combine task and distillation rewards with a Demonstrator anchor, while optional Demonstrator updates use task rewards and a base-model anchor.
  • A.9 Theoretical Background: Group-relative optimization reinforces trajectories that outperform peers and can improve beyond the teacher, while AWM aligns reinforcement learning with flow matching for lower variance and improved training efficiency.AWM uses advantage weighting and a velocity-space KL regularizer to stabilize updates against a reference model.
  • A.10 Method Derivations and Proofs: The derivations and proofs establish the gradient decomposition and Proposition 1 using shared initialization, Lipschitzness, Cauchy–Schwarz, and a coupling between terminal laws.The appendix explicitly derives results from Sec. 3 rather than introducing a separate empirical evaluation.
Loading 2606.12072v1…