💻 computer science

How Well Do Latent World Models Understand Partially Observable Safety Constraints?

This paper investigates how partial observability in latent world models leads to safety failures through estimation and prediction gaps, proposing mutual-information and rollout-based diagnostics alongside multimodal supervision and conformal risk calibration to mitigate these issues in robotic control tasks.

Original authors: Matthew Kim, Kensuke Nakamura, Andrea Bajcsy

Published 2026-06-09

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Matthew Kim, Kensuke Nakamura, Andrea Bajcsy

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are teaching a robot to cook. You want it to melt wax without burning it, or pour rice without spilling it. To do this, the robot uses a "Latent World Model." Think of this model as the robot's internal daydream. It looks at what it sees (like a camera image), creates a simplified mental picture of the world, and then imagines, "If I move my arm this way, what will happen next?"

The paper asks a critical question: What happens if the robot's daydream is missing crucial details?

The authors found that when a robot only sees the world through a single lens (like a standard RGB camera), its internal daydream often suffers from two specific types of "blind spots." These blind spots cause the robot to make dangerous mistakes.

The Two Blind Spots

1. The "Invisible Thermometer" (Estimation Gaps)

The Scenario: The robot is heating wax. The safety rule is "Don't let it get too hot."
The Problem: The robot only has a regular camera. It can see the wax, but it cannot see the temperature. In its daydream, the wax looks fine, even though it's actually burning. The robot's internal state is missing the "temperature" variable entirely because the camera didn't provide that data.
The Metaphor: Imagine trying to guess if a cup of coffee is hot just by looking at a black-and-white photo of it. You can't tell. If you try to drink it based on that photo, you might burn your tongue. The robot is doing the same thing; it's "blind" to the heat.

2. The "Surprise Spill" (Prediction Gaps)

The Scenario: The robot is pouring rice from a bottle. The safety rule is "Don't spill."
The Problem: The robot can see the rice after it spills. However, it cannot tell if the bottle is full or empty before it starts pouring. If the bottle is empty, pouring is safe. If it's full, pouring causes a mess. Because the robot can't see inside the opaque bottle, its daydream can't predict the disaster before it happens.
The Metaphor: Imagine playing a game of "Guess Who" where you can only see the person's face, not their body. If you try to guess if they are holding a heavy box, you might guess wrong. The robot sees the spill happening, but it couldn't predict it because it didn't know the "hidden" state (full vs. empty bottle).

How the Authors Diagnosed the Problem

The researchers didn't just guess; they built two "check-up tools" to measure how well the robot's daydream understood safety:

The "Clue Detector" (Mutual Information): They measured how much a single camera image actually tells the robot about safety. They found that a regular camera image tells the robot very little about temperature (low "clue" score), but adding an infrared camera (which sees heat) gives the robot a huge amount of useful information.
The "Future Simulator" (Rollout Metric): They asked the robot to imagine 16 steps into the future. They found that when the robot lacked the right sensors, its imagination was often wrong. It would imagine a safe path when the path was actually dangerous.

The Solutions: Fixing the Blind Spots

The paper proposes two clever fixes, one for each problem, so the robot can be safe even if it only has a regular camera during the actual task.

Fix for the "Invisible Thermometer": Privileged Supervision

The Idea: You can't give the robot an infrared camera while it's cooking (maybe it's too expensive or bulky), but you can use one while training it.
How it works: During training, the robot looks at both the regular video and the heat map. The researchers force the robot's "daydream" to learn how to predict the heat map just by looking at the regular video.
The Result: The robot learns to "imagine" the temperature even when it only sees the regular video later. It becomes a better guesser. In the experiments, this allowed the robot to lift the pan and stop the wax from burning 100% of the time, whereas the robot without this training failed 85% of the time.

Fix for the "Surprise Spill": Calibrating the Confidence

The Idea: If the robot can't predict the future perfectly, it shouldn't be too confident in its guesses.
How it works: The researchers used a statistical method (Conformal Calibration) to adjust the robot's "risk threshold." They told the robot: "If you aren't 100% sure it's safe, assume it's dangerous and stop."
The Result: The robot became more cautious. When it couldn't tell if a bottle was full or empty, it stopped pouring. This prevented spills. The trade-off is that the robot became a bit more conservative (it might stop pouring even when it was safe), but it was much safer overall.

The Bottom Line

The paper concludes that just having a good video feed isn't enough for a robot to be safe. The robot's internal mental model needs to contain the specific information required to make safe decisions.

If the information is missing (like temperature), the robot needs to be trained with extra sensors so it can learn to "imagine" that missing info.
If the information is unpredictable (like a full bottle), the robot needs to be told to be more cautious and less optimistic.

By using these strategies, the researchers showed that robots can learn to cook safely, even when they are only looking at the world through a standard camera.

Technical Summary: How Well Do Latent World Models Understand Partially Observable Safety Constraints?

Problem Statement

Latent world models (WMs) offer a promising alternative to hand-specified dynamics for learning state representations and control policies directly from high-dimensional observations (e.g., RGB images). However, their efficacy in safety-critical scenarios depends on whether the learned latent state encodes all information necessary to satisfy safety constraints.

This paper investigates the limitations of latent WMs under partial observability, where training data lacks sufficient information to determine safety-critical quantities. The authors identify a critical failure mode distinct from classical state estimation: unknown unknowns. Unlike classical settings where the state space is well-defined (allowing for uncertainty quantification via filters like Kalman filters), latent WMs infer the state representation itself from data. If relevant safety variables are missing from observations, they may be entirely absent from the learned latent space, leading to unsafe control behaviors.

The authors define two specific failure modes:

Estimation Gaps: Current observations are insufficient to determine if a safety constraint is currently being violated (e.g., inferring temperature from RGB images alone).
Prediction Gaps: Failures are observable once they occur, but cannot be reliably anticipated from current observations due to missing information (e.g., predicting a spill without knowing if a bottle is filled).

Methodology

Diagnostics for Observability Gaps

To quantify these gaps, the authors propose two offline diagnostic metrics:

Mutual Information (MI) for Safety Observability: To measure estimation gaps, they estimate the mutual information $I(Y; X)$ between observations $X$ (e.g., RGB images) and safety labels $Y$ . Using the Barber–Agakov variational lower bound, they train a linear probe on the latent embeddings to quantify how much uncertainty about the safety label is reduced by the observation.
Rollout-Based Prediction Accuracy: To measure prediction gaps, they evaluate the world model's ability to predict future safety outcomes via open-loop rollouts. They simulate 16-step action sequences and measure the classification accuracy of the predicted safety labels, specifically analyzing samples where unimodal and multimodal models disagree.

Mitigation Strategies

The paper proposes distinct strategies for each failure mode:

For Estimation Gaps (Privileged Multimodal Supervision): Instead of requiring multimodal sensors at deployment, the authors train the world model using RGB inputs but enforce multimodal supervision. The model is tasked with reconstructing safety-relevant modalities (e.g., IR thermal images) from the RGB input. This "privileged" supervision during training shapes the latent representation to encode safety-critical quantities (like temperature) even if they are not directly observable at runtime.
For Prediction Gaps (Conformal Risk Calibration): When the latent state cannot reliably predict future failures, the authors propose calibrating the safety controller's threshold. Using conformal prediction, they adjust the intervention threshold ( $\epsilon$ ) based on the model's failure prediction accuracy on a calibration set. This ensures the controller adopts a conservative fallback policy when the model's predictive uncertainty is high, bounding the false acceptance rate of unsafe actions.

Experimental Setup

The authors validate their approach on a Franka Research 3 manipulator across two hardware tasks:

Wax Melting (Estimation Gap): A robot must heat wax without exceeding 72°C. The task relies on inferring temperature. Models are trained on RGB-only vs. RGB+IR data.
Rice Pouring (Prediction Gap): A robot must pour rice from an opaque bottle without spilling. The bottle may be empty or full. The task relies on predicting whether a tilt will cause a spill. Models are trained on RGB-only vs. RGB+Tactile data.

Key Results

Diagnostic Findings

Estimation Gaps: In the Wax Melting task, the mutual information between RGB observations and the safety label (overheating) was significantly lower ( $I/H \approx 0.22$ ) compared to IR observations ( $I/H \approx 0.80$ ). Standard accuracy metrics were less sensitive to this gap, suggesting MI is a more robust diagnostic for safety observability.
Prediction Gaps: In the Rice Pouring task, while both models could identify spills after they occurred, the RGB-only model failed to predict them reliably in open-loop rollouts. The multimodal model achieved significantly higher accuracy on samples where the two models disagreed, indicating that the RGB-only latent space lacked the representational capacity to correlate actions with future safety outcomes.

Mitigation Performance

Privileged Supervision (Wax Melting): A world model trained with RGB inputs but supervised by IR reconstruction (WM_RGB-MM) achieved a 100% lift rate (preventing overheating), matching the performance of the fully multimodal model. In contrast, the RGB-only model failed 85% of the time. Crucially, the WM_RGB-MM model learned a conservative policy, intervening earlier than the multimodal model, demonstrating that the latent state successfully encoded the necessary safety dynamics despite lacking IR input at runtime.
Conformal Calibration (Rice Pouring): Applying conformal calibration to the RGB-only controller adjusted its threshold to account for its inability to predict spills. While this made the RGB controller highly conservative (allowing pouring only 18% of the time for empty bottles vs. 60% for the multimodal model), it successfully bounded the failure rate to the target level ( $\alpha = 0.12$ ). This confirmed that calibration can mitigate prediction gaps by trading off control authority for safety guarantees.

Significance and Claims

The paper challenges the assumption that latent world models inherently capture sufficient information for safe control under partial observability. The authors claim:

Identification of Failure Modes: They explicitly identify and characterize "estimation gaps" and "prediction gaps" as distinct causes of unsafe behavior in latent-space control.
Quantitative Diagnostics: They provide practical, offline tools (Mutual Information and rollout classification) to detect these gaps before deployment.
Mitigation Feasibility: They demonstrate that safety can be maintained even with limited runtime modalities (e.g., RGB-only) by leveraging richer data during training (privileged supervision) or by statistically calibrating the controller's risk tolerance (conformal prediction).
Conservatism as a Trade-off: The results highlight that mitigating observability gaps often requires increased conservativeness. When the latent state cannot fully resolve uncertainty, safe control necessitates earlier intervention or reduced task execution rates.

The work concludes that while world models are powerful, their state representations must be explicitly validated for safety observability, and mitigation strategies are essential when training data does not fully reveal the underlying safety-critical state.