Aligning predicted futures with better actions.

A Decision-Aligned Latent World Model

Shuaijun Liu¹Chengyu Wu¹Qifu Wen^2,3Feiyang You¹Chenglong Zhang¹Shuyang Hao¹Xi Lin³Ningxin Su^1,*

¹The Hong Kong University of Science and Technology (Guangzhou)

²Boston University ³Shanghai Jiao Tong University

^*Corresponding author: ningxinsu@hkust-gz.edu.cn

D-JEPA learns decision-relevant relations across pretrained predictive models and realizes them in future representations.

Interactive explainer

Paper Model repository Supervision dataset Code

01 / IN MOTION

Same start. Same goal. A different choice of action.

Matched executions across manipulation and driving, with Cube included as saturated-task coverage. Select a task or scene to watch.

PushT

TD-JEPA / D-JEPA

25 controls, including the complete action sequence and terminal state.

View the timeline ↗

Reacher

TD-JEPA / D-JEPA

25 controls. The target specifies the final joint configuration.

View the timeline ↗

Granular manipulation

DINO-WM / D-JEPA

Five actions with settling steps, showing the evolving particle distribution.

View the timeline ↗

Bimanual grasping

Native VLA / D-JEPA

Matched initial state, complete recorded sequences at 2× playback. The aligned execution stops on success; its final frame is held while the native sequence continues.

View the timeline ↗

Autonomous driving

Drive-JEPA / D-JEPA

Recorded camera context and simulated candidate trajectories, synchronized over four seconds.

View the timeline ↗

Unseen object geometry

Native / Fusion / D-JEPA

An unseen I-shaped object, with the same goal and action candidates.

View the timeline ↗

Changed object appearance

Native / Fusion / D-JEPA

Object-colour change after task-local fitting, evaluated on a fresh start.

View the timeline ↗

Driving · turning comfort

Drive-JEPA / D-JEPA

The appendix turning case: D-JEPA satisfies the comfort criterion; both candidates satisfy the collision and TTC criteria.

View the timeline ↗

Driving · progress

Drive-JEPA / D-JEPA

The appendix progress case: D-JEPA travels farther through the same traffic scene. Logged camera context accompanies paired simulated trajectories.

View the timeline ↗

Cube · task coverage

LeWM / TD-JEPA

Both predictive baselines succeed. This saturated-task example documents coverage; it is not a D-JEPA comparison.

View the timeline ↗

02 / THE QUESTION

A predicted future is useful when it helps choose the action.

The model prefers one future. Execution favors another. A candidate closer to the goal in predicted latent distance can fail, while a slightly farther alternative succeeds. Across matched PushT decisions, ranking fidelity also drops sharply within the shortlist from which actions are selected. D-JEPA learns the relations that distinguish these consequential alternatives.

Our starting point

Learn the decision-relevant relations between candidate futures, then use those relations to align action selection.

The decision boundary is a learning target—not just a property inherited from pretraining.

Recorded LeWM distances for a failing and a successful candidate — (a) Recorded distances

Normalized success–failure distance gaps across 96 starts — (a) Recorded distances

D-JEPA overview: predicted goal distance can favor a failing action; relational learning and geometry write-back align decisions across physical tasks. — **From predictive mismatch to decision-aligned geometry.** The problem, the learning framework and its cross-task evaluation. Click the drawing to view the original PDF ↗

03 / THE FRAMEWORK

From future prediction to decision-aligned representation.

Go inside D-JEPATrace candidate futures through relational alignment and into decision-aligned representations.

Shared candidate futures feed restricted predictor adaptation and bounded relational alignment, with the set-attention and correction architecture shown in detail. — **Learning decision-relevant relations over candidate futures.** Restricted predictor adaptation and a bounded relational operator use the same candidate identities. Shared descriptors and set attention connect predictive evidence to action selection. Click the drawing to view the original PDF ↗

Decision alignment

Learn bounded, set-wise corrections from future–goal descriptors and ordinal structure. Restricted predictor adaptation offers a complementary source of decision evidence.

Explore the modules ↗

Multiple predictive geometries

Combine relational evidence from complementary predictors in one alignment framework, extending beyond a single predictive distance.

Read the protocols ↗

Representation realization

Encode a learned decision ordering in JEPA-compatible goal-distance geometry, and learn bounded, same-action updates across future time steps.

Explore checkpoints ↗

Multi-geometry alignment combines scale-free source ranks; calibrated composition admits a reliable plastic proposal over the relational default while retaining candidate identities. — **Combining geometries and proposals.** Ordinal evidence connects complementary predictive spaces. A calibrated gate combines the relational decision with predictor adaptation to produce the final candidate order. Click the drawing to view the original PDF ↗

Two complementary representation mechanisms

PushT exact realization preserves earlier futures and encodes terminal ordering while retaining its embedded predictive and relational computation. Reacher learns same-action five-step future updates; its reported physical decisions use the relational selector. These experiments establish complementary mechanisms and do not imply a newly evaluated combined checkpoint or single-backbone distillation.

Representation-lifting architecture: a bounded coefficient network produces same-action temporal residuals, while strict ranks set the native terminal goal-distance radius. — **Realizing decision structure through native future geometry.** Learned temporal transport and deterministic ordinal realization expose their inputs, computation and future-latent outputs. The geometric illustration is retained in Appendix D.2 of the paper. Click the drawing to view the original PDF ↗

04 / RESULTS

Better future relations. Better executed decisions.

Matched candidate sets isolate the quality of action selection across simulation, foundation models, autonomous driving and physical manipulation.

PUSHT · SUCCESS83.59→87.89%

+4.30 pp over LeWM
PushT confirmation evaluation

ROBOTWIN · SUCCESS61.72→76.76%

+15.04 pp over native VLA selection
Four language-conditioned tasks

REAL ROBOT · MEAN SUCCESS RATE64.0→81.0%

+17.0 pp over V-JEPA 2-AC
PushT and two-cube stacking

DRIVING · PDMS57.36→95.34

+37.98 over Drive-JEPA ranking
Seven difficult scenes

Core action-selection results
Task	Evaluation	Matched reference	D-JEPA	Gain
PushT	Success · 256 starts	LeWM · 83.59%	87.89%	+4.30 pp
DMC-Reacher	Success · 128 starts	LeWM · 86.72%	93.75%	+7.03 pp
Granular	Strict attainment · 64 starts	DINO-WM · 7.81%	18.75%	+10.94 pp
Each task retains its own evaluation population and success criterion. Granular reports attainment at the stricter Chamfer-distance threshold (0.12).

Cross-domain and generalization evidence
Setting	Metric / scale	Matched reference	D-JEPA	Gain
Held-out PushObj shapes	Success · 300 starts	Calibrated fusion · 44.67%	53.67%	+9.00 pp
Appearance changes	Mean success · 7 conditions	Calibrated fusion · 62.86%	73.14%	+10.29 pp
RoboTwin	Success rate · 4 tasks	Native VLA · 61.72%	76.76%	+15.04 pp
Physical PiPER	Mean success rate · 2 tasks	V-JEPA 2-AC · 64.0%	81.0%	+17.0 pp
Autonomous driving	Mean PDMS · 7 scenes	Drive-JEPA · 57.36	95.34	+37.98
Appearance changes are evaluated on 50 new base starts per condition; RoboTwin uses 128 starts per task and PiPER 50 paired trials per task. PiPER success is averaged equally across its two tasks. Driving reports mean PDMS over seven difficult scenes.

Each row retains its own task metric and evaluation population. Full counts, paired outcomes and task-specific costs remain available in the paper and the decision-supervision dataset.

Results and analyses

PushT confirmation success over 256 starts — PushT · confirmation 256

Reacher formal success by method — Reacher · independent 128

Granular success across distance thresholds — Granular · thresholds

Paired DINO-WM and D-JEPA Granular costs on the same starts — Granular · paired cost

Success on unseen I, small T and square shapes — Held-out shapes · success

Success across seven visual conditions — Appearance changes · success

Paired gains and losses on each held-out shape — Held-out shapes · paired outcomes

Net gains under each visual condition — Appearance changes · net gains

PushT confirmation success with plasticity, relational alignment and calibrated composition — Decision mechanisms · success

Success rates and intervals for multiple predictive geometries — Multiple geometries · success

Median per-start inference latency of the three deployment configurations — Deployment · inference latency

Success across shared candidate budgets — Candidate availability

Recovery of relational ranks through native distance — Exact ordinal recovery

Shared projection of future representations before and after realization — Future–goal geometry

Learned same-action residual coefficients across five future steps — Same-action temporal transport

Click any panel for the full-size vector figure. TD: TD-JEPA; Ours / D: D-JEPA; N: native selection; F: calibrated fusion. Success panels show 95% intervals; paired panels retain gains and losses. Multi-geometry results use 128 mechanism-evaluation starts; PushT deployment and decision ablations use 256 confirmation starts. Budget bands span 16 subset seeds; temporal transport is the Reacher five-step diagnostic.

Decision mechanisms and native-distance deployment
Configuration	Decision readout	Success	Median inference
Relational alignment	Relational score	87.11%	35.03 ms
Calibrated composition	Gated composition	87.89%	52.16 ms
Ordinal realization	Native goal distance	87.11%	34.02 ms
PushT confirmation evaluation, 256 starts. Interface verification: native goal distance recovers 100% of relational action choices and candidate ranks after ordinal realization. Timing includes the predictive and relational computation, measured per start.

Matched RoboTwin baseline failure and D-JEPA grasp success — **RoboTwin.** The same observation and candidate set produce different executed outcomes.

Matched driving candidates selected by Drive-JEPA and D-JEPA — **Autonomous driving.** Camera frames are shared logged observations; bird’s-eye views show the alternative simulated ego trajectories against logged traffic.

The same context and goal under seven visual conditions — **Seven visual conditions.** Clean, blur, salt-and-pepper noise, darkening, object colour, goal-marker colour and pusher colour share the same physical start and goal. These are the model’s recorded inputs.

LeWM and TD-JEPA both complete the Cube task from the same start — **Cube task coverage.** Both LeWM and TD-JEPA attain the goal despite selecting different actions. Success is near saturation across the task screen, leaving little headroom for binary success to distinguish action selection.

05 / BUILD ON D-JEPA

From a checkpoint to a reproducible decision.

Task-specific checkpoints, decision supervision and a compact evaluation package.

Task interfaces: robotic manipulation · autonomous driving · physical-robot workflows. Each guide describes inputs, training and evaluation commands, and environment dependencies.

Full reproduction guide ↗

Run from the repository root

pip install -e '.[download]'
python scripts/download_artifacts.py --profile pusht-relational --dataset
python -m zipfile -e data/D-JEPA-supervision-v1.zip data
bash scripts/reproduce.sh