Disentangling the Geometries of Prediction and Planning in Latent World Models through Trajectory-Graph Control Calibration

Joint-Embedding Predictive Architectures (JEPAs) provide an efficient foundation for latent world models by predicting future representations rather than reconstructing pixels. However, Model Predictive Control (MPC) planners built on these models typically select action sequences by minimizing the terminal latent distance to the goal, which implicitly assumes that latent similarity measures the remaining cost to the goal. In contrast, goal-reaching cost can be directional, violating the symmetry imposed by latent-distance planning and potentially causing the planner to prematurely discard dynamically preferable action sequences. To address this limitation, we propose Trajectory-Graph Control Calibration (TGCC), a lightweight post-hoc method that disentangles the geometries of prediction and planning over the same frozen latent representations. Concretely, TGCC constructs a directed graph over offline trajectories and distills its temporal relations into a control geometry network. During planning, the network's graph-derived control score is combined with terminal latent similarity to rank candidate action sequences, thereby defining a planning geometry beyond the prediction geometry. Without retraining the world model, TGCC achieves superior planning performance across multiple datasets and latent world-model backbones, with higher success rates and fewer steps to success.