Abstract.
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 (TD-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 versus LeWM’s , while shared Euclidean planning on the same temporally trained checkpoint raises OGB-Cube by points over LeWM and improves Push-T. Against LeWM and the concurrent RC-aux baseline under locked evaluation, TD-JEPA matches or exceeds both methods on every environment. Ablations show that the directed head, cross-trajectory negatives, and rollout consistency each contribute. TD-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/TD-JEPA.
1. Introduction
World models aim to capture enough structure about an environment that an agent can reason about futures without acting in the real world (Li, 2026). In the planner-facing role that motivates this paper, a world model encodes observations, imagines the consequences of candidate action sequences, and ranks those imagined futures by how close they bring the agent to a goal. When the only available supervision is offline demonstration logs—observation–action trajectories without dense rewards—predictive dynamics alone do not yield a progress signal that the planner can trust.
Joint-Embedding Predictive Architectures (JEPAs) provide a natural foundation for this setting (LeCun, 2022; Balestriero and LeCun, 2025). Rather than reconstructing pixels, a JEPA learns an encoder that maps observations into a latent space and a predictor that forecasts future embeddings from context and actions. Useful world knowledge resides in the predictive structure of representations: if the model can anticipate what will happen next in embedding space, it has internalized aspects of the world’s dynamics without committing to a generative decoder. LeWM instantiates this recipe for latent control (Maes et al., 2026). It trains an action-conditioned encoder–predictor with next-latent prediction, regularizes the representation with SIGReg to avoid collapse (Balestriero and LeCun, 2025), and at test time searches over action sequences with CEM (de Boer et al., 2005), scoring each imagined terminal latent by Euclidean distance to a goal embedding.
This design yields a strong predictive backbone, yet it leaves a structural gap between what is learned and what planning needs. JEPA-style training optimizes short-horizon prediction in latent space; latent model predictive control (MPC) instead needs a multi-step ranking of futures by goal progress. In LeWM and related JEPA planners, that ranking is typically inherited from embedding geometry: Euclidean distance between latents becomes the de facto planning cost. Geometry can correlate with progress when local appearance and control advance together, but it is not trained to reflect steps-to-goal, directed reachability, or planner ranking. Concurrent planning-aware JEPA variants address pieces of this mismatch: value shaping with offline RL (Destrade et al., 2026), temporal straightening of latent paths (Wang et al., 2026b), horizon-conditioned reachability auxiliaries (Li et al., 2026), and progression subspaces (Thil et al., 2026). These approaches either introduce value estimators beyond the demonstration logs, reshape geometry without an explicit plan cost, or treat temporal structure as an auxiliary rather than as a discoverable ranking signal for MPC.
We therefore ask: how can reward-free demonstration trajectories yield temporal progress structure that narrows the train–plan gap for JEPA world-model planners, and when should that structure be deployed as the plan cost versus used to shape representations for geometric planning? Figure 1 makes the ranking side of this gap concrete on Push-T. On held-out demonstration pairs, LeWM’s latent Euclidean distance correlates with temporal separation at Spearman , whereas a cost mined from the same logs tracks step order at . Whether that mined signal also improves control—and under which plan-time cost—is the empirical question addressed below.
We propose temporal-distance JEPA (TD-JEPA) to address this challenge. TD-JEPA keeps the LeWM encoder–predictor and SIGReg backbone, and mines a goal-conditioned energy that is low when a rolled-out latent is compatible with reaching the goal (LeCun et al., 2006). Because goal reaching under actions is generally asymmetric, the cost is directed rather than Euclidean (Wang et al., 2023): it depends on the ordered pair of current and goal latents, so traveling from toward need not cost the same as the reverse. Supervision comes from trajectory mining rather than rewards. Same-trajectory step order supplies positive temporal targets; cross-trajectory pairs act as heuristic negatives that discourage spuriously low off-trajectory costs; and a rollout-consistency term aligns the learned cost with the open-loop horizon used by the planner (Fig. 2).
The mined signal plays two roles, and plan-time cost selection follows task structure. When progress is mainly topological, as in navigation and reaching, the directed temporal cost is the natural planning objective. When contact geometry dominates, the same temporally trained checkpoint is planned with latent Euclidean distance, so temporal supervision shapes the representation used by a geometric planner. We deploy for topology-dominated control and latent for contact-rich control; later sections refer to this design choice rather than re-arguing it.
Locked evaluations support both roles. Deploying the mined cost improves Two-Room and Reacher over LeWM; shared Euclidean planning on the same temporally trained checkpoint raises OGB-Cube by points and improves Push-T. Against LeWM and RC-aux (Li et al., 2026) under locked evaluation, TD-JEPA matches or exceeds both methods on every environment. Ablations show that the directed head, cross-trajectory negatives, and rollout consistency each contribute, while contact-phase diagnostics clarify when geometry must be selected at plan time.
Contributions.
- •
A clear diagnosis of the JEPA train–plan gap. We show that predictive JEPA world models learn dynamics from offline logs, while plan-time scoring often relies on latent embedding geometry rather than an explicitly mined temporal-progress cost, and that a log-derived directed cost closes much of this ranking mismatch (Fig. 1).
- •
Temporal-distance mining for JEPA planners. We introduce TD-JEPA, which keeps the LeWM backbone and adds directed temporal supervision, cross-trajectory negatives, and horizon-matched rollout consistency so that progress structure discovered in demonstrations serves as either a deployed planning cost or a representation signal (Sec. 3).
- •
Evidence that cost form and deployment should be co-designed. Under matched evaluation, deploying the mined cost helps where topology governs progress, while shared Euclidean planning on the temporally trained checkpoint helps on contact-rich tasks; ablations and contact-phase analyses identify when each choice is appropriate (Tables 2, 5, 4, and 9).
2. Related Work
Energy-based learning and planning. Energy-based models associate low energy with compatible variable assignments and perform inference by energy minimization (LeCun et al., 2006). Latent MPC fits this template by searching for action sequences that minimize terminal energy after imagined rollout. Value-guided and quasimetric methods learn energies or values from offline data, providing a foundation for goal-conditioned planning costs (Wang et al., 2023; Destrade et al., 2026). TD-JEPA follows the same inference view but changes the knowledge source: instead of a value target, it mines a directed temporal cost from reward-free step-count labels and cross-trajectory negative pairs in demonstration logs (van den Oord et al., 2018; Kostrikov et al., 2021).
JEPA and LeWM. Latent-space planning from pixels predates JEPA control. PlaNet learns stochastic latent dynamics and plans online with predicted rewards, while TD-MPC jointly learns task-oriented dynamics and a terminal value through temporal-difference learning (Hafner et al., 2019; Hansen et al., 2022). JEPAs instead learn predictive representations from data without pixel reconstruction; LeJEPA introduces SIGReg for stable self-supervision, and LeWM combines this recipe with latent planning (Balestriero and LeCun, 2025; Maes et al., 2026). Recent extensions explore relational interventions, online adaptation, and temporal abstraction (Nam et al., 2026; Wang et al., 2026a; Zhang et al., 2026; Masip et al., 2026). TD-JEPA leaves the LeWM dynamics backbone unchanged and instead mines the goal-progress cost consumed by the planner.
The central distinction is where task structure is discovered. PlaNet learns rewards with a reconstructive latent model, and TD-MPC learns reward and value predictions through online interaction. LeWM removes reward prediction and treats distance in a self-supervised predictive representation as the goal cost. TD-JEPA stays reward-free, but makes progress supervision explicit by mining temporal order from demonstration logs. It therefore targets the missing objective between latent prediction and action search, rather than replacing the dynamics architecture or introducing a learned policy.
Value-shaped and temporal visual representations. VIP learns a temporally smooth visual embedding whose distance acts as a goal-conditioned reward, and LIV extends value-shaped representation learning to language and image goals (Ma et al., 2023b, a). TLDR learns temporal-distance-aware representations for goal selection, intrinsic reward, and goal-conditioned policy learning (Bae et al., 2024). These methods use values or temporal distances to shape rewards and policies. TD-JEPA instead mines a directed cost inside an action-conditioned JEPA world model and uses that cost to rank imagined rollouts. VIP and LIV provide reward signals to a downstream controller; TLDR trains goal-conditioned policies and exploration objectives. TD-JEPA stays in the offline, reward-free setting: no reward model and no policy learner, only a temporal cost discovered from logs and consumed by planning.
Planning-aware latent geometry. Accurate short-horizon prediction is insufficient when MPC searches over multi-step rollouts in a latent space whose metric is not aligned with reachability. Destrade et al. shape JEPA embeddings so a quasimetric distance approximates a goal-conditioned value function learned with IQL, then plan with MPC on that geometry (Destrade et al., 2026). Unlike Value-Guided JEPA, TD-JEPA needs no IQL value estimator: it mines reward-free step-count labels on the LeWM+SIGReg backbone and evaluates task-dependent deployment of versus latent (Table 1). Wang et al. regularize latent curvature so Euclidean distance better tracks geodesic progress after temporal straightening (Wang et al., 2026b). RC-aux adds multi-horizon open-loop prediction and budget-conditioned reachability supervision on top of LeWM, distinguishing states reachable within the planner’s horizon from merely eventually reachable ones (Li et al., 2026). SD-JEPA partitions the latent into orthogonal progression and content subspaces, training angular coordinates with a cosine-margin triplet loss while retaining SIGReg on the remainder (Thil et al., 2026). Hierarchical planners and latent-planner JEPAs instead modify inference structure, using coarse latent models or action-free subgoal predictors to reduce long-horizon search difficulty (Zhang et al., 2026; Masip et al., 2026).
| Method | Supervision | Directed | Plan use |
|---|---|---|---|
| LeWM (Maes et al., 2026) | none | no | latent |
| Value-guided (Destrade et al., 2026) | IQL value | yes | learned value |
| RC-aux (Li et al., 2026) | horizon budget | no | auxiliary |
| SD-JEPA (Thil et al., 2026) | triplet angle | partly | partly |
| TD-JEPA | mined step count | yes | or |
Quasimetric and goal-conditioned RL. Quasimetric representations model directed reachability distances; optimal goal-reaching cost-to-go is generally asymmetric (Wang et al., 2023). Value-guided JEPA approximates with latent distance via IQL (Destrade et al., 2026; Kostrikov et al., 2021). Recent offline GCRL work further strengthens the link between temporal distances and goal reaching: TMD combines contrastive successor-feature learning with quasimetric constraints to recover stitchable temporal distances (Myers et al., 2025), while ProQ uses learned asymmetric distances as directional costs over latent keypoints for long-horizon planning (Kobanda et al., 2025). Temporal representation learning also studies continuous time structure outside JEPA control, including TD InfoNCE for future-event prediction (Zheng et al., 2024) and trajectory encoders that preserve multi-scale spatiotemporal detail in mined mobility logs (Zhou et al., 2025b). TD-JEPA stays in this log-mining view: it regresses a directed temporal cost on demonstration step counts with pairs and does not estimate a reward or value model.
3. Method
We keep the LeWM encoder–predictor backbone and mine a directed temporal cost from offline demonstration logs. The method has three parts: the LeWM prediction backbone, a directed cost head calibrated by trajectory step order, and a horizon-matched rollout loss on the same open-loop operator used at plan time. By default, latent MPC minimizes after open-loop rollout; on contact-rich tasks the evaluation protocol instead plans with latent on the same temporally trained checkpoint.
Backbone prediction loss.
Given observation and action , the LeWM backbone encodes observations and predicts the next latent:
| (1) |
The local dynamics term is the next-latent prediction loss
| (2) |
TD-JEPA leaves this backbone and SIGReg regularization intact, then adds planning-specific supervision for the cost optimized by MPC.
3.1. Two roles: cost class vs. cost calibration
TD-JEPA uses two related objects: an asymmetric scalar cost and a forward temporal-separation label. By “directed temporal cost” we mean a non-negative function of an ordered state–goal pair, not a vector displacement. The directed head sets the form of the cost; temporal labels calibrate its values on demonstration data.
- (1)
Directed cost head (geometry class). The planner optimizes a directed functional form. Goal-reaching under actions is generally asymmetric: the cost of reaching from can differ from the reverse cost. Optimal goal-reaching value functions and reachability costs are non-negative and asymmetric, motivating quasimetric-style parameterizations (Wang et al., 2023). Raw latent is symmetric (), so it cannot express this directionality by itself. The MRN head in Eq. (5) parameterizes a directed, non-negative cost with a symmetric configuration term plus an asymmetric residual. This parameterization encodes directionality by design. A controlled Push-T ablation keeps the same temporal labels but removes this MRN form: the ablated cost is , a symmetric Euclidean distance between learned projection features. We evaluate this symmetric temporal-distance head under the same locked Push-T protocol as the other component ablations (Table 4).
- (2)
Temporal-distance supervision (calibration signal). Temporal labels assign numeric values to that cost. The directed cost family gives TD-JEPA the right shape, but it still needs labels that tie energy values to task progress. Offline demonstrations supply a reward-free supervision signal: along the same trajectory, the temporal gap is a proxy for steps-to-subgoal on the data manifold. Eq. (6) regresses to , while cross-trajectory hinges penalize spuriously small costs for heuristic negative pairs.
Composition. Temporal labels set the magnitude of ; the quasimetric-style head makes that scalar depend on the order of its arguments. Together they yield an asymmetric ranking function for goal progress. TD-JEPA keeps LeWM’s dynamics and anti-collapse terms, and adds a planning term that supervises the same used in Eq. (7).
Relation to LeWM planning.
LeWM’s cost can be seen as a special case with identity and zero asymmetric residual. TD-JEPA adds a loss that encourages the plan-time cost to rank same-trajectory pairs by . This focus on plan geometry is consistent with recent plan-aware JEPA extensions (Destrade et al., 2026; Wang et al., 2026b; Li et al., 2026).
3.2. Energy-based and value-function foundation
TD-JEPA grounds in two established frameworks for plan geometry: energy-based learning (LeCun et al., 2006) and quasimetric goal-reaching RL (Wang et al., 2023; Destrade et al., 2026). This section states the connection and the assumption used by our temporal supervision.
Goal-conditioned energy.
In the energy-based view (LeCun et al., 2006), compatibility of variables is scored by an energy with low values on desirable configurations. For goal-reaching from latent state to goal , define a directed energy
| (3) |
Latent MPC selects actions that minimize terminal energy after rollout: . TD-JEPA learns and uses the same functional form as the default train-time and plan-time energy.
Link to optimal goal-reaching costs.
Consider deterministic dynamics with unit cost and define as the minimum number of steps needed to reach from . On mutually reachable states, this shortest-path cost is non-negative, vanishes on the diagonal, obeys the triangle inequality, and need not be symmetric; it is therefore a quasimetric (or an extended quasimetric if unreachable pairs are assigned ) (Wang et al., 2023). Value-guided JEPA methods learn with implicit Q-learning and set plan cost to or a learned quasimetric surrogate (Destrade et al., 2026). TD-JEPA instead uses explicit temporal labels from offline demonstrations. For a feasible trajectory segment from to , its length upper-bounds the shortest-path distance:
| (4) |
Equality holds only when the demonstrated segment is a shortest path, so we treat demonstration path length as an on-manifold surrogate whose ordering reflects goal progress rather than as itself. Eq. (4) is therefore an empirical ranking-calibration step. The MRN head supplies a directed parameterization inspired by the quasimetric structure of (Wang et al., 2023); Eq. (6) targets in step units as a training signal, and Fig. 1 shows the resulting monotonic tracking.
3.3. Directed distance and temporal supervision
We instantiate the directed cost with the metric-residual (MRN) parameterization (Wang et al., 2023). Each latent is mapped to symmetric and asymmetric features and , and
| (5) |
The head is non-negative by construction and can satisfy . The symmetric term captures shared configuration factors; the asymmetric residual encodes directed reachability.
Expert demonstrations provide the calibration labels. For indices on the same trajectory, define as forward steps along the observed segment (not a shortest-path metric). We minimize
| (6) |
where is Smooth L1, permutes cross-trajectory goals at the same time index, and is a margin with window length . The regression term fits on-manifold progress; the hinge discourages spuriously low off-trajectory costs, while admitting false negatives when trajectories share reachable states. This is regressive EBM calibration without estimating a partition function: positives are shaped toward , and negatives are pushed above margin (LeCun et al., 2006; van den Oord et al., 2018; Destrade et al., 2026). We treat demonstration step counts as an on-manifold proxy for progress, not as shortest-path distances.
Train–plan coupling.
TD-JEPA trains the energy that serves as the default plan-time cost:
| (7) |
where is the -step open-loop rollout of the learned predictor. CEM therefore ranks candidates with the energy shaped by by default. For manipulation evaluation, the cost composes this trained energy with latent geometry as ; Sec. 4.3 reports this blend as a planning-time diagnostic.
3.4. Latent MPC with CEM and iCEM
At plan time, MPC searches over action sequences that minimize the latent rollout cost in Eq. (7), executes the first action block, and replans from the resulting observation. Algorithms 1 and 2 summarize the two solvers used in our evaluations. Standard CEM updates a diagonal Gaussian over action sequences from elite moments (de Boer et al., 2005). iCEM adds temporally colored noise, elite reuse across refinements, and momentum-smoothed mean/variance updates (Pinneri et al., 2020). Both call the same latent-rollout cost: open-loop prediction for horizon , then scoring with (or latent on contact-rich manipulation).
3.5. Horizon-aligned rollout consistency
CEM scores -step open-loop rollouts. We add a multi-step teacher-forcing loss on expert actions:
| (8) |
with matched to the planner horizon. Minimizing reduces compounding error in the same rollout operator that appears inside Eq. (7); together with , it couples temporal and spatial planning structure.
3.6. Full objective
The complete TD-JEPA objective is
| (9) |
where is the LeWM one-step MSE. The short-window recipe uses , , , prediction window 5, and rollout horizon 5. SIGReg preserves the anti-collapse prior from LeWM (Balestriero and LeCun, 2025; Maes et al., 2026). At plan time, CEM or iCEM (Algs. 1–2) optimizes Eq. (7) with by default; on contact-rich tasks the main protocol uses latent on the same checkpoint (Sec. 4.3).
4. Experiments
We address three questions: (i) Does temporal-distance mining from reward-free logs improve planning over LeWM and RC-aux (Li et al., 2026) under a fixed evaluation budget? (ii) Which mining ingredients drive the gains? (iii) When should the mined temporal cost be deployed, and when should latent geometry still be selected?
We study these questions on the four LeWM environments—Two-Room, Reacher, Push-T, and OGBench-Cube—covering navigation, reaching, and contact-rich manipulation. TD-JEPA retains the LeWM encoder–predictor backbone and is trained for ten epochs; full hyperparameters appear in Appendix C. Unless noted otherwise, primary comparisons report meanstd over ten independent seed runs with the checkpoint and episode manifest held fixed. Diagnostic ablations and phase logs use three seeds; we note this whenever those tables appear. Plan-time defaults follow the two-role design: TD-JEPA deploys pure on Two-Room and Reacher, and latent on Push-T and OGB-Cube. LeWM and RC-aux use latent throughout under the same locked manifests. On Reacher, the planner aggregates terminal and trajectory-mean costs with weight so that smooth progress is credited; terminal-only scoring is a sensitivity check in Sec. 4.5.
4.1. Planning Performance
Table 2 compares TD-JEPA with LeWM and RC-aux under matched episode manifests and planner budgets. RC-aux is evaluated with latent planning on all four environments, matching its role in Table 1.
| Method | Two-Room | Reacher | Push-T | OGB-Cube |
|---|---|---|---|---|
| LeWM | ||||
| RC-aux (Li et al., 2026) | ||||
| TD-JEPA (ours) |
Under this locked protocol, TD-JEPA matches or exceeds both LeWM and RC-aux on every environment. On Two-Room and Reacher, deploying the mined cost yields the gains, showing the value of an explicit temporal plan cost over horizon-budgeted reachability auxiliaries that still plan with geometry. On Push-T and OGB-Cube, all three methods plan with latent , so improvements reflect temporal representation learning rather than direct cost deployment. The largest manipulation gain appears on OGB-Cube, where TD-JEPA improves over LeWM by percentage points and slightly exceeds RC-aux. Thus temporal supervision helps in both roles: as a deployed plan cost on topology-dominated tasks, and as a representation signal for geometric planning on contact-rich tasks.
For broader context, Table 3 places the same locked TD-JEPA numbers beside published offline RL, imitation learning, and latent world-model baselines under their native protocols. Those rows are not a matched head-to-head; all LeWM and RC-aux claims in this paper use Table 2.
| Method | Two-Room | Reacher | Push-T | OGB |
|---|---|---|---|---|
| GCBC (Ghosh et al., 2020) | — | |||
| IVL (Kostrikov et al., 2021) | — | |||
| IQL (Kostrikov et al., 2021) | — | |||
| PLDM (Sobal et al., 2025) | ||||
| DINO-WM (Zhou et al., 2025a) | ||||
| LeWM (Maes et al., 2026) | ||||
| TD-JEPA (ours) |
Relative to published baselines under non-matched protocols, TD-JEPA achieves the highest reported success on Two-Room and Reacher and remains competitive on OGB-Cube. We next ablate training ingredients, then ask when pure should be deployed.
4.2. Which Cost Components Matter?
We ablate the directed head, cross-trajectory negatives, and rollout consistency on Push-T (Table 4). Unlike the ten-seed primary comparisons, these variants use three independent seed runs under a ten-epoch training budget continued from four-epoch checkpoints, with the same 50 validation episodes. The columns report a diagnostic blend, pure , pure latent , and iCEM; the full-model entry of is therefore a three-seed diagnostic and need not match the ten-seed in Table 2. The symmetric-head row keeps the temporal regression and hinge but replaces the MRN residual with a Euclidean projection distance, separating cost form from supervision.
| Variant | Blend | iCEM | ||
|---|---|---|---|---|
| TD-JEPA (full) | ||||
| Sym. Euclidean TD head | ||||
| w/o cross-traj. hinge | ||||
| w/o rollout |
Each ablation hurts every planner setting relative to the full model. Removing the directed residual causes the largest drop: temporal labels alone are not enough without the MRN form. Removing rollout consistency hurts more than removing the hinge. Cost form and supervision must therefore be co-designed: a symmetric Euclidean head, missing negatives, or a missing horizon-matched each weaken the planner.
4.3. When Should the Mined Temporal Cost Be Deployed?
We next fix the TD-JEPA checkpoint, episodes, solver, and budget, and vary only the plan-time objective among , latent , and their blend with . Figure 3 summarizes the comparison; exact values appear in Appendix Table 11. With everything else held fixed, differences in success come from the cost surface seen by CEM.
Plan-cost comparison reveals two regimes. In navigation-dominated tasks, pure matches or exceeds : Two-Room reaches full success, and Reacher favors the temporal cost. In contact-rich manipulation, outperforms pure by a wide margin on Push-T and by a smaller margin on OGB-Cube; a small geometric blend partially closes but does not eliminate the gap.
Table 5 isolates the navigation regime by comparing LeWM’s native planner with the same TD-JEPA checkpoint under and under . Deploying exceeds LeWM in both environments, while planning the TD-JEPA checkpoint with falls slightly below LeWM on Reacher. The navigation gain is therefore not explained by a better Euclidean landscape alone: where topology matters, the temporal cost itself helps. Figure 4 shows the Two-Room training trajectory under the deployed planner. Early epochs are unstable under pure temporal scoring, then success climbs from at epoch 5 to full success by epoch 9, surpassing the locked LeWM baseline. This trajectory supports evaluating the final temporally trained checkpoint for navigation deployment under the two-role protocol, rather than selecting an intermediate epoch by peak online score alone.
| Environment | Planner setting | Success (%) |
|---|---|---|
| Two-Room | LeWM, | |
| Two-Room | TD-JEPA, | |
| Two-Room | TD-JEPA, | |
| Reacher | LeWM, | |
| Reacher | TD-JEPA, | |
| Reacher | TD-JEPA, |
On manipulation tasks, temporal training improves planning even though ranks temporal gaps more accurately. Table 6 reports Spearman correlations on Push-T over three independently trained checkpoints with 1,000 held-out pairs each: exceeds on every split, yet still plans better in Fig. 3. Better temporal ranking is therefore not the same as a better contact-rich plan cost. The two signals carry different information: tracks progress, while preserves local geometry useful near contact (Maes et al., 2026). Offline CEM sweeps in Fig. 5 reinforce the same pattern: manipulation benefits from a geometric plan cost, while temporal training still helps when that geometric cost is used.
| Split | ||
|---|---|---|
| All pairs | ||
| Far from endpoint | ||
| Near endpoint/contact |
The mined temporal signal therefore helps navigation when deployed as , and helps contact-rich control when used as representation shaping for an planner. The next subsection analyzes why pure underperforms on Push-T despite stronger temporal ranking.
4.4. Why Pure Temporal Cost Is Not Enough on Push-T
Unless stated otherwise, Push-T evaluation in this paper uses the primary locked protocol with planning in Table 2. Here we explain why pure underperforms that geometric planner: under the locked cost matrix, pure reaches while latent reaches . We test whether failures occur before contact (approach) or during post-contact alignment (manipulation).
MPC phase logging on the locked 50-episode manifest answers this (Table 7). These logs use three seeds, so per-cost success rates can differ slightly from the ten-seed matrix in Fig. 3. Failed episodes are almost entirely post-contact, whereas latent almost never fails before contact and roughly halves post-contact failures. Pure temporal cost is not enough because Push-T success depends on fine pose and angle control after contact, where local geometry matters more than a progress ranking.
| Plan cost | Pre-contact fails | Post-contact fails |
|---|---|---|
| Latent | ||
| Blend |
Start geometry modulates the size of the gap but not its dominant mode (Table 8). Across locked ten-seed outcomes, uniquely solves far more episodes than . The gap is smaller on above-median close starts and widens on far starts; splits on block–goal pose change the gap much less. Far starts are harder overall, but Table 7 shows that the dominant failure mode remains post-contact control rather than failing to arrive.
| Start stratum | Blend | Gap | ||
|---|---|---|---|---|
| All locked episodes | ||||
| Agent–block close | ||||
| Agent–block far |
A proximity gate that uses temporal cost far from contact and geometry near contact does not recover most of the score without retraining. Table 9 tests that idea with the same CEM-30 checkpoint and ten independent seed runs. A hard gate improves over pure but stays well below pure ; a soft gate matches the blend and still does not beat geometry alone. The planner needs geometric scoring throughout CEM search, not only after contact is detected.
| Plan cost | Success (%) |
|---|---|
| Pure | |
| Hard contact gate (far , near ) | |
| Soft contact gate | |
| Fixed blend | |
| Pure latent |
4.5. Robustness and Calibration
How the cost is aggregated over a rollout matters as much as which cost is used. Table 10 varies planner aggregation on Reacher under the locked validation manifest. The primary protocol mixes terminal and trajectory-mean costs with weight , matching temporal-distance supervision; the single-run rates of for TD-JEPA and for LeWM are consistent with, but not identical to, the ten-seed means in Table 2. Terminal-only scoring with is a sensitivity check that favors geometry and understates the temporal cost. Under the progressive aggregation used in the primary protocol, CEM and iCEM agree and TD-JEPA matches or exceeds LeWM. This choice affects only planner aggregation, not the checkpoint or training objective.
| Protocol | TD-JEPA | LeWM |
|---|---|---|
| CEM-30, (primary) | ||
| iCEM-30, (primary) | ||
| CEM-30, (sensitivity) |
Held-out rises monotonically with demonstrated step separation (Appendix Fig. 6), so the head learns the intended ordering even if its absolute scale need not match the true step count. Together with Table 6, this separates two claims: the cost ranks time well, but good ranking need not suffice for contact-rich control.
5. Conclusion
TD-JEPA narrows the train–plan gap in JEPA world models by mining a directed temporal cost from reward-free demonstration logs on the LeWM encoder–predictor backbone. The mined signal has two roles: as the deployed plan cost when progress is topological, and as a representation signal when contact-rich control prefers latent geometry. Locked evaluation supports both choices: deploying improves Two-Room and Reacher over LeWM, while shared- planning on the same temporally trained checkpoint raises OGB-Cube by points and improves Push-T. Contact-rich control still prefers geometry at plan time even when ranks temporal gaps better, and proximity-gated hybrids do not recover pure geometric scoring. Component ablations show that the directed residual, cross-trajectory hinge, and rollout consistency each contribute under locked planning.
References
- J. Bae, K. Park, and Y. Lee (2024) TLDR: unsupervised goal-conditioned reinforcement learning via temporal distance-aware representations. In Proceedings of the Conference on Robot Learning, Cited by: §2.
- R. Balestriero and Y. LeCun (2025) LeJEPA: provable and scalable self-supervised learning without the heuristics. External Links: 2511.08544 Cited by: §1, §2, §3.6.
- P. de Boer, D. P. Kroese, S. Mannor, and R. Y. Rubinstein (2005) A tutorial on the cross-entropy method. Annals of Operations Research 134, pp. 19–67. Cited by: §1, §3.4.
- M. Destrade, O. Bounou, Q. Le Lidec, J. Ponce, and Y. LeCun (2026) Value-guided action planning with JEPA world models. External Links: 2601.00844 Cited by: §1, §2, §2, §2, §2, §3.1, §3.2, §3.2, §3.3.
- D. Ghosh, A. Gupta, A. Reddy, J. Fu, C. Devin, B. Eysenbach, and S. Levine (2020) Learning to reach goals via iterated supervised learning. External Links: 1912.06088, Link Cited by: §4.1.
- D. Hafner, T. Lillicrap, I. Fischer, R. Villegas, D. Ha, H. Lee, and J. Davidson (2019) Learning latent dynamics for planning from pixels. In Proceedings of the 36th International Conference on Machine Learning, Proceedings of Machine Learning Research, Vol. 97, pp. 2555–2565. Cited by: §2.
- N. A. Hansen, H. Su, and X. Wang (2022) Temporal difference learning for model predictive control. In Proceedings of the 39th International Conference on Machine Learning, Proceedings of Machine Learning Research, Vol. 162, pp. 8387–8406. Cited by: §2.
- A. Kobanda, W. Radji, M. Petitbois, O. Maillard, and R. Portelas (2025) Offline goal-conditioned reinforcement learning with projective quasimetric planning. External Links: 2506.18847 Cited by: §2.
- I. Kostrikov, A. Nair, and S. Levine (2021) Offline reinforcement learning with implicit Q-learning. External Links: 2110.06169 Cited by: §2, §2, §4.1, §4.1.
- Y. LeCun, S. Chopra, R. Hadsell, M. Ranzato, and F. J. Huang (2006) A tutorial on energy-based learning. In Predicting Structured Data, Cited by: §1, §2, §3.2, §3.2, §3.3.
- Y. LeCun (2022) A path towards autonomous machine intelligence. Note: OpenReview preprintVersion 0.9.2 External Links: Link Cited by: §1.
- F. Li (2026) A functional taxonomy of world models: renderers, simulators, planners, and the loop that connects them. Note: SubstackPublished June 3, 2026 External Links: Link Cited by: §1.
- W. Li, G. Li, K. Maeda, T. Ogawa, and M. Haseyama (2026) Predictive but not plannable: RC-aux for latent world models. External Links: 2605.07278 Cited by: Appendix C, §1, §1, §2, §2, §3.1, §4.1, §4.
- Y. J. Ma, V. Kumar, A. Zhang, O. Bastani, and D. Jayaraman (2023a) LIV: language-image representations and rewards for robotic control. In Proceedings of the 40th International Conference on Machine Learning, Proceedings of Machine Learning Research, Vol. 202, pp. 23301–23320. Cited by: §2.
- Y. J. Ma, S. Sodhani, D. Jayaraman, O. Bastani, V. Kumar, and A. Zhang (2023b) VIP: towards universal visual reward and representation via value-implicit pre-training. In International Conference on Learning Representations, Cited by: §2.
- L. Maes, Q. Le Lidec, D. Scieur, Y. LeCun, and R. Balestriero (2026) LeWorldModel: stable end-to-end joint-embedding predictive architecture from pixels. External Links: 2603.19312 Cited by: §1, §2, §2, §3.6, §4.1, §4.3.
- S. Masip, J. Swinnen, Y. Hu, R. Detry, and T. Tuytelaars (2026) FF-JEPA: long-horizon planning in world models with latent planners. External Links: 2606.09311 Cited by: §2, §2.
- V. Myers, B. C. Zheng, B. Eysenbach, and S. Levine (2025) Offline goal-conditioned reinforcement learning with quasimetric representations. Note: Advances in Neural Information Processing Systems Cited by: §2.
- H. Nam, Q. L. Lidec, L. Maes, Y. LeCun, and R. Balestriero (2026) Causal-jepa: learning world models through object-level latent masking. External Links: 2602.11389, Link Cited by: §2.
- C. Pinneri, S. Sawant, S. Blaes, J. Achterhold, J. Stueckler, M. Rolinek, and G. Martius (2020) Sample-efficient cross-entropy method for real-time planning. In Proceedings of the Conference on Robot Learning, Cited by: §3.4.
- V. Sobal, W. Zhang, K. Cho, R. Balestriero, T. G. J. Rudner, and Y. LeCun (2025) Stress-testing offline reward-free reinforcement learning: a case for planning with latent dynamics models. In 7th Robot Learning Workshop: Towards Robots with Human-Level Abilities, External Links: Link Cited by: §4.1.
- L. Thil, J. Read, R. Kaddah, and G. Doquet (2026) Subspace-decomposed JEPAs: disentangling progression and content in latent world models. External Links: 2605.31111 Cited by: §1, §2, §2.
- A. van den Oord, Y. Li, and O. Vinyals (2018) Representation learning with contrastive predictive coding. External Links: 1807.03748 Cited by: §2, §3.3.
- T. Wang, A. Torralba, P. Isola, and A. Zhang (2023) Optimal goal-reaching reinforcement learning via quasimetric learning. In Proceedings of the 40th International Conference on Machine Learning, Proceedings of Machine Learning Research, Vol. 202, Honolulu, HI, USA, pp. 36411–36430. Cited by: §1, §2, §2, item 1, §3.2, §3.2, §3.2, §3.3.
- Y. Wang, O. Bounou, Y. LeCun, and M. Ren (2026a) AdaJEPA: an adaptive latent world model. External Links: 2606.32026 Cited by: §2.
- Y. Wang, O. Bounou, G. Zhou, R. Balestriero, T. G. J. Rudner, Y. LeCun, and M. Ren (2026b) Temporal straightening for latent planning. External Links: 2603.12231 Cited by: §1, §2, §3.1.
- W. Zhang, B. Terver, A. Zholus, S. Chitnis, H. Sutaria, M. Assran, R. Balestriero, A. Bar, A. Bardes, Y. LeCun, and N. Ballas (2026) Hierarchical planning with latent world models. External Links: 2604.03208 Cited by: §2, §2.
- C. Zheng, R. Salakhutdinov, and B. Eysenbach (2024) Contrastive difference predictive coding. Note: International Conference on Learning Representations Cited by: §2.
- G. Zhou, H. Pan, Y. LeCun, and L. Pinto (2025a) DINO-wm: world models on pre-trained visual features enable zero-shot planning. External Links: 2411.04983, Link Cited by: §4.1.
- S. Zhou, Y. Chen, S. Shang, L. Chen, B. He, and R. Shibasaki (2025b) Blurred encoding for trajectory representation learning. In Proceedings of the 31st ACM SIGKDD Conference on Knowledge Discovery and Data Mining, New York, NY, USA, pp. 4132–4143. External Links: Document Cited by: §2.
Appendix A Detailed Plan-Cost Results
Table 11 lists the exact values behind Fig. 3. The epoch-10 TD-JEPA checkpoint, locked 50-episode manifests, environment-specific solvers, and ten independent seed runs are held fixed; only the plan cost changes. The Push-T and OGB-Cube cells match the TD-JEPA row of Table 2.
| Plan cost | Two-Room | Reacher | Push-T | OGB-Cube |
|---|---|---|---|---|
| Directed temporal | ||||
| Latent | ||||
| Blend |
Appendix B Temporal-Cost Calibration
On held-out same-trajectory pairs, mean rises monotonically with demonstrated step separation. Absolute scale need not match the true step count; the planner uses the induced ranking.
Appendix C Training and Evaluation Protocol
Training.
The short-window recipe uses (window with history size 3), rollout horizon 5, , , and . Directed head: hidden dimension 512; . All runs share a ViT-tiny encoder (15M parameters, patch 14, image 224), 192-dimensional latent, history size 3, AdamW with learning rate and weight decay , bf16 mixed precision, and 10 epochs. Push-T uses effective batch 128 ( accumulation).
Planning defaults.
MPC uses horizon , goal offset , 50 episodes, and 300 candidate action sequences per CEM/iCEM iteration unless noted. Two-Room: iCEM-30 with pure and terminal weight . Reacher: iCEM-30 with pure and (terminal mixed with trajectory-mean cost). Push-T and OGB-Cube locked LeWM/RC-aux comparisons (Table 2): latent on the temporally trained checkpoint (CEM-30 and CEM-10; ), matching LeWM’s geometric plan cost. RC-aux uses the same locked manifests and latent planning on all four environments (CEM/iCEM settings as LeWM), with multi-horizon rollout and budgeted-reachability training losses as in (Li et al., 2026). Table 3 places the same locked TD-JEPA numbers in literature context beside published point estimates; it is not a matched head-to-head. Blend and pure appear in the matched cost matrix and contact diagnostics (Fig. 3, Tables 11 and 9). Push-T LeWM uses the released checkpoint; other environments use our LeWM checkpoints. iCEM uses noise_beta, smoothing coefficient , and n_elite_keep.
Statistical scope.
Unless noted otherwise, values are sample standard deviations () over independent seed runs with the checkpoint and episode manifest held fixed. Push-T ablations (Table 4) use three training-seed runs; paired LeWM/RC-aux comparisons, the cost matrix, and contact-gated Push-T results use ten plan seeds (Tables 2, 11, and 9). Push-T rank correlations (Table 6) use independently trained checkpoints.