World Models That Know When They Don't Know - Controllable Video Generation with Calibrated Uncertainty

TL;DR

This paper introduces C3, a novel uncertainty quantification method for controllable video models that localizes and calibrates confidence estimates at the pixel level, enhancing reliability in video generation tasks.

Contribution

The paper presents a new framework for training video models with calibrated uncertainty estimates in latent space, improving trustworthiness and out-of-distribution detection in controllable video generation.

Findings

Provides dense, pixel-level uncertainty heatmaps for generated videos.

Achieves calibrated uncertainty estimates within training distribution.

Enables effective detection of out-of-distribution video frames.

Abstract

Recent advances in generative video models have led to significant breakthroughs in high-fidelity video synthesis, specifically in controllable video generation where the generated video is conditioned on text and action inputs, e.g., in instruction-guided video editing and world modeling in robotics. Despite these exceptional capabilities, controllable video models often hallucinate - generating future video frames that are misaligned with physical reality - which raises serious concerns in many tasks such as robot policy evaluation and planning. However, state-of-the-art video models lack the ability to assess and express their confidence, impeding hallucination mitigation. To rigorously address this challenge, we propose C3, an uncertainty quantification (UQ) method for training continuous-scale calibrated controllable video models for dense confidence estimation at the subpatch level, precisely localizing the uncertainty in each generated video frame. Our UQ method introduces three core innovations to empower video models to estimate their uncertainty. First, our method develops a novel framework that trains video models for correctness and calibration via strictly proper scoring rules. Second, we estimate the video model's uncertainty in latent space, avoiding training instability and prohibitive training costs associated with pixel-space approaches. Third, we map the dense latent-space uncertainty to interpretable pixel-level uncertainty in the RGB space for intuitive visualization, providing high-resolution uncertainty heatmaps that identify untrustworthy regions. Through extensive experiments on large-scale robot learning datasets (Bridge and DROID) and real-world evaluations, we demonstrate that our method not only provides calibrated uncertainty estimates within the training distribution, but also enables effective out-of-distribution detection.