Behavioral Monitoring of JEPA World Models with Jacobian Centroids

2026-09-27

AI Summary

This study addresses the challenge of real-time monitoring for task alignment in world-model-based planning. We propose a behavior monitoring framework based on the Jacobian centroid, which generates saliency maps by computing the geometric structure of the input space to identify behavioral attributes within the internal representations of world models. By integrating the Joint Embedding Predictive Architecture (JEPA) with Jacobian-vector products, this approach reveals the structural decoupling between the encoder and predictor. Crucially, the proposed method enables direct prediction of planning failures prior to action execution and supports goal resampling to accommodate distribution shifts. Evaluated on continuous control tasks, our framework outperforms baseline distribution shift detectors, establishing an effective behavior monitoring stack for robust model-based planning.

Abstract

Detecting failures in World Model (WM)-based planning requires monitoring whether the model is behaviorally aligned with the current task, which in turn requires studying its internal representations. Here, we show that centroids---sub-component Jacobian row-sums---effectively identify the behavioral properties of WMs, complementing traditional activation-based knowledge signals. The centroids of a model are easily computed through Jacobian vector products and characterize how the model organizes the geometry of its input space, yielding an efficient perspective on internal representations, including the generation of task-relevant saliency maps. Evaluated on continuous control tasks using JEPA WMs, this behavioral view reveals a structural dissociation, where the encoder correctly represents the goal while the predictor remains behaviorally unresponsive. This failure mode directly predicts planning failure before any action is taken, allowing for goal resampling to recapture out-of-distribution success. Moreover, centroid-based methods outperform baseline methods as distribution-shift detectors. Together, these tools yield a behavioral monitoring stack that is operational and consequential under distribution shifts.

Problem

• World Models
• Behavioral Monitoring
• Distribution Shift
• Planning Failure
• JEPA

Innovation

• Jacobian Centroids
• JEPA World Models
• Behavioral Monitoring
• Distribution Shift Detection
• Structural Dissociation