Joint-Embedding Predictive Architectures (JEPAs) have become a popular backbone for latent model predictive control, as they learn world models by predicting in representation space rather than reconstructing pixels. However, standard JEPA training focuses on short-horizon latent prediction, while planning requires a multi-step ranking of imagined futures by goal progress. Prior JEPA planners often rely on latent Euclidean distance as a proxy for progress, which is a byproduct of representation learning rather than a cost explicitly derived from data.
To address this, researchers from HKBU and collaborators introduce Temporal-Distance JEPA (TD-JEPA), which retains the LeWM encoder–predictor backbone and mines a directed temporal cost from reward-free trajectories. The method uses same-trajectory step order as positive targets, cross-trajectory pairs as heuristic negatives, and a rollout-consistency term to match the planner horizon. This mined supervision serves dual purposes: as the deployed planning cost when progress is topological, and as a representation signal that improves Euclidean planning when contact geometry dominates.
In locked evaluation, TD-JEPA achieves 100.0% success on Two-Room compared to LeWM's 97.4%, and improves OGB-Cube by 14.2 points over LeWM when using shared Euclidean planning on the same temporally trained checkpoint. It also improves Push-T performance. Against LeWM and the concurrent RC-aux baseline, TD-JEPA matches or exceeds both methods on every environment tested.
Ablation studies confirm that the directed head, cross-trajectory negatives, and rollout consistency each contribute to the gains. The work narrows the train–plan gap for JEPA world-model planners by discovering temporal progress structure in offline logs and co-designing cost form with plan-time deployment. Code is available on GitHub.