SWIM: Learning What Matters from Behaviors for JEPA World Models

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

cs.ROOct 27, 2026

SWIM: Learning What Matters from Behaviors for JEPA World Models

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

Organizations: Boston University · Unity Technologies

Abstract

Joint-embedding predictive architectures (JEPAs) enable world models to plan from images by predicting future embeddings. However, their visual prediction objectives do not explicitly incorporate task-relevant state information, so the physical state that determines task progress may not be reliably captured in the latent space. Supervising world models with simulator state can make planning more informative, but it requires deciding which state features matter. Features are typically selected by hand or weighted uniformly, and a uniform weighting gives every standardized state feature equal influence, so redundant coordinates can overcount the same variation, while irrelevant ones can obscure task-relevant differences. To address this limitation, we propose SWIM (State Weighting with Inverse-Dynamics Model), which learns state weights from offline trajectories by predicting recorded actions from current states and observed future displacements. The learned weights quantify each feature's contribution to action prediction and define a weighted state metric for state-supervised training, which we demonstrate by aligning latent and state distances. The inverse-dynamics model is used only during training, so the resulting world model plans from images with no additional planning-time computation. Across five simulated robotic manipulation tasks, SWIM consistently improves planning success over uniform state weighting.

Figures & tables

Figure 1: SWIM. (a) The behavioral predictor learns coordinate weights from current states, gated future changes, and logged actions. (b) Frozen weights rescale standardized states. Squared latent and state distances from the same frame pairs enter the correlation loss, with state distances held fixed; prediction and SIGReg losses are retained. The inset names physical features before standardization and shows illustrative weight bars; within this inset, $\theta$ denotes a joint angle. The training symbols are defined in Sections 4.1 – 4.2 . Deployment uses image-based latent planning; the behavioral predictor and logged future states are training-only components.

Task	DINO-WM †	LeWM	SCALE-All	SWIM-H ⋆	SWIM-N ⋆	$\Delta_{\text{SCALE-All}}$	$\Delta_{\text{LeWM}}$
Cube	72.5	67.7 $\pm$ 1.0	74.6 $\pm$ 1.1	75.9 $\pm$ 1.5 (.01)	75.8 $\pm$ 1.2 (.03)	+1.3	+8.3
Cube-Double	17.5	16.4 $\pm$ 1.1	18.2 $\pm$ 0.8	19.4 $\pm$ 1.1 (1.0)	19.3 $\pm$ 0.7 (1.0)	+1.1	+3.0
Cube-Triple	34.2	30.6 $\pm$ 0.8	31.3 $\pm$ 1.2	32.6 $\pm$ 1.0 (.05)	32.4 $\pm$ 0.5 (.03)	+1.3	+2.0
Cube-Quadruple	30.2	28.3 $\pm$ 1.0	29.6 $\pm$ 0.7	30.7 $\pm$ 0.9 (.03)	30.5 $\pm$ 0.6 (1.0)	+1.0	+2.3
Scene	22.0	30.5 $\pm$ 0.5	33.4 $\pm$ 1.9	37.3 $\pm$ 3.3 (.1)	39.4 $\pm$ 2.3 (.1)	+6.0	+8.9

Table 1: CEM planning success (%) over five training seeds (mean

\pm

standard deviation). SWIM-H ⋆ and SWIM-N ⋆ use the best

\lambda

for each task, shown in parentheses. DINO-WM † uses pretrained visual features. The

\Delta

columns compare the better SWIM variant with SCALE-All and LeWM. Bold: best; underline: second best.

Figure 2: Learned gates on Cube (InfoNCE with $\lVert\mathbf{g}\rVert_{1}$ regularization). (a) Gate values for the 22 candidate state dimensions as the sparsity coefficient $\lambda$ increases. End-effector height ( eff_z ), block height ( block_z ), and the arm joints ( cos_j1 , cos_j3 , sin_j3 ) are retained longest, while sin_j1 and sin_j2 are suppressed first. (b) Gates learned at $\lambda=0.1$ when 26 structured variables from unrelated Scene trajectories are appended to the Cube state (48 candidate dimensions). Every distractor receives a gate of at most 0.31, and the five Cube dimensions at or above 0.5 are exactly the five retained at $\lambda=0.1$ on the clean 22-dimensional state.

Appendix figures & tables11 assets

Supplementary material from the paper’s appendix.

Appendix

Feature	$\ell_{1}$ + Hinge				$\ell_{1}$ + InfoNCE
	$\lambda=0.01$	$0.03$	$0.05$	$0.10$	$\lambda=0.01$	$0.03$	$0.05$	$0.10$
b0_y	0.94	0.87	0.76	0.51	0.95	0.91	0.86	0.66
b1_y	0.94	0.86	0.75	0.50	0.95	0.92	0.87	0.67
b0_x	0.92	0.79	0.63	0.39	0.93	0.88	0.79	0.55
b1_x	0.91	0.78	0.62	0.39	0.93	0.88	0.79	0.56
eff_z	0.88	0.84	0.80	0.68	0.88	0.85	0.81	0.71

Appendix

Table 2: Stage-1 gate values on Cube-Double (27 dimensions). Values are continuous gates learned with

\ell_{1}

regularization and the two discrimination objectives across the displayed

\lambda

values.

Feature	$\ell_{1}$ + Hinge				$\ell_{1}$ + InfoNCE
	$\lambda=0.01$	$0.03$	$0.05$	$0.10$	$\lambda=0.01$	$0.03$	$0.05$	$0.10$
b0_y	0.93	0.82	0.67	0.41	0.95	0.90	0.82	0.58
b1_y	0.93	0.82	0.66	0.40	0.95	0.90	0.82	0.57
b2_y	0.93	0.81	0.65	0.40	0.95	0.90	0.82	0.57
b2_x	0.89	0.69	0.50	0.30	0.93	0.85	0.73	0.46
b0_x	0.89	0.68	0.50	0.30	0.94	0.86	0.74	0.47

Appendix

Table 3: Stage-1 gate values on Cube-Triple (32 dimensions). Values are continuous gates learned with

\ell_{1}

regularization and the two discrimination objectives across the displayed

\lambda

values.

Feature	$\ell_{1}$ + Hinge				$\ell_{1}$ + InfoNCE
	$\lambda=0.01$	$0.03$	$0.05$	$0.10$	$\lambda=0.01$	$0.03$	$0.05$	$0.10$
b3_y	0.93	0.77	0.58	0.34	0.95	0.89	0.79	0.52
b2_y	0.92	0.76	0.57	0.33	0.95	0.89	0.78	0.51
b0_y	0.92	0.75	0.57	0.33	0.95	0.88	0.77	0.50
b1_y	0.92	0.75	0.57	0.33	0.95	0.88	0.77	0.50
eff_y	0.89	0.80	0.69	0.48	0.89	0.82	0.74	0.53

Appendix

Table 4: Stage-1 gate values on Cube-Quadruple (37 dimensions). Values are continuous gates learned with

\ell_{1}

regularization and the two discrimination objectives across the displayed

\lambda

values.

Feature	$\ell_{1}$ + Hinge				$\ell_{1}$ + InfoNCE
	$\lambda=0.01$	$0.03$	$0.05$	$0.10$	$\lambda=0.01$	$0.03$	$0.05$	$0.10$
drawer	0.89	0.78	0.66	0.44	0.96	0.92	0.86	0.68
eff_z	0.88	0.82	0.76	0.61	0.87	0.83	0.78	0.63
cos_j3	0.87	0.83	0.77	0.60	0.87	0.83	0.78	0.64
window	0.85	0.73	0.60	0.40	0.92	0.87	0.79	0.60
sin_j2	0.81	0.63	0.48	0.27	0.80	0.66	0.52	0.29

Appendix

Table 5: Stage-1 gate values on Scene (26 dimensions). Values are continuous gates learned with

\ell_{1}

regularization and the two discrimination objectives across the displayed

\lambda

values.

$\lambda$	Spearman $\rho_{\mathrm{S}}$	Pearson $\rho_{\mathrm{P}}$	mean $\|\Delta g_{k}\|$	max $\|\Delta g_{k}\|$	Counts	Jaccard
0.01	0.959	0.966	0.040	0.08	21 / 22	0.95
0.03	0.981	0.974	0.037	0.09	20 / 20	1.00
0.05	0.977	0.982	0.029	0.08	13 / 15	0.87
0.10	0.995	0.983	0.023	0.14	5 / 5	1.00

Appendix

Table 6: Gate stability under the Scene-coordinate augmentation. The comparison uses the original 22 Cube coordinates, with InfoNCE and

\lVert\mathbf{g}\rVert_{1}

regularization. Spearman

\rho_{\mathrm{S}}

and Pearson

\rho_{\mathrm{P}}

measure rank and linear association, respectively. Counts give clean / augmented numbers of coordinates at or above the reporting threshold.

Alignment target	State	CEM	iCEM
none (LeWM)	–	67.7	70.8
SCALE-All	clean, 22-d	75.5	76.3
SCALE-All	contaminated, 48-d	74.2	75.5
SWIM ( $\lambda$ =.05)	contaminated, 48-d	75.2	76.7
SWIM ( $\lambda$ =.10)	contaminated, 48-d	74.5	76.8

Appendix

Table 7: Downstream planning under state contamination on Cube. CEM/iCEM success percentages from a single training seed. The reference is clean SCALE-All; bold marks the largest value in each planner column.

Task	LeWM	SCALE-All	SWIM-H ⋆	SWIM-N ⋆	$\Delta_{\text{SCALE-All}}$	$\Delta_{\text{LeWM}}$
Cube	67.8 $\pm$ 1.7	76.8 $\pm$ 0.5	77.1 $\pm$ 0.6 (.03)	76.8 $\pm$ 0.4 (.03)	+0.2	+9.3
Cube-Double	18.0 $\pm$ 1.7	19.5 $\pm$ 1.8	21.6 $\pm$ 1.4 (1.0)	22.6 $\pm$ 1.9 (.05)	+3.1	+4.6
Cube-Triple	35.3 $\pm$ 1.5	36.0 $\pm$ 1.5	36.9 $\pm$ 1.3 (.1)	37.1 $\pm$ 0.9 (1.0)	+1.1	+1.8
Cube-Quadruple	33.9 $\pm$ 1.5	34.7 $\pm$ 1.0	35.2 $\pm$ 1.3 (1.0)	35.7 $\pm$ 0.2 (.05)	+1.0	+1.8
Scene	24.6 $\pm$ 1.8	26.3 $\pm$ 1.2	26.7 $\pm$ 1.1 (.03)	27.1 $\pm$ 1.5 (.01)	+0.7	+2.5

Appendix

Table 8: iCEM planning success (%) over five training seeds (mean

\pm

standard deviation). Notation as in Table 1 .

Variant	CEM success (%)	iCEM success (%)	Dimensions with $g_{k}\geq 0.5$
Full Hinge ( $\kappa=1$ )	$75.5\pm 4.2$	$78.0\pm 3.3$	15/22
Full InfoNCE ( $\tau=1$ )	$77.0\pm 6.6$	$77.3\pm 3.6$	20/22
No discrimination term	$76.2\pm 5.3$	$76.8\pm 4.7$	13/22
No NLL	$75.7\pm 6.2$	$76.7\pm 3.9$	0/22

Appendix

Table 9: Loss-component ablations. Success rates are percentages. The gate-count threshold is used only for reporting; gates remain continuous. Both ablated objectives retain the gate penalty

\lambda\lVert\mathbf{g}\rVert_{1}

Penalty + objective	$\lambda$	CEM (%)	iCEM (%)
LeWM baseline	–	67.7	70.8
SCALE (9 selected dimensions)	–	73.5	76.7
$\ell_{1}$ + Hinge	0.01	76.0	77.5
	0.03	75.5	78.0
	0.05	75.2	75.5
	0.10	76.2	75.2

Appendix

Table 10: Regularization–objective comparison. All entries are planning success percentages under the planning protocol in Appendix E .

Setting	CEM success (%)	iCEM success (%)
Hinge, $\kappa=0.3$	$74.7\pm 4.3$	$77.3\pm 3.1$
Hinge, $\kappa=1$ (default)	$75.5\pm 4.2$	$78.0\pm 3.3$
Hinge, $\kappa=3$	$73.5\pm 2.7$	$77.0\pm 3.6$
Hinge, $\kappa=10$	$75.3\pm 5.1$	$76.8\pm 4.0$
InfoNCE, $\tau=0.3$	$75.0\pm 5.6$	$77.7\pm 3.7$
InfoNCE, $\tau=1$ (default)	$77.0\pm 6.6$	$77.3\pm 3.6$

Appendix

Table 11: Margin and temperature sensitivity. Values report planning success percentages as mean

\pm

standard deviation over six evaluation seeds.

Trajectories	$\lambda=0.1$			$\lambda=0.01$
	Jaccard	Spearman	Selected dims	Spearman
25	0.40	0.69	13.0	0.03
50	0.45	0.71	11.0	0.08
100	0.65	0.79	6.0	0.30
250	0.80	0.81	4.0	0.61
500	0.80	0.86	4.0	0.81

Appendix

Table 12: Gate stability versus offline-data size. Hinge with

\kappa=1

. Values are means over three trajectory-subset draws. The full-data gates from 2,000 trajectories are the reference for Jaccard and Spearman calculations.

World Models for Robotics

World models for robotics predict how scenes and robot states change under actions to support planning, control, and learning.

See all 384 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
LeWAM: End-to-End World and Action Modeling with JEPAs
Oct 8, 2026Minghao Fu, Tavis Siebert, Eryk Halicki +1World Model Learning World Model-Based Planning
UC San Diego · ETH Zurich · University of British Columbia +1
Toward Physically Grounded JEPA World Models for Goal-Conditioned Robotic Planning
Sep 3, 2026Muyuan Liu, Yue Huang, Zheng Liang +1World Models for Robotics World Model-Based Planning
GENISOM AI, Beijing, China

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 for Robotics

Explore similar work

Reperesentation Geometry Matters for Planning with JEPA World Models

LeWAM: End-to-End World and Action Modeling with JEPAs

Toward Physically Grounded JEPA World Models for Goal-Conditioned Robotic Planning