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, θ 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 ⋆
ΔSCALE-All
ΔLeWM
Cube
72.5
67.7 ± 1.0
74.6 ± 1.1
75.9 ± 1.5 (.01)
75.8 ± 1.2 (.03)
+1.3
+8.3
Cube-Double
17.5
16.4 ± 1.1
18.2 ± 0.8
19.4 ± 1.1 (1.0)
19.3 ± 0.7 (1.0)
+1.1
+3.0
Cube-Triple
34.2
30.6 ± 0.8
31.3 ± 1.2
32.6 ± 1.0 (.05)
32.4 ± 0.5 (.03)
+1.3
+2.0
Cube-Quadruple
30.2
28.3 ± 1.0
29.6 ± 0.7
30.7 ± 0.9 (.03)
30.5 ± 0.6 (1.0)
+1.0
+2.3
Scene
22.0
30.5 ± 0.5
33.4 ± 1.9
37.3 ± 3.3 (.1)
39.4 ± 2.3 (.1)
+6.0
+8.9
Table 1: CEM planning success (%) over five training seeds (mean ± standard deviation). SWIM-H ⋆ and SWIM-N ⋆ use the best λ for each task, shown in parentheses. DINO-WM † uses pretrained visual features. The Δ columns compare the better SWIM variant with SCALE-All and LeWM. Bold: best; underline: second best.
Figure 2: Learned gates on Cube (InfoNCE with ∥g∥1 regularization). (a) Gate values for the 22 candidate state dimensions as the sparsity coefficient λ 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 λ=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 λ=0.1 on the clean 22-dimensional state.
Appendix figures & tables11 assets
Supplementary material from the paper’s appendix.
Appendix
Feature
ℓ1 + Hinge
ℓ1 + InfoNCE
λ=0.01
0.03
0.05
0.10
λ=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 ℓ1 regularization and the two discrimination objectives across the displayed λ values.
Feature
ℓ1 + Hinge
ℓ1 + InfoNCE
λ=0.01
0.03
0.05
0.10
λ=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 ℓ1 regularization and the two discrimination objectives across the displayed λ values.
Feature
ℓ1 + Hinge
ℓ1 + InfoNCE
λ=0.01
0.03
0.05
0.10
λ=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 ℓ1 regularization and the two discrimination objectives across the displayed λ values.
Feature
ℓ1 + Hinge
ℓ1 + InfoNCE
λ=0.01
0.03
0.05
0.10
λ=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 ℓ1 regularization and the two discrimination objectives across the displayed λ values.
λ
Spearman ρS
Pearson ρP
mean ∣Δgk∣
max ∣Δgk∣
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 ∥g∥1 regularization. Spearman ρS and Pearson ρ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 ( λ =.05)
contaminated, 48-d
75.2
76.7
SWIM ( λ =.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 ⋆
ΔSCALE-All
ΔLeWM
Cube
67.8 ± 1.7
76.8 ± 0.5
77.1 ± 0.6 (.03)
76.8 ± 0.4 (.03)
+0.2
+9.3
Cube-Double
18.0 ± 1.7
19.5 ± 1.8
21.6 ± 1.4 (1.0)
22.6 ± 1.9 (.05)
+3.1
+4.6
Cube-Triple
35.3 ± 1.5
36.0 ± 1.5
36.9 ± 1.3 (.1)
37.1 ± 0.9 (1.0)
+1.1
+1.8
Cube-Quadruple
33.9 ± 1.5
34.7 ± 1.0
35.2 ± 1.3 (1.0)
35.7 ± 0.2 (.05)
+1.0
+1.8
Scene
24.6 ± 1.8
26.3 ± 1.2
26.7 ± 1.1 (.03)
27.1 ± 1.5 (.01)
+0.7
+2.5
Appendix
Table 8: iCEM planning success (%) over five training seeds (mean ± standard deviation). Notation as in Table 1 .
Variant
CEM success (%)
iCEM success (%)
Dimensions with gk≥0.5
Full Hinge ( κ=1 )
75.5±4.2
78.0±3.3
15/22
Full InfoNCE ( τ=1 )
77.0±6.6
77.3±3.6
20/22
No discrimination term
76.2±5.3
76.8±4.7
13/22
No NLL
75.7±6.2
76.7±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 λ∥g∥1 .
Penalty + objective
λ
CEM (%)
iCEM (%)
LeWM baseline
–
67.7
70.8
SCALE (9 selected dimensions)
–
73.5
76.7
ℓ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, κ=0.3
74.7±4.3
77.3±3.1
Hinge, κ=1 (default)
75.5±4.2
78.0±3.3
Hinge, κ=3
73.5±2.7
77.0±3.6
Hinge, κ=10
75.3±5.1
76.8±4.0
InfoNCE, τ=0.3
75.0±5.6
77.7±3.7
InfoNCE, τ=1 (default)
77.0±6.6
77.3±3.6
Appendix
Table 11: Margin and temperature sensitivity. Values report planning success percentages as mean ± standard deviation over six evaluation seeds.
Trajectories
λ=0.1
λ=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 κ=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.
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.
Joint-embedding predictive architectures (JEPAs) learn compact latent world models without pixel reconstruction. Existing JEPA pipelines typically train visual representations and action-conditioned dynamics before adding a control policy or solver. We introduce LeWAM, an end-to-end JEPA for control with: (1) a joint training objective over next latent state prediction and action flow matching, and (2) an isotropic Gaussian constraint on the latent state, enabling a single model to support direct policy execution and efficient planning by evaluating sampled actions with its learned latent dynamics. Compared with prior JEPAs in visual planning, LeWAM improves success rate from 28.6% to 89.7% on contact-rich tasks and from 10.9% to 46.2% on long-horizon tasks, and plans 32.3 times faster. On behavior cloning tasks in simulation and physical robots, LeWAM is also competitive with established methods. Interpretability analyses show that LeWAM learns control-relevant representations and action distributions effectively guiding test-time planning and control.
Minghao Fu, Tavis Siebert, Eryk Halicki +1
UC San Diego · ETH Zurich · University of British Columbia +1
Action-conditioned JEPA world models enable planning toward visually specified goals without reconstructing future pixels, yet latent prediction alone does not explicitly encourage the learned representations to retain information relevant to robotic control. We introduce an end-to-end JEPA world model that augments latent prediction with inverse dynamics (IDM) and state alignment (SA). While inverse dynamics discourages latent collapse and makes latent transitions informative of the actions that produced them, state alignment grounds consecutive representations in their associated physical configuration and motion. Across four benchmark tasks, our model attains the highest success rates on TwoRoom (100%), PushT (98%), and OGBench-Cube (87%), while performing comparably to LeWorldModel on Reacher. Our ablation further shows that adding state alignment consistently improves planning success over IDM alone across all four tasks. Although LeWorldModel, our primary baseline, attains higher average straightening on OGBench-Cube, transition-subspace analysis shows that its transition energy is concentrated in a substantially lower-dimensional subspace. Our state-aligned model exhibits a higher effective transition dimension than LeWorldModel and improves planning over IDM alone, supporting state alignment as an effective complement to inverse dynamics for robotic planning.