Video World Models Need Predictable Latents
Organizations: New York University
Abstract
Video world models trained on large-scale robot data promise a safe and scalable approach for learning and evaluating policies that act in the physical world. However, they still struggle with long-horizon rollouts and out-of-distribution generalization, and they remain computationally demanding at inference time. Typically, video world models make predictions in the latent space of a reconstruction-trained autoencoder, which however is not directly optimized for dynamics modeling. On the other hand, world models designed for latent planning usually optimize latents to be predictable, but they forgo accurate prediction of observations in pixel space. In this paper, we show that it is possible to achieve the best of both worlds: a latent geometry optimized to make future prediction as simple as possible without sacrificing accurate prediction of future observations. More precisely, we find that temporal objectives for representation learning, paired with a past-conditioned decoder architecture, lead to consistent improvements in generative video world models at the same training budget. The gains include up to 5x cheaper inference, over 7x faster training, and 15-24% lower pixel and perceptual errors out of distribution and over long horizons. Our world models also agree more closely with real-world policy evaluation when using a vision-language model to judge task success from their pixel-space predictions. Overall, our study indicates that temporal predictability is an important design criterion for the latent space of video world models, and provides concrete guidance on how to incorporate it into the training of autoencoders.
Figures & tables
| ID | OOD | 2-step | Long-horizon | Autoencoder | ||||||
|---|---|---|---|---|---|---|---|---|---|---|
| MSE | LPIPS | MSE | LPIPS | MSE | LPIPS | MSE | LPIPS | Updates | Components | |
| Reconstruction (baseline) | 11.65 | 0.164 | 33.87 | 0.308 | 15.95 | 0.214 | 33.26 | 0.287 | 4.4 | 388 |
| 0.20 | 0.001 | 1.28 | 0.008 | 1.03 | 0.004 | 1.84 | 0.008 | |||
| Pixel prediction | 11.43 | 0.159 | 34.23 | 0.300 | 14.53 | 0.199 | 32.63 | 0.279 | 8.9 | 445 |
| 0.22 | 0.001 | 0.67 | 0.002 | 0.75 | 0.002 | 0.69 | 0.003 | |||
| Contrastive prediction | 11.17 | 0.160 | 31.23 | 0.303 | 12.38 | 0.185 | 31.50 | 0.277 | 96.0 | 293 |
| ID | OOD | 2-step | Long-horizon | Policy eval. | Predictability | ||||||
| MSE | LPIPS | MSE | LPIPS | MSE | LPIPS | MSE | LPIPS | 10 steps | 2 steps | Latent MSE | |
| Ordinary decoder | |||||||||||
| Reconstruction (baseline) | 14.34 | 0.194 | 36.15 | 0.302 | 18.35 | 0.232 | 47.93 | 0.394 | 0.51 | 0.21 | 0.232 |
| 0.16 | 0.002 | 0.98 | 0.003 | 0.65 | 0.004 | 1.51 | 0.002 | 0.10 | 0.07 | 0.008 | |
| Pretrained, no fine-tuning | 14.81 | 0.208 | 34.81 | 0.309 | 19.15 | 0.250 | 46.78 | 0.393 | 0.57 | 0.21 | 0.253 |
| 0.26 | 0.001 | 1.97 | 0.008 | 0.89 | 0.003 | 0.83 | 0.003 | 0.07 | 0.09 | 0.001 | |
Appendix figures & tables33 assets
Supplementary material from the paper’s appendix.
Appendix
| Equal reconstruction (Section 5.1 ) | Equal budget and weight sweep (Sections 5.2 and 5.3 ) | |
| Autoencoder fine-tuning | ||
| Frame pairs | consecutive, or two frames apart with probability 1/3 | consecutive |
| Batch | 24 pairs | 48 pairs |
| Learning rate | ||
| Length | until the reconstruction target of Appendix F | 80,000 updates; final checkpoint |
| Validation pairs | one per held-out video, redrawn at each check | a fixed set, one per held-out video |
| Condition | Clips | Scored prediction steps | Denoising chunk | DDIM steps | Sampling seeds |
|---|---|---|---|---|---|
| In distribution | 2,000 | 1–9 | 8 | 10 | 3 |
| Out of distribution | 479 | 1–29 | 1 | 10 | 3 |
| Sampling-step sweep | 512 | 1–9 | 1 | 2, 4, 6, 8, 10 | 1 |
| Long rollouts | 716 | 1–59 | 8 | 10 | 1 |
| ID | OOD | |||||
|---|---|---|---|---|---|---|
| MSE | LPIPS | SSIM | MSE | LPIPS | SSIM | |
| Reconstruction (baseline) | 11.65 | 0.164 | 0.801 | 33.87 | 0.308 | 0.675 |
| 0.20 | 0.001 | 0.002 | 1.28 | 0.008 | 0.009 | |
| Pixel prediction | 11.43 | 0.159 | 0.806 | 34.23 | 0.300 | 0.683 |
| 0.22 | 0.001 | 0.001 | 0.67 | 0.002 | 0.002 | |
| Contrastive prediction | 11.17 | 0.160 | 0.807 | 31.23 | 0.303 | 0.681 |
| ID | OOD | |||||
| MSE | LPIPS | SSIM | MSE | LPIPS | SSIM | |
| Ordinary decoder | ||||||
| Reconstruction (baseline) | 14.34 | 0.194 | 0.779 | 36.15 | 0.302 | 0.682 |
| 0.16 | 0.002 | 0.001 | 0.98 | 0.003 | 0.004 | |
| Pretrained, no fine-tuning | 14.81 | 0.208 | 0.758 | 34.81 | 0.309 | 0.673 |
| 0.26 | 0.001 | 0.001 | 1.97 | 0.008 | 0.009 | |
| Reconstruction error | ||||
|---|---|---|---|---|
| Past frame | Context disabled | Permuted past frame | Ordinary decoder | |
| 0 | 0.0231 | 0.0407 | 0.0394 | 0.0294 |
| 0.1 | 0.0237 | 0.0430 | 0.0424 | 0.0314 |
| 0.3 | 0.0257 | 0.0496 | 0.0481 | 0.0345 |
| 0.5 | 0.0288 | 0.0700 | 0.0565 | 0.0379 |
| 0.7 | 0.0375 | 0.1947 | 0.1072 | 0.0451 |
| OpenVLA | Octo | RT-1-X | Three policies pooled | |||||||
| 10 sampling steps | ||||||||||
| Ordinary decoder | ||||||||||
| Reconstruction (baseline) | 0.40 | 0.27 | 0.58 | 0.53 | 0.11 | 0.19 | 0.51 | 0.45 | – | – |
| 0.20 | 0.27 | 0.21 | 0.21 | 0.24 | 0.24 | 0.10 | 0.12 | |||
| Pretrained, no fine-tuning | 0.56 | 0.56 | 0.58 | 0.55 | 0.01 | 0.10 | 0.57 | 0.55 | +0.06 | 0.261 |
| Lower than baseline (%) | Lower than all reconstruction-only (%) | Median change (%) | |||||||
| Setting | Steps | Clips | Latent | Inverse | Latent | Inverse | Both | Latent | Inverse |
| ID | 10 | 224 | 92.0 | 93.3 | 80.4 | 79.9 | 69.6 | ||
| ID | 2 | 224 | 96.0 | 98.7 | 89.3 | 92.9 | 87.5 | ||
| OOD | 10 | 479 | 95.0 | 96.7 | 60.5 | 86.2 | 53.2 | ||
| OOD | 2 | 479 | 100.0 | 99.2 | 97.9 | 97.9 | 96.2 | ||
| ID | OOD | |||||
| MSE | LPIPS | SSIM | MSE | LPIPS | SSIM | |
| 200,000 world-model updates | ||||||
| Reconstruction | 12.14 | 0.163 | 0.813 | 34.53 | 0.280 | 0.707 |
| 0.36 | 0.001 | 0.002 | 2.33 | 0.006 | 0.006 | |
| Inverse dynamics | 11.56 | 0.148 | 0.823 | 31.19 | 0.242 | 0.737 |
| 0.02 | 0.001 | 0.001 | 1.42 | 0.005 | 0.005 | |
| OpenVLA | Octo | RT-1-X | Three policies pooled | |||||||
|---|---|---|---|---|---|---|---|---|---|---|
| 10 sampling steps | ||||||||||
| Reconstruction | 0.50 | 0.45 | 0.55 | 0.53 | 0.24 | 0.41 | 0.61 | 0.61 | – | – |
| 0.16 | 0.22 | 0.21 | 0.21 | 0.23 | 0.21 | 0.08 | 0.10 | |||
| Inverse dynamics | 0.54 | 0.47 | 0.57 | 0.53 | 0.34 | 0.30 | 0.66 | 0.65 | +0.05 | 0.081 |
| 0.15 | 0.20 | 0.21 | 0.21 | 0.19 | 0.19 | 0.07 | 0.10 | 0.03 | ||
| Components for 95% | Participation ratio | Turning- angle ratio | Within / across video distance | Spectral slope | High-frequency spectral share | |
| Equal reconstruction, three seeds | ||||||
| Reconstruction | 388 0.9 | 52.7 0.03 | 1.100 0.0001 | 0.423 0.0004 | 0.837 0.0003 | 0.073 0.00008 |
| Latent prediction | 177 8 | 49.4 1.1 | 1.019 0.007 | 0.344 0.005 | 0.907 0.010 | 0.061 0.002 |
| Pixel prediction | 445 0.3 | 64.4 0.2 | 1.110 0.0004 | 0.443 0.001 | 0.763 0.0008 | 0.094 0.0002 |
| Contrastive prediction | 293 6 | 46.7 0.7 | 1.086 0.0009 | 0.384 0.002 | 0.933 0.010 | 0.052 0.002 |
| Equal budget, past-conditioned decoder, three seeds | ||||||
| ID | OOD | |||||
|---|---|---|---|---|---|---|
| MSE | LPIPS | SSIM | MSE | LPIPS | SSIM | |
| Reconstruction | 10.78 | 0.158 | 0.797 | 35.11 | 0.278 | 0.698 |
| Latent prediction, | 10.57 | 0.150 | 0.805 | 30.31 | 0.257 | 0.715 |
| Latent prediction, | 11.67 | 0.163 | 0.790 | 29.68 | 0.271 | 0.704 |
| Latent prediction, | 11.98 | 0.178 | 0.763 | 28.21 | 0.299 | 0.673 |
| 2-step | Long-horizon | |||||
| ID | OOD | |||||
| MSE | LPIPS | SSIM | MSE | LPIPS | SSIM | |
| No KL penalty | ||||||
| Reconstruction | 13.76 | 0.178 | 0.801 | 31.81 | 0.274 | 0.713 |
| Latent prediction | 14.16 | 0.170 | 0.808 | 26.58 | 0.260 | 0.738 |
| KL penalty, weight | ||||||
| Reconstruction | 13.71 | 0.169 | 0.810 | 36.35 | 0.258 | 0.725 |
Video World Models
Video world models generate and predict action-conditioned visual futures to simulate changing scenes for embodied agents.
See all 257 papers