Inside-JEPA: From Prediction to Relational Grounding in Latent World Models

Learning structured and control-relevant latent representations remains a key challenge for world models. Recent JEPA-based world models learn action-conditioned predictive latent dynamics from observation sequences. However, their forward-prediction objectives do not explicitly enforce reliable identifiability of robot-centric physical state from individual latents or state changes from latent pairs, which can limit downstream planning and policy performance. We propose PSG-JEPA, a physically grounded JEPA world model that shapes its latent space with two complementary grounding objectives beyond forward prediction: grounding individual latents in robot proprioceptive state, and grounding latent pairs in multi-horizon joint-angle changes. Both objectives are applied only during training, leaving the inference architecture and computational cost unchanged. To comprehensively evaluate PSG-JEPA, we conduct experiments at three levels: (1) latent identifiability via probing, (2) goal-conditioned planning on frozen latents, and (3) policy learning in simulation and on a real robot. Experiments demonstrate that our PSG-JEPA consistently outperforms state-of-the-art latent world-model baselines at all three levels.

Latent world models predict the consequences of actions, but accurate prediction does not guarantee that latent distance reflects which candidate will execute successfully. We identify a decision-local prediction gap: among the few futures competing for execution, a candidate predicted closer to the goal can produce a worse realized outcome than an available alternative. We introduce D-JEPA, a decision-aligned latent world model that learns decision-relevant relations among candidate futures from executed outcomes. A bounded, permutation-equivariant operator jointly reasons over goal-relative predictive features and ordinal evidence, refining pretrained predictive geometry where action choices are most consequential. Restricted predictor adaptation and a shared ordinal interface extend this alignment across complementary predictive geometries. D-JEPA further realizes the learned decision structure in JEPA-compatible future representations, enabling deployment through native latent-distance planning. Evaluations across latent control, manipulation, pretrained action-producing models, physical robots and autonomous driving demonstrate improved action selection, including 87.89% success on PushT, a 15.04-point average gain on RoboTwin, and a 17-point gain on physical robot tasks. These results establish decision-relevant relational structure as a direct bridge between predictive world modeling and effective control.

World models used for planning learn to predict how a scene evolves, yet they rarely hold a global view of how far a task has progressed. We argue that such a view should be part of the representation, and that it can be compact: describing a scene takes many dimensions, but where the scene stands on its way to the goal can be summarised by only a few. We propose SD-JEPA, a joint-embedding predictive world model whose latent is split into a small progression subspace and a large content subspace. A temporal triplet loss keeps consecutive frames of a trajectory close in the progression coordinates and separates different trajectories, while a distributional regulariser keeps the content coordinates from collapsing. On four control benchmarks, the model plans as well as the same architecture without the split. In return, its latent becomes easier to read. A few percent of its dimensions are enough to predict progress within an episode, and the angle they form follows the task, moves backward when the agent backtracks and changes fastest around events such as grasps and releases. Surprise also becomes anchored in the trajectory: when the scene changes mid-episode, the angle spikes and then relocalizes on a new part of the progression, anchoring surprise in the agent's own map of the task.