Reperesentation Geometry Matters for Planning with JEPA World Models

Jiaming Hu; Yan Zheng; Shi Bo; Tian Wang; Florian Dubost; Alejandro Mottini; Junze Liu; Arvind Srinivasan

cs.LGAug 17, 2026

Reperesentation Geometry Matters for Planning with JEPA World Models

Authors: Jiaming Hu, Yan Zheng, Shi Bo, Tian Wang, Florian Dubost, Alejandro Mottini, Junze Liu, Arvind Srinivasan, +4 more

Organizations: Boston University · Unity Technologies

Abstract

Joint-embedding predictive world models support planning through latent predictions, but unconstrained joint training can collapse distinct observations to identical embeddings. Two prominent strategies for avoiding collapse are to inherit pretrained features, as in DINO-WM, or to learn representations end-to-end with anti-collapse regularization, as in LeWorldModel (LeWM). Yet avoiding collapse does not ensure that latent distances distinguish outcomes in ways that matter for the task. In object manipulation, for example, success depends on the object's position and orientation relative to the goal. Such task-relevant state information can remain accurately decodable while barely influencing latent distance. The resulting planning cost may fail to reflect how close a predicted outcome is to the task goal. In this paper, we propose SCALE (State-CAlibrated Latent Embeddings), a method that correlates sampled pairwise latent distances with distances in task-relevant state space. Added to LeWM's existing objective, SCALE preserves its architecture, requires privileged state only during training, and adds no planning-time computation. We show that SCALE improves planning success over LeWM across manipulation and navigation tasks with multiple solvers and provide a comprehensive analysis of how SCALE reshapes representation geometry to support planning.

Figures & tables

Figure 1: (a) A fixed encoder and predictor support action search using the terminal–goal latent distance. The planner executes the first action block and replans. (b) During training, SCALE correlates sampled latent and task-state distance using the correlation-based regularizer.

Task	LeWM	DINO-WM
Push-T	.13	.70
Reacher	.18	.46
Cube	.00	.46
Two-Room	.41	.85
PointMaze	.52	.50

Table 1: Spearman correlation

\rho_{\mathrm{S}}

between latent and task-state distances for LeWM and DINO-WM on held-out frame pairs.

Task	Method	T1	T2	T3	T4	T5	Overall
Push-T	LeWM	$86.0{\scriptstyle\pm}2.4$	$79.8{\scriptstyle\pm}2.1$	$73.4{\scriptstyle\pm}5.9$	$55.2{\scriptstyle\pm}4.8$	$25.1{\scriptstyle\pm}5.1$	$63.9{\scriptstyle\pm}3.0$
	SCALE	$89.1{\scriptstyle\pm}2.0$	$86.3{\scriptstyle\pm}2.8$	$79.7{\scriptstyle\pm}3.6$	$61.3{\scriptstyle\pm}6.2$	$25.2{\scriptstyle\pm}5.5$	$68.3{\scriptstyle\pm}2.8$
	Aux	$89.9{\scriptstyle\pm}2.0$	$86.2{\scriptstyle\pm}2.3$	$79.6{\scriptstyle\pm}3.4$	$62.2{\scriptstyle\pm}5.3$	$27.6{\scriptstyle\pm}5.2$	$69.1{\scriptstyle\pm}2.6$
	DINO-WM	$70.0{\scriptstyle\pm}4.5$	$65.3{\scriptstyle\pm}4.5$	$58.8{\scriptstyle\pm}5.1$	$41.5{\scriptstyle\pm}8.8$	$23.1{\scriptstyle\pm}3.7$	$51.7{\scriptstyle\pm}4.3$
Reacher	LeWM	$84.9{\scriptstyle\pm}2.6$	$84.2{\scriptstyle\pm}3.4$	$83.8{\scriptstyle\pm}3.4$	$79.0{\scriptstyle\pm}6.1$	$62.2{\scriptstyle\pm}3.6$	$78.8{\scriptstyle\pm}1.7$
	SCALE	$86.7{\scriptstyle\pm}3.6$	$87.0{\scriptstyle\pm}3.1$	$87.3{\scriptstyle\pm}3.2$	$81.7{\scriptstyle\pm}3.0$	$64.4{\scriptstyle\pm}3.0$	$81.4{\scriptstyle\pm}2.0$

Table 2: Pooled iCEM success rates (%) across tasks and compute tiers. Each entry reports mean

\pm

SD over

12

set-level measurements: two independently trained runs, each evaluated on six held-out sets of 100 episodes. Overall pools T1–T5. Bold marks the best method per column within each task. Appendix E reports both CEM and iCEM results.

Panel of Figure 2 — Figure 2: Normalized PCA eigenvalue spectra of LeWM, SCALE, Aux, and DINO-WM across five tasks. Eigenvalues are normalized by total variance, so each principal component’s share corresponds to its relative metric leverage. SCALE concentrates more variance in the leading principal components than LeWM and Aux on most tasks.

Panel of Figure 3 — Figure 3: Task-state information in leading principal components. Held-out $R^{2}$ for predicting task-relevant state from the leading k principal components across five tasks. SCALE improves state predictability compared with LeWM and outperforms Aux on four of five tasks. This suggests that SCALE concentrates task-relevant information in directions with greater metric leverage.

Task	Metric	LeWM	SCALE	Aux	DINO-WM
Push-T	RollErr $\downarrow$	.0350	.0277 ( $-21.0\%$ )	.0269 ( $-23.1\%$ )	.2819 ( $+704.9\%$ )
Push-T	Kendall $\tau$ $\uparrow$	.7316	.7857 ( $+7.4\%$ )	.7576 ( $+3.6\%$ )	.7000 ( $-4.3\%$ )
Reacher	RollErr $\downarrow$	.1812	.1548 ( $-14.6\%$ )	.1686 ( $-7.0\%$ )	.5860 ( $+223.4\%$ )
Reacher	Kendall $\tau$ $\uparrow$	.6233	.6435 ( $+3.2\%$ )	.6328 ( $+1.5\%$ )	.6414 ( $+2.9\%$ )
Cube	RollErr $\downarrow$	.3195	.3174 ( $-0.7\%$ )	.2389 ( $-25.2\%$ )	.3227 ( $+1.0\%$ )
Cube	Kendall $\tau$ $\uparrow$	.3654	.4433 ( $+21.3\%$ )	.4150 ( $+13.6\%$ )	.7775 ( $+112.8\%$ )

Table 3: Within-model rollout quality. Lower rollout error and higher Kendall

\tau

are better. Parentheses report change relative to LeWM; when LeWM’s

\tau

is near zero, we report the absolute difference

\Delta

. Results average 20 starts.

Appendix figures & tables14 assets

Supplementary material from the paper’s appendix.

Appendix

Task	$\lambda_{\mathrm{corr}}{=}0$	$0.03$	$0.1^{\dagger}$	$0.3$	$1.0$
Push-T	63.9 [-1pt]86/81/74/54/26	65.8 [-1pt]88/83/76/55/26	68.2 [-1pt]90/87/79/60/26	67.4 [-1pt]89/85/78/58/26	68.3 [-1pt]91/86/80/58/26
Reacher	78.0 [-1pt]86/82/83/78/62	79.6 [-1pt]84/85/85/81/64	81.1 [-1pt]87/86/87/82/64	77.3 [-1pt]82/84/81/78/62	79.4 [-1pt]84/84/86/80/64
Cube	65.3 [-1pt]71/68/70/63/55	67.8 [-1pt]73/74/72/64/57	70.5 [-1pt]77/76/75/65/60	70.8 [-1pt]77/77/75/65/60	70.1 [-1pt]75/76/76/65/59
Two-Room	82.9 [-1pt]91/93/89/78/63	91.7 [-1pt]100/100/96/88/76	92.9 [-1pt]100/100/97/90/77	94.3 [-1pt]100/100/98/92/81	95.5 [-1pt]100/100/98/95/84
PointMaze	80.2 [-1pt]82/83/81/82/74	91.7 [-1pt]97/96/94/88/83	83.4 [-1pt]87/83/84/83/79	80.2 [-1pt]84/83/78/80/76	75.2 [-1pt]79/76/75/76/71

Appendix

Table 4: iCEM sensitivity to the SCALE correlation weight. Each cell reports the success rate (%) averaged across T1–T5 on the first line, followed by the five tier values.

\dagger

marks the value used in the main experiments; its averages are bold. We use a single value of

0.1

across all tasks for a controlled comparison. Selecting

\lambda_{\mathrm{corr}}

separately for each task gives higher averages on four of five tasks. Results are from a single training run.

Model	Training (it/s)	CEM (ms/plan)	iCEM (ms/plan)	Planning vs. LeWM
LeWM	$4.35$	$1{,}213.3$	$1{,}173.0$	$1.00\times$
SCALE	$3.85$	$1{,}195.8$	$1{,}184.9$	$1.00\times$
Aux	$4.20$	$1{,}229.4$	$1{,}218.8$	$1.03\times$
DINO-WM	$8.95$	$91{,}290.7$	$91{,}654.3$	$76.7\times$

Appendix

Table 5: Training throughput and median per-plan latency on Cube. Planning uses the largest matched compute tier T1.

Probe	State component	LeWM	SCALE	Aux	DINO-WM
Linear	agent $x$	.9191	.9656	.9851	.9342
	agent $y$	.9194	.9715	.9845	.9394
	block $x$	.9686	.9820	.9904	.9291
	block $y$	.9603	.9746	.9886	.9722
	${\color[rgb]{0,0,0}\cos\vartheta}$	.9019	.9023	.9767	.8353
	${\color[rgb]{0,0,0}\sin\vartheta}$	.8694	.8855	.9778	.5943

Appendix

Table 6: Per-component held-out

R^{2}

for Push-T.

\vartheta

is the block angle.

Probe	State component	LeWM	SCALE	Aux	DINO-WM
Linear	${\color[rgb]{0,0,0}\cos\vartheta_{0}}$	.9996	.9998	.9996	.9972
	${\color[rgb]{0,0,0}\sin\vartheta_{0}}$	.9997	.9998	.9998	.9981
	${\color[rgb]{0,0,0}\cos\vartheta_{1}}$	.9994	.9997	.9995	.9953
	${\color[rgb]{0,0,0}\sin\vartheta_{1}}$	.9980	.9996	.9986	.9702
	Mean	.9992	.9997	.9994	.9902
MLP	${\color[rgb]{0,0,0}\cos\vartheta_{0}}$	.9837	.9875	.9854	.9479

Appendix

Table 7: Per-component held-out

R^{2}

for Reacher.

\vartheta_{0}

and

\vartheta_{1}

denote the two joint angles; the four displayed sine and cosine components constitute the selected joint-only state

\mathbf{q}

Probe	State component	LeWM	SCALE	Aux	DINO-WM
Linear	effector $x$	.9773	.9854	.9934	.9814
	effector $y$	.9661	.9778	.9943	.9980
	effector $z$	.9825	.9899	.9928	.9818
	${\color[rgb]{0,0,0}\cos 2\varphi}$	.7599	.9888	.9938	.7627
	${\color[rgb]{0,0,0}\sin 2\varphi}$	-.2577	.9880	.9938	.4213
	gripper	.9072	.9783	.9742	.6638

Appendix

Table 8: Per-component held-out

R^{2}

for Cube.

\varphi

denotes the gripper yaw; the doubled-angle encoding accounts for the parallel gripper’s

\pi

symmetry. The nine displayed components constitute the selected state

\mathbf{q}

Probe	State component	LeWM	SCALE	Aux	DINO-WM
Linear	$x$	.9957	.9982	.9963	.9957
	$y$	.9942	.9977	.9981	.9975
	Mean	.9950	.9979	.9972	.9966
MLP	$x$	.9581	.9742	.9722	.9315
	$y$	.9660	.9829	.9714	.9209
	Mean	.9621	.9786	.9718	.9262

Appendix

Table 9: Per-component held-out

R^{2}

for Two-Room.

Probe	State component	LeWM	SCALE	Aux	DINO-WM
Linear	$x$	.9988	.9992	.9989	.9982
	$y$	.9989	.9990	.9989	.9982
	Mean	.9989	.9991	.9989	.9982
MLP	$x$	.9703	.9869	.9717	.9288
	$y$	.9824	.9733	.9796	.9609
	Mean	.9764	.9801	.9757	.9449

Appendix

Table 10: Per-component held-out

R^{2}

for PointMaze.

Task	SCALE state $\mathbf{q}$	$d_{q}$	CEM	iCEM
Push-T	Selected subset	6	$73.9{\scriptstyle\pm}2.2$ [-1pt] $96{\scriptstyle\pm}2.0/93{\scriptstyle\pm}2.6/85{\scriptstyle\pm}4.0/63{\scriptstyle\pm}4.6/32{\scriptstyle\pm}4.4$	$68.2{\scriptstyle\pm}2.3$ [-1pt] $90{\scriptstyle\pm}2.0/87{\scriptstyle\pm}2.5/79{\scriptstyle\pm}3.9/60{\scriptstyle\pm}4.7/26{\scriptstyle\pm}5.8$
Push-T	Full state	8	$74.5{\scriptstyle\pm}2.6$ [-1pt] $94{\scriptstyle\pm}2.9/92{\scriptstyle\pm}2.6/85{\scriptstyle\pm}3.0/67{\scriptstyle\pm}5.0/34{\scriptstyle\pm}5.8$	$68.3{\scriptstyle\pm}2.8$ [-1pt] $90{\scriptstyle\pm}1.6/87{\scriptstyle\pm}1.8/79{\scriptstyle\pm}4.8/59{\scriptstyle\pm}6.9/26{\scriptstyle\pm}1.2$
Reacher	Selected subset	4	$71.5{\scriptstyle\pm}3.9$ [-1pt] $82{\scriptstyle\pm}4.6/79{\scriptstyle\pm}3.8/75{\scriptstyle\pm}5.9/64{\scriptstyle\pm}8.4/57{\scriptstyle\pm}6.1$	$81.1{\scriptstyle\pm}2.0$ [-1pt] $87{\scriptstyle\pm}3.6/86{\scriptstyle\pm}2.9/86{\scriptstyle\pm}3.4/82{\scriptstyle\pm}3.2/64{\scriptstyle\pm}2.5$
Reacher	Full state	8	$72.1{\scriptstyle\pm}1.0$ [-1pt] $82{\scriptstyle\pm}2.4/81{\scriptstyle\pm}1.0/73{\scriptstyle\pm}2.8/66{\scriptstyle\pm}3.4/59{\scriptstyle\pm}6.1$	$82.3{\scriptstyle\pm}1.4$ [-1pt] $89{\scriptstyle\pm}2.3/86{\scriptstyle\pm}2.4/89{\scriptstyle\pm}2.9/84{\scriptstyle\pm}4.9/64{\scriptstyle\pm}2.7$
Cube	Selected subset	9	$64.1{\scriptstyle\pm}4.0$ [-1pt] $74{\scriptstyle\pm}3.4/71{\scriptstyle\pm}4.2/66{\scriptstyle\pm}4.8/56{\scriptstyle\pm}5.2/54{\scriptstyle\pm}5.0$	$70.5{\scriptstyle\pm}3.3$ [-1pt] $77{\scriptstyle\pm}3.9/76{\scriptstyle\pm}4.5/75{\scriptstyle\pm}3.5/65{\scriptstyle\pm}3.7/60{\scriptstyle\pm}4.5$
Cube	Full state	22	$64.8{\scriptstyle\pm}3.1$ [-1pt] $76{\scriptstyle\pm}5.0/71{\scriptstyle\pm}3.5/64{\scriptstyle\pm}3.9/58{\scriptstyle\pm}4.1/55{\scriptstyle\pm}4.8$	$69.7{\scriptstyle\pm}3.4$ [-1pt] $76{\scriptstyle\pm}3.7/76{\scriptstyle\pm}4.2/74{\scriptstyle\pm}4.1/64{\scriptstyle\pm}4.5/57{\scriptstyle\pm}3.5$

Appendix

Table 11: Sensitivity to the task-state choice used by SCALE. Each CEM and iCEM cell reports the success rate (%) averaged across T1–T5 on the first line, followed by the five tier values (mean

\pm

SD). Results are from a single training run.

Task	LeWM	SCALE	Aux	DINO-WM
Push-T	.13	.77	.18	.70
Reacher	.18	.53	.17	.46
Cube	.00	.56	.03	.46
Two-Room	.41	.77	.41	.85
PointMaze	.52	.80	.52	.50

Appendix

Table 12: Held-out latent–state distance rank alignment (Spearman

{\color[rgb]{0,0,0}\rho_{\mathrm{S}}}

) over

1.12

M frame pairs per task. LeWM and DINO-WM repeat Table 1 for comparison.

Solver	Method	T1	T2	T3	T4	T5	Overall
CEM	LeWM	$92.3{\scriptstyle\pm}2.3$	$89.3{\scriptstyle\pm}3.2$	$80.5{\scriptstyle\pm}4.1$	$57.8{\scriptstyle\pm}5.6$	$33.8{\scriptstyle\pm}3.6$	$70.7{\scriptstyle\pm}2.7$
	SCALE	$94.8{\scriptstyle\pm}3.2$	$92.1{\scriptstyle\pm}2.6$	$84.2{\scriptstyle\pm}3.6$	$61.6{\scriptstyle\pm}6.9$	$32.8{\scriptstyle\pm}5.5$	$73.1{\scriptstyle\pm}2.7$
	Aux	$94.5{\scriptstyle\pm}2.0$	$92.6{\scriptstyle\pm}2.4$	$84.4{\scriptstyle\pm}4.1$	$65.1{\scriptstyle\pm}6.9$	$34.9{\scriptstyle\pm}6.4$	$74.3{\scriptstyle\pm}3.4$
	DINO-WM	$74.8{\scriptstyle\pm}4.6$	$72.1{\scriptstyle\pm}4.5$	$65.2{\scriptstyle\pm}4.9$	$46.1{\scriptstyle\pm}4.6$	$24.3{\scriptstyle\pm}6.1$	$56.5{\scriptstyle\pm}4.0$
iCEM	LeWM	$86.0{\scriptstyle\pm}2.4$	$79.8{\scriptstyle\pm}2.1$	$73.4{\scriptstyle\pm}5.9$	$55.2{\scriptstyle\pm}4.8$	$25.1{\scriptstyle\pm}5.1$	$63.9{\scriptstyle\pm}3.0$
	SCALE	$89.1{\scriptstyle\pm}2.0$	$86.3{\scriptstyle\pm}2.8$	$79.7{\scriptstyle\pm}3.6$	$61.3{\scriptstyle\pm}6.2$	$25.2{\scriptstyle\pm}5.5$	$68.3{\scriptstyle\pm}2.8$

Appendix

Table 13: Pooled Push-T success rates (%). Mean

\pm