D-JEPA: A Decision-Aligned Latent World Model
Shuaijun Liu, Chengyu Wu, Qifu Wen, Feiyang You, Chenglong Zhang, Shuyang Hao, Xi Lin, Ningxin Su
cs.RO, cs.LG
2026-09-21
D-JEPA aligns competing latent futures with executed outcomes and plans via native latent distance: 87.89% on PushT, +15.04 points on RoboTwin, +17 on physical robots.
Latent world models in the JEPA family (LeWM, TD-JEPA, DINO-WM, V-JEPA 2) predict action-conditioned futures in representation space, then execute the candidate whose predicted terminal sits closest to the goal. Pixel reconstruction is not the planning interface. Distance is.
That distance can rank the wrong action among the few futures that actually compete for execution. On a 96-start PushT audit, within-start Spearman correlation between predicted and realized latent costs falls from 0.90 to 0.11 for LeWM and from 0.80 to 0.13 for TD-JEPA once the pool is cut from 63 candidates to each model's top four. Among mixed-outcome top-four shortlists, a failing candidate is ranked strictly closer than a successful one 49.0% of the time for LeWM and 38.2% for TD-JEPA. Each model picks a failure on 8 of 96 starts even though a success sits in the same pool. One recorded pair: failing candidate RMS 0.2027, succeeding alternative 0.2167.
The paper calls this a decision-local prediction gap. Average prediction quality does not decide the next action.
D-JEPA leaves the pretrained world model in place. It treats the full set of candidate futures that share a context and a goal as one set, and uses executed success/failure labels to supervise relations among them.
Each candidate contributes two kinds of evidence. A goal-relative descriptor is the LayerNorm of (terminal latent minus goal), so each backbone keeps its own directions. An ordinal coordinate is the within-set rank of that backbone's native planning cost, scaled to [0, 1]. The dual-model token concatenates LeWM and TD-JEPA descriptors and ranks. Two permutation-equivariant Transformer layers (four heads, no positional encoding) see all candidates, 63 on PushT. A zero-initialized low-rank head, tanh-bounded at 0.2, adds a residual to the fused base score. A base winner ahead by more than 0.4 cannot be overturned, so the operator only reshuffles near the decision boundary.
Training mixes three terms: success mass over successful candidates, success/failure ordering on the current low-score subset, and an L2 penalty on the residual. At test time a calibrated gate falls back to the base fusion if the relational winner's margin is too small.
Predictive plasticity updates only the last TD-JEPA predictor block and the projection; its native-distance winner is a proposal, admitted only when a calibrated margin says so. An ordinal interface appends JEPA-WM and DINO-WM ranks so extra geometries vote without forcing a shared metric.
The learned order can be written back into JEPA-compatible futures. On PushT, exact ordinal realization places each terminal at radius π/(K+1) from the goal along its original direction, so native latent distance recovers the aligned choice. On Reacher, a coefficient bounded by 0.1 transports same-action futures along the difference between the two models.
Same starts, same candidate pools:
| Setting | D-JEPA | Main baselines |
| PushT success (n=256) | 87.89% | LeWM 83.59%, TD-JEPA 76.95%, four-geometry fusion 85.94% |
| Reacher (n=128) | 93.75% | LeWM 86.72%, TD-JEPA 68.75% |
| RoboTwin mean of four tasks (n=512) | 76.76% | native VLA 61.72%, Oracle 83.40% |
| PiPER physical mean of two tasks | 81.0% | same V-JEPA 2-AC planner 64.0% |
| Driving PDMS (0-100 trajectory score), seven scenes | 95.34 | Drive-JEPA 57.36 |
Paired counts: PushT versus TD-JEPA is +10.94 points (37 gains, 9 losses). Physical PushT is 12/3 on 50 paired trials, stacking 10/2. Drop the preservation gate and RoboTwin overall success falls to 71.48%.
Relational alignment carries the ablation. TD-JEPA 76.95%, plasticity alone 79.69%, relational alignment 87.11%, calibrated composition 87.89%. On a 128-start mechanism split, dual descriptors hit 93.75% against 84.38% (LeWM) and 83.59% (TD-JEPA). Growing the pool from 8 to 63 candidates lifts relational selection from 81.96% to 87.11%.
Exact ordinal realization holds 87.11% success at 34.02 ms median end-to-end, versus 52.16 ms for calibrated composition. Native goal distance recovers 100% of aligned actions and ranks. That recovery is by construction: ranks are written into radii.
Held-out PushObj shapes (I, small T, square) average 53.67% against fusion 44.67% and native 28.33%. Object-colour change is the largest appearance gain, 30% to 54% on 50 fresh starts. On Granular, D-JEPA's mean Chamfer is 0.358, worse than DINO-WM's 0.341, while strict-threshold success rises from 7.81% to 18.75%. Mean distance and crossing a threshold are different outcomes.
The JEPA line has been tightening prediction. This paper cuts the problem at the few futures a planner will actually run: full-set correlation can sit near 0.9 while top-4 correlation falls to 0.11. If a system already proposes candidates, this is a decision-stage module. VLA action chunks, Drive-JEPA trajectories, and V-JEPA 2-AC search results can be re-ranked without replacing the predictor or the low-level controller.
Bounded residuals and native-distance write-back exist so the existing latent-distance planning interface still works. The cost is training labels from counterfactual execution of the candidate set, which is not pure offline representation learning.
Against four-geometry fusion the PushT gap is 1.95 points. The larger claims are the locked-planner transfers: +17 points on the physical robot and +15.04 on RoboTwin.
The main text has no limitations section. The protocol supplies the constraints.
Training needs counterfactual rollouts of the full candidate set. The 256-start PushT confirmation cohort is filtered for mixed outcomes and a tight pretrained-score gap, so it is not a random start distribution. Driving uses seven hard scenes from three logs; Oracle is 100 on every scene and Drive-JEPA averages 57.36. The jump is large and still untested on a full driving benchmark. Appearance tests reuse the same known transform families, with new starts, not new distortion types.
The 0.2 bound can promote only candidates within 0.4 of the base minimum. Far mistakes in the pretrained ranking stay mistakes. RoboTwin Oracle still sits at 83.40% against D-JEPA 76.76%. Physical trials are 50 pairs per task. Mean Chamfer on Granular does not beat DINO-WM. Cube is near-saturated; the paper uses it as coverage, not as a gain.
The 100% native-distance recovery after exact ordinal realization checks that ranks were written into radii correctly. It does not show that a new geometry was learned.