Does LeWorldModel Preserve Local Action Geometry?

Vasil Vasilev; Siddharth Nayak

cs.LGOct 22, 2026

Does LeWorldModel Preserve Local Action Geometry?

Authors: Vasil Vasilev, Siddharth Nayak

Organizations: Technological School Electronic Systems (TUES) · Massachusetts Institute of Technology

Abstract

We examine how closely pretrained LeWorldModel (LeWM) reproduces the effects of small action changes. Using matched simulator and model responses, we measure action sensitivity with Gram matrices at three perturbation scales. Across nearly a thousand states from Reacher and OGBench Cube, the most sensitive action directions differ by median angles of about $36^\circ$ and $42^\circ$ , respectively, at the smallest scale. Increasing the perturbation tenfold reduces median normalized Gram discrepancy by roughly 70% on Reacher and 47% on Cube. Geometry measured in pixels and embeddings varies more across scales than the chosen physical-state geometry. On held-out directions, environment Grams achieve median embedding response errors of 17% and 28%, compared with 26% and 38% for an isotropic reference with the same trace. LeWM's errors are higher at this scale. These findings support scale-specific evaluation of action sensitivity; its effect on control remains an open question.

Figures & tables

Environment	Model	Isotropic	Rotation	Shuffle
Reacher	0.465	0.449	0.632	0.530
Cube	0.508	0.609	0.885	0.560

Table 1

Figure 2

Environment	Env. Gram	Isotropic	Model	Model + oracle gain
Reacher	17.0	26.2	33.2	26.6
Cube	28.0	38.2	56.4	32.3

Table 3

Appendix figures & tables6 assets

Supplementary material from the paper’s appendix.

Appendix

Env.	$d_{a}$	$D_{G}$	$\kappa_{\rm env}$	$\kappa_{\rm model}$	Top-1	Top-2
Reacher	2	$0.465\,[0.270,0.708]$	4.49	4.50	$36.03\,[18.04,53.85]^{\circ}$	–
Cube	5	$0.508\,[0.397,0.633]$	$2.84{\times}10^{3}$	$1.07{\times}10^{4}$	$41.52\,[25.03,66.09]^{\circ}$	$40.27\,[27.13,58.12]^{\circ}$

Appendix

Table 4

$\epsilon$	Median $D_{G}$	Top-1 angle	Top-2 max angle
$10^{-3}$	0.508	$41.52^{\circ}$	$40.27^{\circ}$
$3{\times}10^{-3}$	0.408	$34.82^{\circ}$	$30.38^{\circ}$
$10^{-2}$	0.271	$20.96^{\circ}$	$18.48^{\circ}$

Appendix

Table 5

Figure 6

Environment	$\epsilon$	Env. Gram	Isotropic	Model	Oracle gain	One-sided
Reacher	0.001	17.0	26.2	33.2	26.6	32.2
Reacher	0.003	11.3	22.2	20.6	18.0	23.9
Reacher	0.01	6.8	20.2	10.3	9.2	13.8
Cube	0.001	28.0	38.2	56.4	32.3	42.6
Cube	0.003	21.6	34.5	39.1	25.1	33.9
Cube	0.01	13.5	32.9	23.9	16.9	22.7

Appendix

Table 7

Environment	Output metric	Gram change	WAPE
Reacher	Physical pose	$9.32\times 10^{-7}$	$2.97\times 10^{-9}$
Reacher	Pose + velocity	$2.11\times 10^{-6}$	$6.37\times 10^{-9}$
Reacher	Pixels	0.1125	0.0498
Reacher	Embedding	0.4331	0.1702
Cube	Physical pose	0.0270	0.0004
Cube	Pose + velocity	0.0266	0.0004

Appendix

Table 8

Figure 9

World Models

World models represent and predict how environments change, helping agents plan, act, and reason about possible futures.

See all 366 papers

Reperesentation Geometry Matters for Planning with JEPA World Models
Aug 17, 2026Jiaming Hu, Yan Zheng, Shi Bo +9Representation Learning Neural Representation Geometry
Boston University · Unity Technologies
Same Goal, Different Words: Spatial Instruction Sensitivity in World-Action Models
Oct 12, 2026Ali-Adeeb Abbas, Esmaeil Seraj, Behrad Toghi +1World Model Evaluation Vision-Language-Action Models
RIVeR Laboratory, Northeastern University, Boston, MA, USA. · General Motors Robotics, Mountain View, CA, USA.
PrismWAM: Decoupling State Representations in World Action Models
Oct 19, 2026Tianyu Xin, Wanqi Zhong, Bowen Hu +8World Models for Robotics World Models
Tsinghua University · NVIDIA

Aug 17, 2026J/K move · Enter open · S save

Reperesentation Geometry Matters for Planning with JEPA World Models

Jiaming Hu, Yan Zheng, Shi Bo, Tian Wang+8

Boston University · Unity Technologies

Representation Learning Neural Representation Geometry World Models for Robotics World Model-Based Planning Joint-Embedding Predictive Architecture Model-Based Planning

Open paper AskPDF arXiv

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.

Abstract

World Models

Explore similar work

Reperesentation Geometry Matters for Planning with JEPA World Models

Same Goal, Different Words: Spatial Instruction Sensitivity in World-Action Models

PrismWAM: Decoupling State Representations in World Action Models