Keeping JEPA World Models Plannable When Little of the Frame Moves

• The vanilla LeWM succeeds on 3 of 1200 SLIM trials despite reaching 98% on PushT, with probes locating failure in an action-insensitive encoder.Neither pusher nor object positions can be decoded from the latent, and the planner’s cost landscape is flat.Source: §1 · p. 1
• LeWM+I raises SLIM success from 0.003 to 0.348 ± 0.016 and improves PushT success at twice the trained horizon by +30 pp.Controls attribute the repair to the gradient into the encoder: the encoder-side term alone recovers most of the gain, the predictor-side term alone none, and together they work best.Source: §1 · p. 2
• A language-goal head reaches 0.84 on navigation versus the visual-goal oracle’s 1.00, and staged sentences raise medium and hard pushing success to the goal-frame oracle’s level.It reaches the named zone in all 300 decoy episodes and retains 84% of its success with unseen nouns; one sentence alone reaches about a fifth of the oracle on pushing tiers.Source: §1 · p. 2
• SLIM supplies multi-object 2D pushing across four difficulty tiers, with paired visual and language goals on identical scenes.Its scenes expose settings where little of the frame responds to actions, while remaining at PushT-level compute.Source: §1 · p. 2
• The paper’s mechanistic account combines probes and a response-to-actions sweep with controls that isolate encoder shaping as the repair mechanism.The controls include encoder-versus-predictor disentanglement, a shuffled-action null, and a frozen-encoder cap.Source: §1 · p. 2

• JEPA control models predict embeddings rather than pixels and regularize against collapse; prior inverse-model work uses action prediction to retain controllable state, while noting one-step models can miss parts of it.The paper reports that its one-step IDM makes the pusher decodable and empirically also restores objects, which actions do not determine; weak pushing tiers align with known limits of one-step inverse models.Source: §2 · p. 3
• Concurrent 2026 work adds inverse-dynamics objectives to LeWM-style models, but does not study the low-action-response regime or language-goal planning described here.The cited work also differs in its treatment of encoder and predicted-latent objectives and in reported controls.Source: §2 · p. 3
• Prior language-goal methods include learned rewards, shared image-language goal spaces, language-token prediction, and latent subgoal mapping, whereas this paper trains a separate head on a frozen world model.The paper distinguishes staged instructions as subgoal latents for a latent planner, rather than language-model steps supplied to a low-level controller.Source: §2 · p. 3
• Existing benchmarks include single-object pushing and richer multi-object language tasks; SLIM adds a controlled low-action-response regime at PushT-level compute.OG-Bench has multi-object visual tasks at 64 × 64 with language labels added by LAGO; Language-Table and CALVIN pair tabletop manipulation with richer language.Source: §2 · p. 3

• The four tiers vary the geometry of the pusher, target object, and zone rather than object appearance, with open-arena and wall-chain scene variants.The easy tier asks only to move the pusher; the other tiers require pushing the target object toward a zone.Source: Figure 2 · p. 4, §3 · p. 4
• SLIM provides several small objects per scene, controlled distractor counts, tiers from solved to unsolved, tunable action response, a scripted-controller ceiling, and paired visual and language goals.The benchmark targets a regime in which standard world-model benchmarks may not test whether planning works when little moves.Source: §3 · pp. 3–4
• Scenes use a top-down 2D arena with a disc pusher, small objects, two colored zones, and randomized positions and arena symmetries.Push tiers contain 1–3 distinct-color objects; success is measured after a 20-control-step episode, with a templated language goal and rendered goal frame for each scene.Source: §3 · p. 4
• Training data comprise 10 000 episodes of 200 steps collected with a hand-coded primitive policy that never sees the goal.The policy mixes six maneuvers, including orbiting, nudging, firm pushing, glancing hits, direction changes, and inter-object pushes.Source: §3 · p. 4
• Under shared random actions, 0.37% of the SLIM frame changes per predictor step versus 0.64% on PushT.A SLIM object moves 1.1 px in a training scene and 0.54 px in a single-object scene, compared with 27 px for PushT’s T-block over 25 steps.Source: §3 · pp. 4–5

4 AN INVERSE-DYNAMICS AUXILIARY LOSS FOR JEPA WORLD MODELS

• LeWM predicts next latents from a context window of latents and actions, while CEM minimizes rollout distance to an encoded goal-frame latent.The planner uses a path cost over the rollout; terminal-only cost is reported as worse.Source: §4 · p. 5
• When movable bodies contribute little to frame changes, the encoder may drop pusher and object states, leaving the planner’s cost surface flat.The paper presents this as a hypothesis for the SLIM failure.Source: §4 · p. 5
• A shared inverse-dynamics head predicts the frameskip-packed action from latent pairs and is applied to consecutive encoder latents and current/predicted latent pairs.The same head weights are used for both applications.Source: §4 · pp. 5–6
• Encoder-IDM back-propagates through fθ to retain action-decodable content; predictor-IDM updates pψ and the head while stop-gradient leaves fθ unchanged.The head is anchored by its encoder-side application on real latent pairs.Source: §4 · pp. 5–6
• At test time, the head is discarded, so planning uses vanilla LeWM modules; the auxiliary head adds under 1% of the parameters.The objective coefficients are held fixed across both benchmarks and every seed.Source: §4 · p. 6

5.1 SETUP

• The encoder is ViT-Tiny and the predictor a six-layer transformer; SLIM results average 400 paired trials across four tiers and two planner horizons.The setup holds model, optimizer, and planner fixed across cells, varying the auxiliary-loss switch, dataset, and seed.Source: §5.1 · p. 6

5.2 PLANNING RESULTS

• LeWM+I reaches 0.348 ± 0.016 on SLIM versus 3 successes in 1200 trials for plain LeWM, and reaches 0.83 on PushT at twice the trained horizon versus 0.53 for plain.On SLIM’s hard tier, success rises from 0 to 0.16; a goal-agnostic policy scores zero there.Source: §5.2 · p. 6

5.3 ABLATIONS: WHICH GRADIENT DOES THE WORK?

• Encoder-IDM alone reaches 0.298 ± 0.005, predictor-IDM alone is indistinguishable from plain, and predictor-IDM adds 5 pp only when paired with encoder-IDM.Shuffling action targets leaves performance at plain’s level; the frozen DINOv2-S encoder with I improves from 0.048 to 0.163 but remains below half of LeWM+I.Source: §5.3 · p. 7, Table 2 · p. 7

5.4 MECHANISM ANALYSIS

• On SLIM-plain, spread/motion is 3.8 × 10^-4, rollout error is 12.4× true motion, pusher-position error is 163 px, and CEM cost variation is 0.001; LeWM+I restores all four probes.LeWM+I has spread/motion 0.56–0.61, rollout-error ratio 1.6–1.7, pusher error 8.5 px, and cost variation 0.14; PushT-plain lacks these failure modes.Source: Figure 4 · p. 7, §5.4 · p. 8
• Across 28 configurations, every model below spread/motion ≈0.05 fails to plan, but success above that threshold ranges from 0.01 to 0.83.The action-sensitivity probe is an empirical gate on plannability, not a ranking of planners.Source: §5.4 · p. 8
• At fixed data volume and training budget, making 1.10% or 1.65% of the frame respond lifts plain-model success to 0.100 or 0.530; doubling object radius alone leaves it collapsed.The default changing fraction is 0.35%; PushT is at 0.57%. The sweep indicates response fraction, rather than object size alone, tracks when vanilla planning works.Source: §5.4 · p. 8

6.1 LANGUAGE-GOAL HEAD

• A 1.1M-parameter head maps the current latent and frozen sentence embedding to a goal latent, leaving the encoder, predictor, and planner unchanged.It is trained on 40,000 simulator scenes to predict the next stage’s encoded latent from a frame and that stage’s sentence.Source: §6.1 · p. 8, Figure 5 · p. 9, §6.2 · p. 9

6.2 RESULTS

• The head reaches 0.84 on navigation against the visual-goal oracle’s 1.00 and ends closer to the named zone than a swapped decoy in all 300 episodes.It retains 84% of its success on held-out templates with unseen object and zone nouns, mostly on navigation.Source: §6.2 · pp. 8–9
• A single final-goal sentence reaches only about a fifth of the oracle on pushing tiers and stays below the goal-agnostic policy on medium.The passage attributes the difficulty to specifying from the start a push ending the planner cannot reach in one step.Source: §6.2 · p. 9

6.3 STAGED INSTRUCTIONS

• Stage sentences raise medium- and hard-tier success from 0.04 to 0.25, reaching the goal-frame oracle on medium and exceeding it on hard on average.The gain lies on pushing tiers; navigation is unchanged.Source: §6.3 · p. 10, Figure 6 · p. 10
• Stage-frame subgoals lift medium success from 0.27 to 0.62 and hard success from 0.16 to 0.55, showing that nearby subgoals are easier for the planner to reach.The language head recovers 40–45% of this stage-frame oracle on pushing; its outputs are nearest to the intended stage frame in 94–100% of pushing queries for two seeds and 76–100% for the weakest seed.Source: §6.3 · p. 11
• Switching stages when latent progress stalls reaches 0.335 ± 0.037, not significantly different from a fixed schedule; a simulator judge that waits for stage completion reaches 0.215.The progress-based rule reads only the world model’s latent.Source: §6.3 · p. 11
• Staged pushing largely survives paraphrases using the training vocabulary, but its gain largely disappears with held-out object and zone nouns.With held-out paraphrases, medium and hard success are 0.09 and 0.06, while navigation keeps most of its success.Source: §6.3 · p. 11

7 CONCLUSION

• The inverse-dynamics loss restores plannability, and a small language-goal head plans from sentences on the repaired latent, including staged pushes with latent-based stage switching.The controls attribute the repair to the gradient into the encoder.Source: §7 · p. 11
• Next steps include improving staged pushing with unseen wording and transferring the repair to larger encoders, whose probes recover but whose planning has not transferred.The paper also identifies closing the gap between stage sentences and stage frames as a next step.Source: §7 · p. 11

A ADDITIONAL EXPERIMENTAL DETAILS

• The appendix reports final checkpoints after 30 epochs, with a random 90/10 clip split whose validation loss is a training diagnostic, not a held-out measure.Clips from the same episode can appear in both splits.Source: A ADDITIONAL EXPERIMENTAL DETAILS · p. 16
• LeWM+I reaches a SLIM trial mean of 0.348 ± 0.016 and improves PushT success at twice the trained horizon by +30 pp.Across three seeds, plain LeWM records 0, 2, and 1 successes out of 400 trials; the encoder-IDM-only variant records means of 0.290, 0.308, and 0.298.Source: A ADDITIONAL EXPERIMENTAL DETAILS · p. 16
• Inverse dynamics restores the probes, while patch decoding adds only a small refinement and residual prediction is redundant.Without inverse dynamics, all tested auxiliary combinations retain plain-level probes.Source: A ADDITIONAL EXPERIMENTAL DETAILS · p. 16

A.1 PLANNER CADENCE CONTROL

• SLIM loses performance under PushT’s full-plan cadence, while applying SLIM’s replan-every-block cadence to PushT lowers success for both models.Each benchmark is most usable under its native cadence, and the LeWM+I-versus-plain gap keeps its sign under either cadence.Source: §A.1 · p. 17
• At three times the trained horizon, PushT success is 0.04 for both models; path cost scores 0.335 versus 0.083 for terminal cost.The very-hard tier also has a lower cost-landscape coefficient of variation than the easy, medium, and hard tiers.Source: §A.1 · p. 17

B ROBUSTNESS TO OBSTACLE WALLS

• Across four converged walled LeWM+I seeds, success is 0.90±0.02 on easy, 0.27±0.03 on medium, 0.04±0.02 on hard, and 0 on very hard.Unconditionally, the five-seed mean is 0.245 ± 0.131 with a 1-in-5 failure; wall chains obstruct steered pushes on hard.Source: B ROBUSTNESS TO OBSTACLE WALLS · pp. 17–18

C ENCODER SCALE: VIT-SMALL

• At ViT-Small, LeWM+I reaches only 0.033 trial mean versus 0.008 for plain, despite recovering representation probes.This matched-epoch comparison supports the mechanism at the representation level, but the authors say the absolute success ceiling is ViT-Tiny-specific at their budget.Source: C ENCODER SCALE: VIT-SMALL · p. 18

D OGBENCH CUBE: AN EXPECTED NULL

• On OGBench cube-single, plain and LeWM+I are indistinguishable at 4/50 versus 3/50; cube-double scores 0/50 for both.The cube-double cell is described as uninformative because neither model reaches the capability floor.Source: D OGBENCH CUBE: AN EXPECTED NULL · p. 18

E A SCRIPTED REFERENCE CONTROLLER

• A scripted controller with simulator state solves every tier within 20 steps except very hard, which reaches 0.28; it solves every tier within 60 steps.The authors present this as evidence that the scenes are solvable by a state-informed policy, not as an optimum.Source: E A SCRIPTED REFERENCE CONTROLLER · p. 18

F DISPLACEMENT STATISTICS OF THE TRAINING DISTRIBUTIONS

• SLIM objects move under 5 px in 48% of predictor-step transitions and under 25 px in 99%, while the median pusher displacement is 22 px.The T-block is static in 54% of PushT transitions, but moves further when it moves; its larger pixel footprint and greater response to sampled actions distinguish the benchmarks.Source: F DISPLACEMENT STATISTICS OF THE… · p. 19

G RESPONSE-TO-ACTIONS SWEEP

• Increasing object size alone barely changes the frame-response fraction, whereas lighter objects and a larger, faster pusher raise it to 1.10% and 1.65% in the response arms.The sweep holds data volume, training, and evaluation fixed; both objects and the enlarged pusher contribute to the changed fraction.Source: G RESPONSE-TO-ACTIONS SWEEP · p. 19, Figure 10 · p. 20
• At the low-response arm, plain LeWM is action-sensitive at 0.88 yet plans poorly, with trial mean 0.100.The authors therefore treat action sensitivity as a gate on plannability, not a ranking of planners.Source: G RESPONSE-TO-ACTIONS SWEEP · p. 20
• Random-policy trial means rise to 0.045 for the high-response arm and 0.040 for the low-response arm, while zero actions still score 0.00.The authors report these as below the repaired model’s performance.Source: G RESPONSE-TO-ACTIONS SWEEP · p. 20

H LOSS-WEIGHT SENSITIVITY AND THE SIGREG WEIGHT

• Trial mean success rises from 0.000 to 0.368 as α increases from 0 to 1; the hard tier rises from 0.01 to 0.17 between α = 0.1 and 1.Easy saturates at α = 0.3, while very hard remains at 0; changing the SIGReg weight leaves plain-model success at the floor.Source: H LOSS-WEIGHT SENSITIVITY AND THE… · pp. 20–21
• Action sensitivity separates usability from quality: α = 0.1 reaches spread/motion 0.55, while α = 0.3 scores 0.293 trial mean despite the highest spread/motion, 0.80.The probe is an empirical gate below which configurations fail to plan, not a predictor of planning quality above it.Source: H LOSS-WEIGHT SENSITIVITY AND THE… · p. 20

I DOES THE LATENT COST RANK ACTIONS LIKE THE SIMULATOR?

• LeWM+I ranks easy-tier simulator costs with Spearman ρ = 0.84 ± 0.01, versus plain’s −0.02 ± 0.05; medium-tier ranking remains signal-starved.Encoder-IDM only and the full model score +0.88 ± 0.01 and +0.84 ± 0.01, while predictor-IDM only is at +0.03 ± 0.03; cost ranking and planner success disagree on their ordering.Source: I DOES THE LATENT COST RANK ACTIONS… · p. 21

J CROSS-DISTRIBUTION GENERALISATION

• Without distractors during training, plain still collapses at 1/400, while LeWM+I reaches 0.190 mean success on default multi-object tiers.The repaired model’s spread/motion is 0.24, about 400× plain’s 6 × 10−4 on the same data, but its success is about half the in-distribution 0.368.Source: J CROSS-DISTRIBUTION GENERALISATION · p. 21

K GROUNDING CONTROLS FOR THE LANGUAGE HEAD

• The encoded-goal-frame head matches the oracle at 0.363 versus 0.368, while the cross-attention language head succeeds in 252/300 swapped-zone episodes and never ends at the decoy.A goal-agnostic policy scores zero on hard; the global-latent language MLP reaches 0.195 and is less reliable under the zone swap.Source: K GROUNDING CONTROLS FOR THE LANGUAGE… · pp. 21–22

L LANGUAGE HEAD: ADDITIONAL RESULTS

• The cross-attention head has 1.1M parameters and predicts the next-stage encoded latent from the current latent, sentence embedding, and encoder patch tokens.It uses two pre-LN transformer layers and is trained for 90 epochs; the global-latent MLP has about 5M parameters and is trained for 30 epochs.Source: L LANGUAGE HEAD: ADDITIONAL RESULTS · p. 22
• Training pairs each current-stage latent and next-stage sentence with the next-stage latent, including approach, midway, and goal transitions.At test time, the final-goal sentence is provided at episode start and again after 10 planner steps, each query using the current frame.Source: L LANGUAGE HEAD: ADDITIONAL RESULTS · p. 22
• With held-out templates whose object and zone nouns never occur in training, success is 0.78/0.01/0.03 on easy/medium/hard, with very hard at 0.Canonical success is 0.88/0.05/0.05, and the 107-template train pool yields 0.84/0.04/0.02.Source: L LANGUAGE HEAD: ADDITIONAL RESULTS · p. 22
• On easy-tier swapped-zone episodes, the head scores 0.92/0.86, 0.84/0.68, and 0.92/0.82 across three seeds at Ha=5/10, with zero decoy hits.Every episode ends closer to the named zone; without the swap, the same heads score 0.88, 0.78, and 0.87.Source: L LANGUAGE HEAD: ADDITIONAL RESULTS · p. 22

M STAGED INSTRUCTIONS: PROTOCOL AND ADDITIONAL RESULTS

• Pushing is queried with approach, midway, then goal sentences at planner steps 0, 7, and 14; navigation receives an approach sentence then goal sentences.Each query uses the current frame, and the planner keeps the latest head output as its target.Source: M STAGED INSTRUCTIONS: PROTOCOL AND… · p. 22
• The latent switch advances after the target-distance criterion or stalled progress; it enters midway after about 5 steps and goal after about 11 on medium and hard, with at least 98% reaching goal stage.The selected rule is ρ = 0.1, P = 3; navigation reaches goal stage after about 3 steps.Source: M STAGED INSTRUCTIONS: PROTOCOL AND… · pp. 22–23
• With simulator-judged switching, 91% of medium and 80% of hard episodes end at midway without receiving the goal sentence; 12 of 13 that reach goal stage succeed.The reference advances using simulator distances to approach, midway, and target poses.Source: M STAGED INSTRUCTIONS: PROTOCOL AND… · p. 23
• Stage mapping measures whether the output for the next-stage sentence is nearest to the intended encoded frame; goal sentences from the start frame remain farther from goal than from midway.The remaining goal-frame distance is 0.10–0.12 from the start frame versus 0.03–0.05 from the midway frame; Table 7 reports the fraction of scenes mapping nearest to the intended stage.Source: M STAGED INSTRUCTIONS: PROTOCOL AND… · p. 23, Table 7 · p. 23
• Across head seeds, staged instructions are compared with a single goal sentence and the goal-frame oracle; on pushing tiers, no seed falls below the oracle.Stage sentences also avoid decoy endpoints across all three seeds, with every episode ending closer to the named zone.Source: M STAGED INSTRUCTIONS: PROTOCOL AND… · p. 23, Table 8 · p. 23