Video world models improve prediction with predictable latent spaces
Video World Models Need Predictable Latents
Machine LearningArtificial IntelligenceComputer Vision and Pattern RecognitionRobotics
Summary
Video world models help robots learn from videos to predict what might happen next, but they usually struggle to imagine far into the future or with new situations. The authors show that by training these models to create hidden “latent” spaces that are easier to predict over time, while still being able to accurately recreate video frames, the models become much better. This approach makes training and predicting faster and more reliable, especially for longer periods and unfamiliar scenarios, helping robots plan actions by imagining outcomes more effectively.
What this means in practice
- •For robotics engineers: Improve robot planning by using video world models that predict future observations more accurately and efficiently over long timeframes.
- •For computer vision developers: Create more reliable generative video models that maintain accuracy when dealing with unfamiliar video inputs and save computation during inference.
Authors
Mattia Scardecchia, Sherry Yang
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.