Keeping JEPA World Models Plannable When Little of the Frame Moves

Florian Strohm; Patrick Wagner; Jannik Schwab; Marco Huber

cs.AIOct 2, 2026

Keeping JEPA World Models Plannable When Little of the Frame Moves

Authors: Florian Strohm, Patrick Wagner, Jannik Schwab, Marco Huber

Organizations: Fraunhofer IPA, Stuttgart, Germany · University of Stuttgart, Germany

Abstract

Specifying a goal in language rather than as a goal frame is a natural interface for planning with a latent world model, but testing it needs scenes in which language must discriminate between several objects. We build SLIM, a pushing benchmark with several small objects and paired visual and language goals on identical scenes. On SLIM a LeWM world model that solves PushT succeeds on under 1% of trials, although a scripted controller with simulator state solves every tier. Probes locate the failure in the encoder: its latent is nearly action-insensitive, neither pusher nor object positions can be decoded from it, and rollouts are no better than copying the current latent forward. One inverse-dynamics auxiliary loss, applied to encoder latents and to predicted latents through a shared head discarded at test time, restores every probe and raises success from 0.003 to 0.35 (0.16 on the hard pushing tier, where a goal-agnostic policy scores zero), and improves PushT at twice the trained horizon. Controls attribute the repair to the gradient into the encoder, and a response sweep shows that the vanilla model plans once enough of the frame responds to actions. A cheap action-sensitivity probe, computable without environment access, acts as an empirical necessary condition: all configurations below its threshold failed to plan. On the repaired latent, a small language-goal head plans from sentences without retraining the world model: it reaches 0.84 on navigation (visual-goal oracle 1.00), follows the named zone when it is swapped with a decoy, and degrades gracefully to unseen nouns. A single goal sentence rarely completes a push, but given the push as a sequence of stage sentences the head raises success on the medium and hard pushing tiers from 0.04 to 0.25, on par with the goal-frame oracle, also when the switch between stages is read from the latent alone.

Figures & tables

Figure 1: Why JEPA world models collapse on SLIM. Under the same random actions, $0.64\%$ of the frame changes per predictor step on PushT (left), where the large T-block moves whenever it is touched, and $0.37\%$ on SLIM (right; the same scene as in Fig. 2 ), where the pusher rarely hits one of several objects. When so little of the frame responds to actions, the prediction objective can be minimised by an action-insensitive encoder ( Section 5.4 ).

Figure 2: SLIM scenes for the four tiers (columns) in the open arena (top) and the identical configuration with procedurally generated wall chains added (bottom). The templated language goal of each scene is quoted under the top row; the easy tier asks only to move the pusher. Tiers differ in the geometry of pusher, target object and zone, not in object appearance.

Panel of Figure 3 — Figure 3: LeWM+I. Left (training): the LeWM backbone is unchanged — encoder $f_{\theta}$ , action encoder $g_{\phi}$ , predictor $p_{\psi}$ , latent prediction loss $\mathcal{L}_{\text{pred}}$ . We add one shared inverse-dynamics head $h_{\xi}$ (orange), applied twice per batch with the same weights: to the encoder pair $(z_{t},z_{t+1})$ , whose gradient flows back into $f_{\theta}$ and forces consecutive latents to retain the action, and to the predictor pair $(\mathrm{sg}(z_{t}),\hat{z}_{t+1})$ , which forces the predicted latent to carry an action signature the same head can read. Right (test): the head is discarded; CEM rolls out encoder and predictor under sampled actions and minimises the latent goal cost averaged over the rollout, with $z_{\text{goal}}$ from an encoded goal frame or, in Section 6 , from the language head.

	SLIM (three training seeds, both horizons)					PushT
model	easy	medium	hard	very hard	trial mean	$H_{a}{=}5$	$H_{a}{=}10$
LeWM (plain)	0.007	0.000	0.003	0.000	$0.003\pm 0.001$	$0.98\pm 0.01$	$0.53\pm 0.03$
LeWM+I (ours)	0.990	0.243	0.157	0.000	$\mathbf{0.348\pm 0.016}$	$0.97\pm 0.01$	$\mathbf{0.83\pm 0.01}$
random or zero-action policy	0.00	0.00	0.00	0.00	0.00	—	—
goal-agnostic, one direction (tuned in hindsight)	0.12	0.20	0.00	0.00	0.08	—	—
scripted controller, privileged state	1.00	1.00	1.00	0.28	0.82	—	—

Table 1: Headline results. Per-tier values pool three training seeds and both horizons (300 trials each); trial-mean and PushT columns are mean

\pm

SEM over seeds (McNemar plain vs. LeWM+I, canonical seed:

p<10^{-43}

; per-seed values in Appendix A ). Reference rows use the same scenes, budgets and thresholds: random or zero actions, the best single direction and duration chosen per tier in hindsight (it never reads the goal), and a scripted controller with simulator state ( Appendix E ).

cell	$\alpha_{\text{inv}}$	$\alpha_{\text{pid}}$	encoder	trial mean	vs. plain (McNemar)
plain	0	0	ViT-Tiny, trained	0.000	—
predictor-IDM only	0	1	ViT-Tiny, trained	0.010	$+1.0$ pp ( $p=0.13$ )
encoder-IDM only	1	0	ViT-Tiny, trained	0.290	$+29.0$ pp ( $p<10^{-34}$ )
shuffled-action null	1	1	ViT-Tiny, trained	0.003	$+0.3$ pp ( $p=1$ )
frozen DINOv2-S, no I	0	0	DINOv2-S, frozen	0.048	$+4.8$ pp
frozen DINOv2-S + I	1	1	DINOv2-S, frozen	0.163	$+16.3$ pp †

Table 2: Disentangle. Canonical training seed (the highest of three for LeWM+I), on which McNemar pairs trials; three-seed means are

0.003

(plain),

0.348

(LeWM+I),

0.298

(encoder-IDM) and

0.002

(shuffled null), per-seed values in Appendix A . † Against frozen DINOv2-S without I:

+11.5

pp,

p<10^{-9}

Figure 4: Mechanism probes on SLIM (and PushT for reference), LeWM vs. LeWM+I. (a) Plain’s latent is almost action-insensitive, two orders of magnitude below the usability threshold and three below LeWM+I; PushT-plain is action-sensitive without I. (b) Pusher and object positions cannot be decoded from plain’s latent (dashed line: constant mean-position predictor) and are recovered to within one patch by LeWM+I; errors in the 512-px simulator frame. (c) Plain’s CEM cost surface is flat; LeWM+I restores a usable descent direction. (d) Planning success against the fraction of the frame that changes per step under random actions, over four environment variants ( Appendix G ): plain stays near zero until enough of the frame responds.

Panel of Figure 5 — Figure 5: Language-goal head. Top (inference): a small trained head maps the frozen encoder’s $z_{t}$ and a frozen sentence embedding of the goal description to a goal latent $z_{\text{goal}}$ , which replaces the encoded goal frame in the CEM cost; encoder, predictor and planner are unchanged. Bottom (training): from each frame the head predicts the encoded latent of the next stage of a push, given that stage’s sentence; at test time it is always asked for the final goal.

goal channel	easy	medium	hard	trial mean	$\Delta$ vs. oracle (McNemar)
goal-agnostic, one direction (hindsight)	0.12	0.20	0.00	0.08	—
visual-goal oracle	1.00	0.30	0.17	0.368	—
goal-frame-conditioned head (MLP)	0.99	0.29	0.17	0.363	$-0.5$ pp ( $p=0.80$ )
language head over the global latent (MLP)	0.53	0.16	0.09	0.195	$-17.3$ pp ( $p<10^{-11}$ )
language head, patch attention, 3 seeds	0.84	0.06	0.04	$\mathbf{0.235\pm 0.026}$	$-13.3$ pp ( $p<10^{-9}$ each seed)
zero sentence embedding	0.00	0.00	0.01	0.003	vs. head: $p<10^{-27}$

Table 3: Language goals on the LeWM+I world model of the canonical seed, whose visual-goal planner is the oracle (per-tier values average both horizons;

\pm

: SD over head seeds). Indented rows ablate head seed 42 (

0.245

), against which their p-values are computed; per-seed values in Appendix A .

subgoals	easy	medium	hard	very hard	trial mean	paired test (McNemar)
goal frame (visual-goal oracle)	0.98	0.27	0.16	0.00	0.353	—
stage frames (oracle)	0.99	0.62	0.55	0.01	0.543	vs. goal frame: $p<10^{-18}$
goal sentence, 3 head seeds	0.84	0.04	0.04	0.00	$0.229\pm 0.021$	—
stage sentences , every 7 steps	0.82	0.25	0.25	0.01	$\mathbf{0.334\pm 0.059}$	vs. goal sentence: $p<0.002$ each seed
stage sentences , latent switch	0.82	0.29	0.23	0.00	$\mathbf{0.335\pm 0.037}$	vs. goal sentence: $p<10^{-4}$ each seed
stage sentences, paraphrases, train pool	0.78	0.24	0.18	0.01	$0.305\pm 0.055$	vs. goal sentence: $p<10^{-4}$ , $0.08$ , $<10^{-4}$

Table 4: Staged instructions on the LeWM+I world model of Table 3 ; the language rows use the cross-attention head (per-tier values average both horizons and the three head seeds;

\pm

: SD over head seeds). Indented rows use head seed 42 (stage sentences every 7 steps:

0.340

). All rows come from one evaluation campaign on one machine; per-seed values in Appendix M .

Figure 6: Goal vs. stage subgoals, as frames and as sentences. Stage frames lift the planner on the pushing tiers far above the goal frame; stage sentences lift the language head from

0.04

to the level of the goal-frame oracle there. Bars: means; dots: the three head seeds.

Appendix figures & tables8 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 7: Action sensitivity vs. CEM success for 28 trained configurations (circles SLIM, squares PushT, triangles walled SLIM, diamonds ViT-Small, downward triangles the response-to-actions arms, plus signs the loss-weight and SIGReg cells; filled with I, hollow without; the cross is LeWM+I trained on the single-object variant and evaluated cross-distribution). Below

0.05

every planner fails; above it success ranges from

0.01

0.83

: the frozen-DINOv2 cells, the non-converged walled seed and ViT-Small clear the threshold yet plan poorly.

probe	SLIM plain	SLIM LeWM+I	PushT plain	PushT LeWM+I
action sensitivity, spread/motion	$3.8\times 10^{-4}$	$0.56$ / $0.61$	$0.67$	$0.57$
rollout error / true motion ( $h{=}10$ )	$12.4$	$1.69$ / $1.62$	$1.02^{\dagger}$	$0.95^{\dagger}$
pusher position MLP error (px)	$163\pm 2$	$8.5\pm 0.1$	—	—
object position MLP error (px)	$168\pm 0.4$	$15.1\pm 0.2$	—	—
cost-landscape cv	$0.001$	$0.14$	—	—

Appendix

Table 5: Mechanism probes. † For PushT only the predicted-over-true latent displacement was measured, not the error ratio, so these two entries are not comparable with the SLIM column. Two values for SLIM LeWM+I are two training seeds (within

10\%

on spread/motion,

5\%

on the rollout-error ratio). Position errors (mean

\pm

SEM over three probe seeds, 80/10/10 split) are in the 512-px simulator frame, where a constant mean-position predictor scores

\approx 165

px.

Figure 8: Validation encoder-IDM loss per epoch for the five walled LeWM+I seeds; seed 9999’s spike (peak

3.9

at epoch 17) exceeds the axis range.

Figure 9: Per-predictor-step displacement of pusher/agent and of the most-moved object in the SLIM and PushT training data.

Figure 10: Response-to-actions sweep. Left: CEM trial mean; right: action-sensitivity probe (dotted line: usability threshold). Arms ordered by the fraction of the frame that changes per predictor step under native random actions (0.35, 0.36, 1.10, 1.65%).

arm	radius	pusher	mass / damping	scale	frame changed	model	easy	medium	hard	very hard	trial mean	spread/motion
r15	15	14	8 / 0.85	40	0.35%	plain	0.01	0.01	0.00	0.00	0.005	$8\times 10^{-4}$
						LeWM+I	0.96	0.20	0.13	0.00	0.323	0.60
r30	30	14	8 / 0.85	40	0.36%	plain	0.00	0.00	0.00	0.00	0.000	$3\times 10^{-4}$
						LeWM+I	0.99	0.22	0.12	0.01	0.335	0.69
lo	15	28	0.5 / 0.95	80	1.10%	plain	0.09	0.19	0.12	0.00	0.100	0.88
						LeWM+I	0.96	0.82	0.44	0.01	0.558	0.81

Appendix

Table 6: Response-to-actions sweep. One training seed per cell; per-tier values average the two planner horizons; “frame changed” is the model-free fraction of the frame that changes per predictor step under native random actions; the r15 rows are the canonical cell of Table 1 retrained.

	easy		medium			hard			very hard
head seed	appr.	goal	appr.	mid.	goal	appr.	mid.	goal	appr.	mid.	goal
42	0.74	0.98	0.96	1.00	0.94	0.94	1.00	0.94	1.00	1.00	0.98
7	0.66	1.00	0.90	0.98	0.92	0.76	1.00	0.92	1.00	1.00	0.98
11	0.64	0.98	1.00	1.00	0.96	1.00	1.00	0.94	1.00	1.00	0.98

Appendix

Table 7: Stage mapping : fraction of the 50 evaluation scenes per tier in which the head’s output for the next-stage sentence, queried from the rendered frame of the previous stage, lies nearest to the encoded frame of the intended stage.

trial mean (easy / medium / hard / very hard)	head seed 42	head seed 7	head seed 11
goal sentence	$0.245$ (0.88/0.04/0.06/0)	$0.205$ (0.76/0.04/0.02/0)	$0.238$ (0.87/0.04/0.04/0)
stage sentences, every 7 steps	$0.340$ (0.83/0.27/0.25/0.01)	$0.273$ (0.72/0.17/0.19/0.01)	$0.390$ (0.91/0.32/0.31/0.02)
stage sentences, latent switch	$0.363$ (0.90/0.33/0.21/0.01)	$0.293$ (0.76/0.21/0.20/0)	$0.350$ (0.80/0.32/0.28/0)
paraphrases, train pool	$0.353$ (0.87/0.30/0.22/0.02)	$0.245$ (0.68/0.20/0.10/0)	$0.318$ (0.80/0.23/0.22/0.02)
paraphrases, unseen nouns	$0.238$ (0.74/0.15/0.06/0)	$0.215$ (0.69/0.09/0.08/0)	$0.218$ (0.80/0.04/0.03/0)

Appendix

Table 8: Staged instructions per head seed (400 trials each).

Language-Conditioned Robot Planning

See all 16 papers

Supervise What Decides Success: Criterion-Aligned Auxiliary Losses for Latent World-Model Planning
Oct 1, 2026Takumi Hara, Kanata SuzukiWorld Model Learning World Models for Robotics
Kyoto University Kyoto, Japan · Fujitsu Limited Kanagawa, Japan
Grounded World Model: Latent Planning with Language Goals
Apr 13, 2026Quanyi Li, Lan Feng, Haonan Zhang +4Language-Conditioned Robot Planning World Models for Robotics
EPFL · NUS
The Planning Limits of Latent World Models
Sep 30, 2026Ali Alrasheed, Basim Azam, Naveed AkhtarWorld Model-Based Planning Action-Conditioned World Models
The University of Melbourne, Australia

Oct 1, 2026J/K move · Enter open · S save

Supervise What Decides Success: Criterion-Aligned Auxiliary Losses for Latent World-Model Planning

Takumi Hara, Kanata Suzuki

Kyoto University Kyoto, Japan · Fujitsu Limited Kanagawa, Japan

World Model Learning World Models for Robotics World Model-Based Planning Model-Based Planning Latent World Models

Open paper AskPDF arXiv

Latent world models plan by scoring candidate action sequences with distances in latent space. However, task success is judged by physical quantities, which we call the success-criterion quantities. In all four latent world models we examine, the end-effector position is encoded in the latent state with an error larger than the success criterion allows. Such a latent state cannot separate successful candidates from failing ones. We propose an auxiliary loss that uses success-criterion quantities as training targets, whereas existing latent world models take them only as inputs. During training, a linear head on the encoder and predictor outputs regresses the success-criterion quantities, and the regression error is added to the training loss. The head is discarded after training, so the model, its cost, and its inputs at test time are unchanged. This loss alone improves the success rate on PushT and cube by 3.5% and 3.4% (absolute), respectively, and both improvements are statistically significant. A success criterion thus specifies what a world model must retain in its latent state, and we show that it can serve directly as a training target.

Abstract

Language-Conditioned Robot Planning

Explore similar work

Supervise What Decides Success: Criterion-Aligned Auxiliary Losses for Latent World-Model Planning

Grounded World Model: Latent Planning with Language Goals

The Planning Limits of Latent World Models