SWIM: Learning What Matters from Behaviors for JEPA World Models

1 Introduction

• JEPA latent prediction does not explicitly require representing the physical state that determines task progress, despite evidence that latents encode physical quantities.Forward prediction alone does not ensure that task-relevant state is reliably identifiable or well separated for planning.Source: §1 · p. 3
• Equal weighting can overemphasize redundant coordinates and introduce irrelevant variation, while hand-selecting state variables requires task-specific expertise.This motivates learning relative feature contributions instead of treating every coordinate equally.Source: §1 · p. 3
• Recorded actions provide a behavior-grounded surrogate for state relevance because state changes that help infer actions reflect agent behavior.Redundant or unaffected coordinates add little information about the action.Source: §1 · p. 3
• SWIM learns continuous state weights by predicting recorded actions from current states and future displacements, then uses the weighted metric for state-supervised JEPA training.The framework is demonstrated by calibrating latent geometry through SCALE and requires neither reward nor success labels.Source: §1 · p. 3
• Across five OGBench manipulation tasks, SWIM consistently improves over uniform state weighting with CEM and iCEM planners.The evaluation uses five training seeds.Source: §1 · pp. 3–4

Latent World Models and Planning

• JEPA world models predict in latent space without reconstructing observations, but their prediction objectives do not specify which task-relevant information representations should capture.The passage contrasts pretrained visual features in DINO-WM with jointly learned visual representations and latent dynamics in PLDM and LeWM.Source: §2 · p. 4

State Supervision and Planning-Oriented Representations

• TC-WM, PhyLatent, PSG-JEPA, and SCALE use fixed state-coordinate choices, whereas SWIM learns each coordinate’s contribution and demonstrates it with SCALE.Other planning-oriented approaches use value guidance, temporal distance, reachability, or action-sequence distance ranking.Source: §2 · p. 4

Inverse Dynamics and Behavioral Relevance

• Prior inverse-dynamics and action-based methods shape representations directly, but do not learn per-coordinate weights over recorded state variables for state supervision.SWIM uses gated inverse dynamics to learn those weights, which then define world-model state supervision.Source: §2 · p. 4

3 Preliminaries

• The preliminaries establish vector and matrix notation and introduce LeWM and SCALE as the planning and state-supervision framework used to demonstrate SWIM.SCALE correlates latent distances with state distances.Source: §3 · p. 5

3.1 Latent-Space Planning with LeWM

• LeWM jointly learns an image encoder and action-conditioned latent dynamics, then rolls predicted embeddings forward for candidate action sequences at test time.CEM and iCEM minimize a planning cost based on the predicted terminal embedding and goal embedding; LeWM trains from observed transitions and uses SIGReg to maintain global spread and avoid collapse.Source: §3.1 · p. 5

3.2 SCALE

• SCALE correlates latent distances with standardized task-state distances for sampled image-frame pairs, encouraging larger state differences to map to larger latent separations.Its training objective combines latent prediction, SIGReg, and a Pearson-correlation loss.Source: §3.2 · p. 6
• Uniform state weighting can make redundant or irrelevant coordinates influence SCALE’s target geometry, motivating a learned weighted state metric.The section calls SCALE applied to all supplied coordinates with uniform weights SCALE-All, the main experimental baseline.Source: §3.2 · p. 6

4 SWIM

• SWIM separates behavior-informed weight learning from state-distance calibration, producing a state metric that does not depend on the downstream objective.Stage 1 learns weights from offline trajectories; Stage 2 demonstrates the metric with SCALE.Source: §4 · p. 6

4.1 Learning State-Feature Gates via Inverse Dynamics

• A gated inverse-dynamics predictor uses current states and future displacements to predict recorded actions, with discrimination against mismatched displacements and an L1 gate penalty.The objective combines action negative log-likelihood, hinge or InfoNCE discrimination, and gate regularization.Source: §4.1 · pp. 6–8

4.2 Transferring Learned Gates to Latent Geometry

• Frozen gates rescale standardized state coordinates and replace SCALE’s uniform state-distance targets in its latent-distance correlation loss.The inverse-dynamics model is discarded after training, and deployment retains image-based planning without simulator states or extra computation.Source: §4.2 · p. 8

4.3 Why Learned Weights Help Planning

• The planning bound separates alignment error, rollout error, and mismatch from reference task weights; the mismatch term remains under uniform weights even with perfect alignment.Learned weights moving toward the reference metric reduce that term.Source: §4.3 · pp. 8–9

5.1 Setup

• Five OGBench tasks with redundant state encodings are compared under shared data and planning protocols, with five training seeds and 600 episodes per task, configuration, and planner per seed.The tasks are Cube, Cube-Double, Cube-Triple, Cube-Quadruple, and Scene; evaluation uses CEM and iCEM.Source: §5.1 · p. 9

5.2 Does Behavioral Weighting Improve Planning?

• SWIM beats LeWM and SCALE-All on all five tasks, by 1.0–6.0 and 2.0–8.9 percentage points, respectively.Across ten configurations, the mean gain over SCALE-All is 0.7 points (p = 0.006); DINO-WM remains stronger on Cube-Triple, while the two SWIM objectives perform comparably overall.Source: §5.2 · p. 9, Table 1 · p. 10

5.3 What Happens When Irrelevant Variables Are Supplied?

• Appending 26 Scene variables to Cube preserved the original gate ranking (Spearman correlation 0.995; distractors ≤0.31), while weighted alignment raised CEM success from 74.2 to 75.2.Under iCEM, weighted alignment reaches 76.7, above contaminated (75.5) and clean (76.3) uniform alignment.Source: Figure 2 · p. 10, §5.3 · p. 11

• SWIM improves over uniform state weighting by 0.9–1.0 percentage points (p < 0.05) across five OGBench manipulation tasks and two planners.The structured-distractor experiment also shows downweighting of unrelated variables appended to the state.Source: §6 · p. 11
• The learned metric depends on the state variables supplied during training and the behaviors covered by offline data, and the paper demonstrates it through SCALE.Source: §6 · p. 11
• Because the learned metric does not depend on the downstream objective, applying it to other state-supervised objectives is a proposed next step.The paper names subspace alignment and physical-state grounding as examples.Source: §6 · p. 11
• The authors expect larger gains when supplied states contain more redundant or noisy coordinates and identify state-dependent weighting as future work.Source: §6 · p. 11

• Stage-1 examples pair a state, a later state from the same offline trajectory, and the recorded action chunk; behavioral horizon h is distinct from planning horizon H.The validity filter uses the maximum behavioral offset, 45 raw frames, regardless of the sampled horizon.Source: A Training Details and the Behavioral… · p. 14
• Stage 1 standardizes states and actions separately, clips both to [−10, 10], and computes displacements as direct differences of standardized states.State standard deviations have a 10−3 floor; Stage 2 does not clip standardized states.Source: A Training Details and the Behavioral… · p. 14
• The Gaussian action-chunk predictor is a two-hidden-layer, 1024-unit MLP whose outputs are the chunk’s mean and diagonal log-standard-deviation.Gate logits start at 2 (gk ≈ 0.88), and learned coordinates are not hard-thresholded.Source: A Training Details and the Behavioral… · p. 14
• In-batch discrimination contrasts each observed displacement with mismatched later states, but does not establish that those alternatives are unreachable or behaviorally incompatible.InfoNCE uses 255 alternatives from a batch of 1024, so correlated or similar examples can weaken the contrast.Source: A Training Details and the Behavioral… · p. 14
• The reported Stage-1 grid varies ℓ1 gate-penalty coefficients for hinge and InfoNCE objectives, with fixed optimization settings and separate batch sizes.The reference hinge setting uses margin κ = 1 and λ = 0.03; InfoNCE uses τ = 1 and λ = 0.10.Source: A Training Details and the Behavioral… · p. 14
• Stage 2 detaches learned gates and applies them after state standardization; its batches mix within-trajectory and cross-episode frame pairs.The optional state-regression head is disabled; Appendix gate values are original sigmoid gates before planning-analysis normalization.Source: A Training Details and the Behavioral… · pp. 14–15

• Applying gates after state standardization preserves their weighting; a second coordinate-wise standardization would cancel it and restore uniform weighting.Source: B Metric Transfer and Interpretation · p. 15
• Stage 1 does not identify an absolute gate scale, but common positive rescaling leaves Stage-2 Pearson correlation unchanged, so relative weights matter.Source: B Metric Transfer and Interpretation · p. 15

C State-Gate Values Across Tasks

• The tables report continuous, unthresholded Stage-1 gate values for Cube-Double (27 dimensions), Cube-Triple (32), Cube-Quadruple (37), and Scene (26), using ℓ1-regularized Hinge and InfoNCE models at four sparsity coefficients.Source: C State-Gate Values Across Tasks · p. 15, Table 2 · p. 15, Table 3 · p. 16
• Corresponding coordinates across repeated cube slots often receive closely matched gates, including 0.95 for the listed y coordinates at λ = 0.01.This pattern suggests shared coordinate roles, without implying exact equality at every coefficient.Source: C State-Gate Values Across Tasks · p. 15

D Structured-Distractor Diagnostics

• The contamination diagnostic compares gates on the original 22 Cube coordinates between clean and augmented training, using rank and linear correlations, gate-change summaries, and thresholded-set overlap.Table 6 specifies InfoNCE with ℓ1 regularization; Table 7 covers contaminated-state Cube planning success percentages from a single training seed.Source: D Structured-Distractor Diagnostics · p. 18, Table 6 · p. 18, Table 7 · p. 18

E Experimental Details and Planning Protocol

• The five OGBench tasks have state dimensions 22, 27, 32, 37, and 26; each LeWM-based method uses five training seeds and 600 episodes per task, configuration, and planner.CEM evaluates 300 action sequences per iteration and retains 30; OGBench uses 10 iterations, a five-decision-step horizon with frame skip 5, and at most 50 environment steps per episode.Source: E Experimental Details and Planning… · p. 19, Table 8 · p. 19

F.1 Loss-Component Ablations

• Removing NLL leaves no gate at the reporting threshold, while NLL-retaining variants keep 13–20; removing Hinge discrimination reduces retained dimensions from 15 to 13.Ablations use λ = 0.03 for both discrimination objectives; success rates are similar across the four variants.Source: §F.1 · pp. 19–20, Table 9 · p. 19, F Ablation Studies · p. 19

F.2 Gate Penalty and Discrimination Objective

• On Cube, all 16 SWIM configurations exceed SCALE’s nine manually selected dimensions under CEM (74.5–77.0% vs. 73.5%) and are comparable under iCEM (75.2–78.0% vs. 76.7%).The ℓ1- and squared-ℓ2-regularized gates have Spearman rank correlations of 0.97 at λ = 0.01 and 0.99 at λ = 0.1; at λ = 0.3 and 1.0, squared-ℓ2 gates remain graded rather than approaching zero.Source: §F.2 · p. 20, Table 10 · p. 20

F.3 Hinge-Margin and InfoNCE-Temperature Sensitivity

• Hinge success changes modestly across tested margins, while InfoNCE τ = 1 has the highest mean across CEM and iCEM among tested temperatures; the defaults are κ = 1 and τ = 1.Larger InfoNCE temperatures do not improve CEM success.Source: §F.3 · p. 20

F.4 Data Efficiency of Gate Learning

• For each subset size, the analysis samples three subsets of complete trajectories from the 2,000-trajectory dataset and averages the selected-dimension count across them.The full-data gates are the reference for Jaccard and Spearman comparisons.Source: §F.4 · pp. 20–21
• With λ = 0.1, 250–500 trajectories reach 0.80 Jaccard overlap and 1,000 reach Spearman 0.94 against the 2,000-trajectory reference.The full-data gates provide the reference for both comparisons.Source: §F.4 · p. 21, Table 12 · p. 21
• At λ = 0.01, ranking is less stable at small sample sizes, with Spearman 0.81 at 500 trajectories; 25 trajectories retain 13 dimensions on average versus 5 using all 2,000.Selected dimensions are coordinates with g_k ≥ 0.5.Source: §F.4 · pp. 20–21

G State-Weight Comparison and Planning Proofs

• Normalizing weights preserves candidate ordering and, for nonconstant distance vectors, leaves ideal Pearson correlation unchanged.In the appendix, g denotes normalized weights in the planning analysis; elsewhere it denotes Stage-1 sigmoid gates.Source: G State-Weight Comparison and Planning… · p. 21
• For exact state-metric candidate selection, R^2 εβ ≤ Δg is sufficient for the stated comparison with uniform weighting; strict inequality gives strict improvement.The result concerns exact state-metric minimization, an analytical benchmark rather than the deployed latent-cost search.Source: G State-Weight Comparison and Planning… · pp. 22–23
• For comparison with SCALE-All, SWIM’s total upper bound is strictly smaller when the weight-discrepancy reduction exceeds any increase in combined calibration and rollout terms.The comparison assumes the stated conditions hold for both models on the same candidates, standardized coordinates, and reference metric.Source: G State-Weight Comparison and Planning… · p. 22
• The planning proof bounds selected-candidate cost using latent-distance calibration and rollout error, then translates the result back to the reference metric.The conclusion also covers arbitrary ties among minimizers.Source: G State-Weight Comparison and Planning… · p. 25