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

Train–plan gap

• JEPA world models predict short-horizon latent dynamics, whereas latent MPC must rank multi-step imagined futures by goal progress.LeWM and related planners typically use latent Euclidean distance as the planning cost, creating a structural mismatch between training and planning.Source: §1 · p. 1

Method

• Temporal-Distance-JEPA retains the LeWM encoder–predictor and SIGReg backbone while adding a directed, goal-conditioned temporal cost.The directed cost depends on the ordered pair of current and goal latents, reflecting asymmetric goal-reaching dynamics.Source: §1 · pp. 1–2
• Cross-trajectory negatives discourage spuriously low off-trajectory costs, while rollout consistency aligns the learned cost with the planner’s open-loop horizon.These components provide temporal supervision matched to the planning procedure.Source: §1 · p. 2

Plan-time cost selection

• Temporal supervision is deployed directly for topology-dominated control and used to shape representations for Euclidean planning when contact geometry dominates.The method uses dψ for topology-dominated control and latent ℓ2 for contact-rich control on the same temporally trained checkpoint.Source: §1 · p. 2

Results

• 14.2 points: shared Euclidean planning on the same temporally trained checkpoint raises OGB-Cube performance over LeWM and improves Push-T.Locked evaluations also report improvements on Two-Room and Reacher, while Temporal-Distance-JEPA matches or exceeds LeWM and RC-aux on every environment.Source: §1 · p. 2

JEPA and LeWM

• PlaNet and TD-MPC learn reward-related signals, while LeWM removes reward prediction and uses distance in a self-supervised predictive representation as the goal cost.Source: §2 · p. 2
• Temporal-Distance-JEPA remains reward-free but explicitly mines temporal order from demonstration logs to supervise progress, targeting the objective missing from prior approaches.Source: §2 · p. 2

Planning-aware latent geometry

• Planning-aware methods reshape JEPA geometry with quasimetric, curvature, or reachability objectives because short-horizon prediction alone may not align latent distance with multi-step reachability.Source: §2 · p. 3
• Unlike Value-Guided JEPA, Temporal-Distance-JEPA needs no IQL value estimator and mines reward-free step-count labels on the LeWM+SIGReg backbone.Source: §2 · p. 3

Quasimetric and goal-conditioned RL

• Quasimetric and goal-conditioned reinforcement-learning methods model directed or asymmetric reachability, linking temporal distances to goal-reaching costs and policies.Source: §2 · pp. 2–3

Backbone and objectives

• Temporal-Distance-JEPA keeps the LeWM prediction and SIGReg backbone, adding planning-specific supervision for the MPC cost.The additions include a directed temporal-distance objective and a horizon-matched rollout-consistency loss.Source: §3 · p. 3, Figure 2 · p. 4

Planning deployment

• On contact-rich tasks, the evaluation protocol plans with latent ℓ2 using the same temporally trained checkpoint rather than d_ψ.This preserves the temporally trained representation while changing the plan-time ranking cost.Source: §3 · p. 3, §3.6 · p. 5

Directed cost head

• The MRN head produces a non-negative asymmetric cost d_ψ(z_s→z_g), allowing directed reachability costs that latent ℓ2 cannot express.The symmetric component captures shared configuration factors, while the asymmetric residual encodes directed reachability.Source: §3.1 · p. 3, §3.3 · pp. 4–5

Temporal supervision

• Same-trajectory temporal gaps τ(i,j)=j−i calibrate cost values to on-manifold progress, while cross-trajectory hinge negatives penalize spuriously small costs.The step count is treated as a progress-ranking proxy rather than a shortest-path metric.Source: §3.1 · p. 3, §3.2 · p. 4, §3.3 · p. 5

Train–plan coupling

• Train–plan coupling makes the temporally supervised energy the default plan-time cost for open-loop latent rollouts ranked by CEM or iCEM.For manipulation evaluation, planning can instead blend the trained energy with latent geometry as (1−α)d_ψ+αℓ2.Source: §3.3 · p. 5, §3.4 · p. 5, §3.6 · p. 5

Locked comparisons

• Temporal-Distance-JEPA matches or exceeds both LeWM and RC-aux on every environment under matched episodes and planner budgets.The study covers Two-Room, Reacher, Push-T, and OGBench-Cube, with primary comparisons averaged over ten independent seed runs.Source: §4.1 · p. 6, §4 · p. 6

Component ablations

• Each component ablation hurts every planner setting, with the largest drop from removing the directed residual and a larger penalty from removing rollout consistency than the hinge.The results support co-designing cost form and supervision rather than relying on temporal labels alone.Source: §4.2 · p. 6

Plan-cost regimes

• Navigation favors deploying the mined temporal cost, while contact-rich manipulation favors latent ℓ2 planning after temporal representation learning.On Two-Room and Reacher, the mined cost drives gains; on Push-T and OGB-Cube, all methods plan with latent ℓ2, so improvements reflect representation learning.Source: §4.1 · p. 6, §4.3 · p. 7

Navigation analysis

• 97.4%: locked LeWM ℓ2 success on Two-Room is exceeded by Temporal-Distance-JEPA with deployed 𝑑𝜓, whereas its ℓ2 planner falls slightly below LeWM on Reacher.This shows that navigation gains are not explained solely by a better Euclidean representation when topology matters.Source: §4.3 · p. 7, Figure 4 · p. 7

Failure analysis

• On Push-T, 𝑑𝜓 ranks temporal gaps more accurately than ℓ2 on every held-out split, yet ℓ2 plans better because local geometry supports fine pose and angle control after contact.Phase logs show that failed 𝑑𝜓 episodes are almost entirely post-contact, while ℓ2 roughly halves post-contact failures.Source: §4.3 · p. 7, §4.4 · pp. 7–8

Push-T cost selection

• 86.0±4.2: pure latent ℓ2 outperforms pure 𝑑𝜓 at 69.0±1.9 on Push-T, while soft and fixed blends reach 84.0±4.8 and 83.4±3.7.A hard contact gate reaches 77.6±5.2, indicating that geometry is especially important after contact.Source: §4.4 · pp. 7–8

Conclusion

• Temporal-Distance-JEPA mines a directed temporal cost from reward-free demonstration logs on the LeWM encoder–predictor backbone.The mined signal is designed to narrow the train–plan gap in JEPA world models.Source: §5 · p. 9
• The mined signal can be deployed as the planning cost for topological progress or used as a representation signal when contact-rich control prefers latent geometry.Locked evaluation supports both choices: deploying dψ improves Two-Room and Reacher over LeWM, while shared-ℓ2 planning uses the same temporally trained checkpoint.Source: §5 · p. 9

Controlled evaluation

• The comparison fixes the epoch-10 checkpoint, 50-episode manifests, environment-specific solvers, and ten independent seed runs while varying only plan cost.This setup isolates plan-cost effects; the Push-T and OGB-Cube ℓ2 entries match Temporal-Distance-JEPA’s Table 2 row.Source: A Detailed Plan-Cost Results · p. 9

Results table

• Table 11 presents mean±std success over ten independent seed runs for the matched plan-cost matrix.Source: Table 11 · p. 9

• Mean d_ψ rises monotonically with demonstrated step separation on held-out same-trajectory pairs.Source: B Temporal-Cost Calibration · p. 9
• The learned temporal cost need not match the true step count in absolute scale.Source: B Temporal-Cost Calibration · p. 9
• The planner uses the induced ranking of temporal costs rather than their absolute scale.Source: B Temporal-Cost Calibration · p. 9

Training

• The short-window recipe uses 𝑁pred=5, window 𝑇=8, history size 3, rollout horizon 5, 𝜆roll=0.5, 𝜆td=1.0, and 𝜆sig=0.09.These settings define the temporal-distance training configuration.Source: C Training and Evaluation Protocol · p. 10
• All runs use a ViT-tiny encoder with ∼15M parameters, 14-pixel patches, 224-pixel images, 192-dimensional latents, AdamW at 5×10−5 learning rate, 10−3 weight decay, bf16 precision, and 10 epochs.The directed head has hidden dimension 512 with 𝜙sym,𝜙asym ∈R128; Push-T uses effective batch 128 (32×4 accumulation).Source: C Training and Evaluation Protocol · p. 10

Planning defaults

• MPC defaults use horizon 𝐻=5, goal offset 𝐺=25, 50 episodes, and 300 candidate action sequences per CEM/iCEM iteration unless noted.Two-Room and Reacher use iCEM-30 with pure 𝑑𝜓, with terminal weights 𝑤=1 and 𝑤=0.3, respectively.Source: C Training and Evaluation Protocol · p. 10
• Push-T and OGB-Cube locked comparisons use latent ℓ2 on the temporally trained checkpoint with CEM-30 and CEM-10, terminal weight 𝑤=1, matching LeWM’s geometric plan cost.RC-aux uses the same locked manifests.Source: C Training and Evaluation Protocol · p. 10

Statistical scope

• Unless noted, ± values are sample standard deviations over independent seed runs with checkpoints and episode manifests held fixed.Push-T ablations use three training seeds; paired comparisons, the cost matrix, and contact-gated Push-T use ten plan seeds, while rank correlations use independently trained checkpoints.Source: C Training and Evaluation Protocol · p. 10