Unified World Models: Coupling Video and Action Diffusion for Pretraining on Large Robotic Datasets

1. Quick overview of the paper

Difficulty rating: ★★★★☆. Requires understanding of DDPM, conditional diffusion, Diffusion Policy, world model, inverse dynamics, latent diffusion and robotic pre-training/fine-tuning experimental protocols.

Keywords: Unified World ModelAction DiffusionVideo DiffusionIndependent TimestepsAction-Free Videos

Core contribution list

• Unify policy and world modeling.UWM uses a model to cover $p(a|o)$, $p(o'|o, a)$, $p(a|o, o')$, and $p(o'|o)$ to connect the paradigms of imitation learning and world model.
• Independent diffusion timesteps.Different from PAD's shared timestep, UWM independently samples/sets $t_a$ and $t_{o'}$ to achieve conditioning or marginalization of variables.
• Take advantage of action-free video.For the non-action video $t_a=T$, use Gaussian noise to compensate for the action, and still use the same training objective so that the video data can enter policy pretraining.
• Real and simulated verification.The authors validate the design on five real-world tasks of DROID setup, LIBERO OOD task, forward/inverse dynamics, OOD classification evaluation, from-scratch vs pretraining and Internet video cotraining.

Figure 1. UWM can switch between policy, forward dynamics, inverse dynamics and video prediction by controlling action/video diffusion timesteps.

2. Motivation

2.1 Problems to be solved

Behavioral cloning and Diffusion Policy have demonstrated that reliable manipulation policies can be learned as long as they are demonstrated by high-quality experts. But robot data collection is expensive, and policies usually only learn the mapping from observations to actions, without explicitly modeling "how actions will change future observations." On the other hand, world models and video models can learn dynamics, but videos often do not have action labels, making it difficult to directly use them to train policy.

2.2 Limitations of existing methods

• General imitation learning: Relying on action-labeled expert demonstrations, it is difficult to use action-free videos; it is easily fragile under OOD conditions.
• Separate world model: It is possible to predict future images, but how this translates into better action policies is not straightforward.
• PAD class joint video-action diffusion: Sharing the same timestep to generate action and video can only be sampled from joint distribution and lacks flexible marginal/conditional inference.

2.3 Solution ideas of this article

The author's key insight is to treat diffusion timestep as a continuous version of masking: the closer $t$ is to $T$, the closer the variable appears to be completely masked; the closer $t$ is to $0$, the closer the variable is to clean input. Therefore, as long as action and future observation are given independent timesteps, a certain variable can be "masked" for marginalization during reasoning, or fixed to a clean value for conditioning.

3. Related work context

4. Detailed explanation of method

4.1 Problem setting

There are two types of data: expert data with motion $\mathcal{D}_e=\{(o_i, a_i, o'_i)\}_{i=1}^N$, and video data without motion $\mathcal{D}_{af}=\{(o_i, o_{i+1})\}_{i=1}^M$. The model expects four distributions from the same set of mechanisms:

4.2 Coupled Video-Action Diffusion

UWM's score/noise prediction network inputs the current observation $o$, the noisy action $a_{t_a}$, the next noisy observation $o'_{t_{o'}}$, the action timestep $t_a$ and the next observation timestep $t_{o'}$, and outputs two noise predictions $\epsilon_a^\theta, \epsilon_{o'}^\theta$.

Figure 2. UWM training and inference process: robot data is randomly sampled $t_a, t_{o'}$; video data is fixed $t_a=T$; timestep is used to control marginal/conditional inference during inference.

4.3 Four reasoning modes

4.4 Architecture

UWM uses diffusion transformer + AdaLN conditioning. The current observation is encoded by ResNet-18; $t_a, t_{o'}$ is encoded with sinusoidal timestep encoder; image diffusion uses latent diffusion. The original $224\times224\times3$ image is encoded into $28\times28\times4$ latent by frozen SDXL VAE, and then cut into patch embeddings using the $(4, 4, 2)$ spatiotemporal patchifier. Each step of the action chunk is encoded with shallow MLP, concatenated with image patch embeddings and learnable register tokens and then input into the transformer.

Figure 3. UWM block: AdaLN injects the current observation and two timesteps; register tokens are used for cross-modal information exchange.

Figure 4. Forward dynamics visualization: given the current observation and action, the predicted next observation is close to the true next observation.

4.5 Training and implementation details

Appendix Implementation Key hyperparameters are given: observation length $h_o=2$, action length $h_a=16$, rollout length $h'_a=8$, embedded dim 768, 12-layer transformer, 12 heads, 8 registers; training diffusion steps 100, inference steps 10, DDIM sampler; batch size is 36×4 pretraining, 36×2 finetuning; AdamW, learning rate $10^{-4}$, weight decay $10^{-6}$.

Training 100K steps DROID pretraining uses 4 NVIDIA A100s, about 24 hours. When deploying, execute the replan after predicting the first $h'_a=8$ actions.

5. Experiments and results

5.1 Experimental setup

Real robotic experiments using Franka Panda and DROID manipulation platform. The observation includes two scene cameras and a wrist camera, plus an overhead evaluation camera for alignment initialization; the control frequency is 10 Hz; the action space is delta end-effector pose plus continuous gripper state.

Figure 5. Real robot pre-training/fine-tuning data: 2000 tracks with motion in DROID are used for pretraining, and the other 2000 without motion are used as cotraining videos.

Figure 6. Real robot camera setup: two scene cameras, a wrist camera, and an evaluation camera.

The five real tasks are Stack-Bowls, Block-Cabinet, Paper-Towel, Hang-Towel, and Rice-Cooker. The first four tasks evaluate 50 initializations, while Rice-Cooker only evaluates 20 initializations close to the data distribution due to its high difficulty; each initialization gives each method 3 attempts.

Figure 7. Start/end/OOD settings for five real tasks. OOD mainly adds visual disturbances such as unseen interference objects and opening drawers/cabinet doors.

5.2 Real Robot Results

Figure 8. Average success rate of real robots. UWM pretraining is already better than the baseline, and action-free video cotraining improves further.

The paper explains the real results as follows: UWM learns the causal relationship between actions and image observations through action/video feature sharing and training of different conditional/marginal distributions. GR1 is a strong regression baseline, but video cotraining sometimes dilutes the action learning signal; the poor performance of PAD is attributed by the author to raw-pixel concatenation conditioning which is difficult under the same model capacity.

5.3 LIBERO Simulation

The simulation experiment uses 4500 trajectories of LIBERO-100 and LIBERO-90 for pre-training. Five tasks in LIBERO-10 are randomly selected for evaluation, and each task is fine-tuned with 50 expert demos. The author expanded the object initialization range to 0.03 and removed the background objects to create OOD.

Figure 9. LIBERO pre-training and fine-tuning tasks: Book-Caddy, Soup-Cheese, Bowl-Drawer, Moka-Moka, Mug-Mug.

5.4 Analysis and Ablations

Figure 10. Classification OOD: lighting, background, clutter.

Figure 11. From scratch vs pretrained: UWM and DP are similar in training from scratch, but UWM benefits more from pretraining.

Figure 12. Robot perspective: camera observations under ID, standard OOD, and categorized OOD conditions.

Figure 13. Internet video dataset from Kinetics-400 and Something-Something-v2.

Figure 14. Real robot evaluation tracker: Overlaying reference initial frame and current frame to systematically align random initialization.

6. Analysis and discussion within the paper

6.1 Core explanation given by the author

• Additional supervision comes from the same data.UWM not only learns actions from demonstrations, but also predicts future observations, thus obtaining more dynamic supervision from the same trajectory.
• Independent timesteps allow the model to learn causal relationships.$t_a, t_{o'}$ is randomly combined during training. The model needs to recover actions and future observations under different noise/occlusion conditions, so it is exposed to different marginal/conditional distributions.
• A natural entry into video cotraining.For non-action videos, the action timestep only needs to be fixed to $T$, and no additional architecture or special pseudo-action annotation is required.
• The real-world improvement over simulation is even more obvious.The authors believe that current simulation dynamics are simpler and therefore the OOD improvement is smaller than in the real world.

6.2 Failure modes

Appendix Real-World Failure Modes Point out: Multiple cameras may still have difficult viewing angles where certain objects are only seen by one camera; the behavior of the object itself can also cause failure. For example, the Paper-Towel will fall over after being placed on the wooden platform at a bad angle, and the baseline in Stack-Bowls will confuse the blue bowl with distractors after picking up the pink bowl.

7. Analysis, Limitations and Boundaries

7.1 The most valuable part of this paper

The most valuable thing is to combine "learning dynamics with videos" and "learning actions with demonstrations" into a controllable continuous diffusion probability model, rather than adding an independent video predictor. Through independent $t_a, t_{o'}$, UWM allows policy, forward dynamics, inverse dynamics, and video prediction to be different inference slices of the same training target. This directly explains why action-free video can enter training: the missing action is a complete noise variable of $t_a=T$ in a diffusion sense.

7.2 Why the results hold up

The evidence chain is relatively complete: real robot tasks cover rigid bodies, deformable objects and long-term tasks; each task has ID/OOD; DROID robot video cotraining and Internet video cotraining are independently verified; LIBERO provides a standard simulation benchmark; forward/inverse dynamics shows that UWM does not only improve policy scores, but can indeed work in different inference modes. Ablation further demonstrates that future obs reconstruction, registers, and AdaLN conditioning all contribute quantitatively.

7.3 Author's statement of limitations

• Human video at scale has not really been exploited yet.The authors explicitly state that the embodiment gap has not been resolved; the Internet videos experiments are only preliminary validation of the Kinetics-400/SSV2 hybrid.
• Forward dynamics reconstruction has artifacts.The author believes that these artifacts may reduce the effectiveness of using models for planning, and new advances in the field of generative models need to be absorbed.
• Dense video prediction is needed.The paper predicts that denser video prediction may further improve performance.
• The real mission is still a controlled platform.DROID setup, Franka Panda, 5 tasks and a limited number of demos support the conclusion, but do not mean that large-scale generalization across embodiments has been proven.

7.4 Applicable boundaries

UWM is suitable for scenarios where you have action-labeled robot trajectories and want to take advantage of additional action-free robot/video data at the same time. It assumes that observation/action/future observations can be modeled by continuous diffusion, and that future images can serve as useful dynamic supervision. The paper does not yet provide sufficient evidence for long-term planning, multi-robot shape transfer, high-precision contact force control, or long-span transfer of human video to robot actions.

8. Reproducibility Audit