Table of Contents
Fetching ...

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

Zhiting Mei, Tenny Yin, Micah Baker, Ola Shorinwa, Anirudha Majumdar

TL;DR

Controllable video models frequently hallucinate, undermining trustworthy policy evaluation and planning in robotics. The authors introduce C3, a latent-space uncertainty quantification framework that trains video models to produce continuous-scale calibrated confidence estimates and maps them to interpretable pixel-level heatmaps. They leverage proper scoring rules, a latent diffusion transformer-based UQ probe, and latent-to-pixel decoding to achieve dense, well-calibrated uncertainty that correlates with actual errors and detects OOD inputs. Across Bridge and DROID datasets and real-world experiments, C3 yields reliable calibration and interpretable uncertainty localizations, enabling safer decision-making in vision-guided robotics.

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.

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

TL;DR

Controllable video models frequently hallucinate, undermining trustworthy policy evaluation and planning in robotics. The authors introduce C3, a latent-space uncertainty quantification framework that trains video models to produce continuous-scale calibrated confidence estimates and maps them to interpretable pixel-level heatmaps. They leverage proper scoring rules, a latent diffusion transformer-based UQ probe, and latent-to-pixel decoding to achieve dense, well-calibrated uncertainty that correlates with actual errors and detects OOD inputs. Across Bridge and DROID datasets and real-world experiments, C3 yields reliable calibration and interpretable uncertainty localizations, enabling safer decision-making in vision-guided robotics.

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.

Paper Structure

This paper contains 21 sections, 2 theorems, 16 equations, 21 figures.

Key Result

Proposition 1

Given the input actions and video frames, the predicted confidence $\hat{\mathbf{q}}$ provides a calibrated measure of uncertainty of the video diffusion model in the generated video, provided that $\phi$ converges to an optimal solution.

Figures (21)

  • Figure 1: We present $\hbox{$C^3$}$, the first method for training video models that know when they don't know. Using proper scoring rules, $\hbox{$C^3$}$ generates dense confidence predictions at the subpatch (channel) level that are physically interpretable and aligned with observations.
  • Figure 2: Model Architecture. $\hbox{$C^3$}$ enables simultaneous video generation and uncertainty quantification (visualized as a heatmap), quantifying the model's confidence in its accuracy using a UQ probe acting on the video latents.
  • Figure 3: Latent space error. We visualize the latent-space video error in the RGB space, showing the observable range of the errors.
  • Figure 4: (a) Average calibration error. All three architectures have very low ECE, and relatively low MCE. (b) Aggregated reliability diagrams. All methods are well-calibrated, closely tracking the line of perfect calibration.
  • Figure 5: Reliability diagrams for FSC and CS-BC. At the same threshold $\varepsilon_v$ ($0.5$), both models achieve the similarly well-calibrated confidence predictions.
  • ...and 16 more figures

Theorems & Definitions (4)

  • Proposition 1: Uncertainty Decomposition
  • Remark 1: Velocity-space Accuracy
  • Proposition 1: Uncertainty Decomposition
  • proof