SCALE: State-Calibrated Latent Embeddings for JEPA Planning in the Right Geometry

2026-08-17

AI Summary

This study addresses the failure of state guidance in JEPA planning caused by suboptimal representational geometry. We propose SCALE, a method that aligns the latent space with the task state space via pairwise distance correlation regularization to achieve state-calibrated embeddings, thereby enabling state information to effectively dominate planning cost computation. Our findings reveal that planning performance depends critically on representational geometry rather than merely information presence. Experiments demonstrate that SCALE consistently outperforms LeWM across five tasks and multiple solvers without additional inference overhead. Consequently, this work establishes a novel paradigm for lightweight optimization in JEPA-based planning by ensuring geometrically faithful state representations.

Abstract

Joint-embedding predictive world models plan by scoring predicted terminal embeddings against a goal embedding using a cost defined on the representation itself. Two prominent strategies for obtaining non-collapsed representations are to inherit a pretrained feature space, as in DINO-WM, and to learn an embedding end to end with anti-collapse regularization, as in LeWorldModel (LeWM) with SIGReg. These strategies show complementary strengths across tasks. Although task-relevant state is decodable from the full embeddings of both models, DINO-WM's leading principal components usually retain substantially more state information than LeWM's. Because Euclidean planning costs are dominated by high-variance directions, this difference affects how strongly state can influence candidate selection. We propose SCALE (State-CAlibrated Latent Embeddings) to give the end-to-end LeWM representation the favorable geometric property observed in DINO-WM. SCALE induces this property by correlating sampled pairwise latent distances with distances in a standardized task-relevant state space, without replacing LeWM's learned encoder. Across five tasks, three planning solvers, and five compute budgets, SCALE improves every task--solver average over LeWM. A latent-to-state regression control matches or exceeds SCALE's full-embedding decodability yet leaves latent--state distance alignment essentially unchanged and yields less consistent planning gains. SCALE adds a single lightweight training-time regularizer and no planning-time overhead. These results show that planning depends not only on whether task-relevant information is present, but also on whether it shapes the geometry consumed by the planner.

Problem

• JEPA Planning
• Latent Representation Geometry
• State-Calibrated Embeddings
• World Models
• Representation Learning

Innovation

• Joint-embedding predictive architecture
• Latent space geometry
• State-calibrated embeddings
• World models
• Planning