Time-to-Goal: Learned Temporal Metrics for Long-Horizon Planning with JEPA World Models

Ruilong Yu; Fei Ye

cs.RONov 4, 2026

Time-to-Goal: Learned Temporal Metrics for Long-Horizon Planning with JEPA World Models

Authors: Ruilong Yu, Fei Ye

Organizations: University of Electronic Science and Technology of China

Abstract

Joint-Embedding Predictive Architectures (JEPAs) support planning by comparing predicted latent states with an encoded goal. For this comparison to guide temporal progress, states closer in time to the goal should receive lower costs. However, training representations for prediction and distribution regularization does not directly enforce this ordering. We examine this mismatch in flat LeWorldModel (LeWM) by analyzing how its Euclidean goal cost changes with temporal separation. Our analysis identifies conditions under which states at different long offsets receive similar average costs, weakening their temporal distinction. To learn the missing ordering explicitly, we introduce Time-to-Goal (T2G), a lightweight metric trained on latent triples from offline trajectories. A margin ranking loss encourages smaller distances for temporally nearer states, requiring no rewards or policy labels. The learned metric then replaces the Euclidean planning cost while the JEPA encoder and predictor remain frozen. Diagnostics on AntMaze and FourRoom show that T2G preserves stronger temporal contrast and improves near--far ranking at long offsets. Across four navigation and manipulation benchmarks, T2G improves the average planning success rate from 32.7% (flat LeWM) to 63.0% and approaches or even surpasses the performance of the hierarchical H-JEPA baselines. These results show that learning temporal ordering for the goal cost can substantially improve latent planning even with a flat world model.

Figures & tables

Figure 1: On successful FourRoom executions, Euclidean cost weakly signals initial progress (i.e., distant to goal), whereas T2G provides a more consistent temporal cost.

Figure 2: Changes in the components of Euclidean temporal contrast on AntMaze (left) and FourRoom (right), using 50 recorded planner executions per environment. The horizontal axis is the separation

k

from a sampled future goal in model steps, with a fixed comparison gap of

\Delta=1

model step (five environment steps). Curves show

|\delta_{\Delta}B(k)|

and

|\delta_{\Delta}q(k)|

, each divided by the environment’s mean Euclidean squared distance

\overline{D}_{E}

over the tested offsets. In both environments, the

B

term varies little, while changes in the

q

term become smaller at longer offsets.

Figure 3: Overview of T2G metric training. A latent triple $({\bf z}_{t},{\bf z}_{t+k},{\bf z}_{t+k+\Delta})$ is sampled from an offline trajectory, where the near and far states are separated from the anchor by $k$ and $k+\Delta$ steps, respectively. The three encoder blocks denote the same frozen JEPA encoder applied to the corresponding observations. T2G evaluates the learned temporal distances from the anchor to the near and far latents and optimizes the metric with a margin ranking objective, so that the near state is closer to the anchor than the far state by at least a margin $m$ . Only the metric network and its scale are updated, and the world model remains frozen.

Environment	Episode budget	$H$	Last- $n$	$S$	$\eta$
FourRoomDistractors	210	23	2	200	0.1
Visual AntMaze	350	23	2	100	0.3
Push-T	90	15	2	800	0.05
OGBench Cube	30	2	1	800	0.05

Table 1: Planning settings of the proposed T2G. Episode budgets are in environment steps,

H

is the beginning planning horizon in predicted steps,

S

is the number of candidate action sequences, and

\eta

is the AdamW learning rate. Last-

n

is the number of final predicted states averaged by the T2G last-

n

cost. Candidate counts, learning rates, and episode budgets follow the flat LeWM baseline of H-JEPA [ 5 ] ; Cube uses an initial horizon of two.

Figure 4: Overview of T2G planning. The current observation $o_{1}$ and goal observation $o_{g}$ are encoded into latents ${\bf z}_{1}$ and ${\bf z}_{g}$ by the frozen JEPA encoder. A solver optimizes an action sequence $a_{1:H}$ , and the frozen predictor rolls out the resulting latent trajectory ${\bf z}_{1}\!\to\!{\bf z}_{2}\!\to\!\cdots\!\to\!{\bf z}_{H}$ . Each predicted latent is scored by its learned temporal distance to the goal, $d_{\theta}({\bf z}_{t},{\bf z}_{g})$ , and the solver updates the actions to minimize the aggregated cost. The updated actions are executed and planning is repeated from new observations in a receding-horizon loop, while the encoder, predictor, and solver remain unchanged.

Method	New Params $\downarrow$	Visual AntMaze $\uparrow$	OGBench Cube $\uparrow$	FourRoomDistractors $\uparrow$	Push-T $\uparrow$
H-JEPA (2 levels)	10.79	$39.3\pm 3.7$	$47.3\pm 2.7$	$81.3\pm 6.7$	$45.3\pm 1.3$
H-JEPA (3 levels)	21.58	$73.3\pm 3.5$	$60.0\pm 5.0$	$96.0\pm 1.1$	$17.3\pm 3.5$
H-JEPA (4 levels)	32.37	$63.3\pm 7.7$	$55.3\pm 4.4$	$80.0\pm 4.0$	$0.7\pm 0.7$
LeWM (flat)	--	$18.0\pm 3.5$	$32.0\pm 6.1$	$40.7\pm 6.8$	$40.0\pm 2.3$
+ T2G dense	0.18	$48.3\pm 2.4$	$4.0\pm 1.2$	$81.3\pm 2.9$	$42.0\pm 1.2$
+ T2G lastn	0.18	$58.7\pm 2.4$	$63.3\pm 0.7$	$78.7\pm 2.9$	$51.3\pm 5.9$

Table 2: Planning success rate (%) on the evaluation tasks, averaged over multiple seeds. LeWM (flat) uses the native level-1 cost; the 2-, 3-, and 4-level rows use hierarchical H-JEPA planning, and the T2G rows use the same frozen flat world model with either a full-rollout (dense) or final-state-window (last-

n

) temporal cost. Window sizes are given in Tab. 1 . New Params denotes the number of additional parameters in millions (M) compared to LeWM (flat). Cell shading continuously interpolates from light red to deep red using min–max normalization independently within each column. Higher mean success rates receive darker shading; the scale is reversed for New Params , where lower values are better. Standard errors do not affect shading. The missing entry ( -- ) is unshaded and excluded from normalization.

Figure 5: Selected successful LeWM+T2G planner executions using our T2G cost, on one evaluation task per environment.

Figure 6: Temporal guidance of Euclidean squared and T2G costs on recorded AntMaze (top) and FourRoom (bottom) executions. Left: mean cost to the same sampled future goals, normalized by each cost’s mean over the tested offsets. Middle: normalized temporal contrast between offsets separated by one model step (five environment steps). Right: the fraction of comparisons in which the temporally nearer state receives a lower cost, with ties receiving half credit; the dashed line marks 50%. Horizontal axes report temporal separation in environment steps. At long offsets, T2G maintains a larger positive contrast and higher ranking accuracy in both environments.

Appendix figures & tables1 asset

Supplementary material from the paper’s appendix.

Appendix

Environment	Latent dimension $D$	Image size	$T_{\mathrm{clip}}$	$K_{\max}$
AntMaze	384	$64\times 64$	24	23
FourRoom	192	$224\times 224$	24	23
Cube	192	$224\times 224$	16	15
Push-T	192	$224\times 224$	16	15

Appendix

Table 3: Input and sampling hyper-parameters for T2G training. Image size is in pixels. Clip length

T_{\mathrm{clip}}

counts latent states, and

K_{\max}

is measured in model steps. All environments use

k_{\min}=1

\Delta_{\min}=2

, and five environment steps per model step.

World Model-Based Planning

World model-based planning uses models of how environments change to compare imagined actions and choose goal-directed behavior.

See all 274 papers

Jul 28, 2026

Temporal-Distance JEPA: Plan-Aware Representation Learning for Latent World Model Predictive Control

Joint-Embedding Predictive Architectures (JEPAs) learn world models by predicting in representation space rather than reconstructing pixels, making them a natural backbone for latent model predictive control from offline demonstration logs. JEPA-style training optimizes short-horizon latent prediction, whereas planning requires a multi-step ranking of imagined futures by goal progress. Prior JEPA planners often inherit that ranking from embedding geometry, typically latent Euclidean distance, which arises as a byproduct of representation learning rather than as a progress cost mined from the logs. We propose Temporal-Distance-JEPA, which retains the LeWM encoder--predictor backbone and mines a directed temporal cost from reward-free trajectories: same-trajectory step order supplies positive targets, cross-trajectory pairs act as heuristic negatives, and a rollout-consistency term matches the planner horizon. The mined supervision serves two roles: as the deployed planning cost when progress is topological, and as a representation signal that improves Euclidean planning when contact geometry dominates. Under locked evaluation, deploying the mined cost raises Two-Room success to 100.0% versus LeWM's 97.4%, while shared Euclidean planning on the same temporally trained checkpoint raises OGB-Cube by 14.2 points over LeWM and improves Push-T. Against LeWM and the concurrent RC-aux baseline under locked evaluation, Temporal-Distance-JEPA matches or exceeds both methods on every environment. Ablations show that the directed head, cross-trajectory negatives, and rollout consistency each contribute. Temporal-Distance-JEPA 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 at https://github.com/HKBU-KnowComp/Temporal-Distance-JEPA.

Jiaxin Bai, Jiaxuan Xiong

Hong Kong Baptist University Hong Kong, China

World Model-Based Planning Joint-Embedding Predictive Architecture Model Predictive Control Model-Based Planning Predictive Representation Learning Latent World Models+1

Ask PDF

Oct 4, 2026

How Long, Not How Close: A Learned Temporal Metric for Planning in Latent World Models

Latent world models plan by rolling a frozen predictor forward under candidate action sequences and ranking the candidates by the latent distance between their imagined end state and the goal. However, this ranking breaks down when the goal lies several plans away, because the latent distance measures how closely an end state resembles the goal rather than how far it remains from reaching it. To address this, we propose TEMPO, a temporal-distance planning objective that leaves the world model untouched, learns only from the recorded trajectories already used to train it, and adds negligible cost to the planner's search. TEMPO learns a small map of the frozen latent in which the distance between two states of an episode reflects the number of environment steps between them, and blends this distance into the planner's cost. It requires no rewards, policies or success labels and, being a cost rather than a model, applies to frozen world models with one latent vector per state that plan by a latent distance. We evaluate TEMPO on eleven simulated environments (e.g., maze navigation, tabletop pushing, robotic arm control and three-dimensional manipulation) with the LeWM and PLDM planners. With a small MLP that adds at most 0.3% to a plan's arithmetic, TEMPO improves both planners at every goal distance, including the one-plan setting of their evaluations, raises LeWM from 36% to 99% on TwoRoom three plans from the goal, and remains competitive on a broad range of 2D and 3D navigation, reaching and manipulation tasks.

Lama Moukheiber, Haotian Xue, Yongxin Chen

Georgia Institute of Technology

World Model-Based Planning Long-Horizon Planning Latent World Models Latent-Space Planning

Ask PDF

Mar 12, 2026

Temporal Straightening for Latent Planning

Learning good representations is essential for latent planning with world models. While pretrained visual encoders produce strong semantic visual features, they are not tailored to planning and contain information irrelevant -- or even detrimental -- to planning. Inspired by the perceptual straightening hypothesis in human visual processing, we introduce temporal straightening to improve representation learning for latent planning. Using a curvature regularizer that encourages locally straightened latent trajectories, we jointly learn an encoder and a predictor of a Joint-Embedding Predictive Architecture (JEPA) world model. We show that reducing curvature this way makes the Euclidean distance in latent space a better proxy for the geodesic distance and improves the conditioning of the planning objective. We demonstrate empirically that temporal straightening makes gradient-based planning more stable and yields significantly higher success rates across a suite of goal-reaching tasks. Our code is available at https://agenticlearning.ai/temporal-straightening.

Ying Wang, Oumayma Bounou, Gaoyue Zhou +4

New York University · Brown University · University of Toronto

Latent-Space Planning World Model-Based Planning Trajectory Representation Learning Joint-Embedding Predictive Architecture

Ask PDF

Abstract

World Model-Based Planning

Explore similar work

Temporal-Distance JEPA: Plan-Aware Representation Learning for Latent World Model Predictive Control

How Long, Not How Close: A Learned Temporal Metric for Planning in Latent World Models

Temporal Straightening for Latent Planning