PaperlayerView exact revision on arXiv

Source-linked AI summary

No Gaussian Required: Contrastive Inverse Dynamics for JEPA World Models

Jack Boylan, Chris Hokamp

arXiv:2608.17542v1cs.LGcs.AI

TL;DR

JEPA world models can collapse to constant representations, while existing anti-collapse methods prescribe latent distributions that may not fit the environment. This paper replaces that prescription with contrastive inverse dynamics and reports stronger performance on the harder OGBench Visual Scene task while remaining competitive across pixel-control tasks.

  • Problem

    Existing JEPA world models need anti-collapse mechanisms, but it remains unclear whether transition data can provide this pressure without imposing a fixed latent distribution.

  • Method

    AC-MTM adds a training-only Action-NCE inverse-dynamics head that identifies each transition’s action while retaining LeWM’s forward latent-prediction objective.

  • Results

    AC-MTM shows mixed pixel-control results but reaches 80.0% success versus 58.0% for SIGReg on OGBench Visual Scene under matched trajectory-goal planning.

  • Takeaways & Limitations

    Contrastive inverse dynamics provides a distribution-free, dynamics-native anti-collapse signal, while PushT shows that weakly action-controlled task-relevant state can remain underrepresented.

  • Takeaways & Limitations

    OGBench Visual Scene uses a matched trajectory-goal protocol rather than the official fixed-goal leaderboard, so its 80% result is not a public benchmark score.

Abstract

from arXiv · show

Joint-Embedding Predictive Architectures (JEPAs) learn world models by predicting future embeddings, but the objective admits a trivial solution of a constant encoder, so every practical system adds an anti-collapse mechanism (LeCun, 2022; Assran et al., 2023; Bardes et al., 2022; 2024). LeWorldModel (LeWM) prevents collapse with SIGReg, a regularizer that forces the latent distribution to match an isotropic Gaussian: the representation is stabilized by prescribing what it must look like, independently of the environment it models. We argue that the anti-collapse pressure can instead come from the transition data itself. Action-Contrastive Masked Transition Modeling (AC-MTM) keeps LeWM's forward latent-prediction objective and adds a training-only inverse-dynamics head trained with Action-NCE: each latent transition must identify the action that produced it among the other actions in the batch, a discrimination task that a collapsed encoder provably fails. The inverse branch is discarded after training, leaving test-time encoding, forward prediction, planning, and compute identical to LeWM. On four standard pixel-control tasks under a matched planning protocol, AC-MTM trains stably from scratch and matches SIGReg on average. On the harder multi-object OGBench Visual Scene task, results are consistent with the prescribed geometry becoming a bottleneck: AC-MTM reaches 80.0$\pm$2.0% success versus 58.0$\pm$2.0% for SIGReg, improving by 20-24 points in each training seed. A single 50-episode random-policy run gives a 52% baseline estimate. Contrastive inverse dynamics thus provides a distribution-free anti-collapse signal that requires no target network, stop-gradient, pretrained encoder, or reconstruction objective, and we characterize the action-space and observability assumptions under which it holds. We make our code available at https://github.com/jackboyla/action-contrastive-jepa

1 INTRODUCTION

AC-MTM derives anti-collapse pressure from transition data by adding a training-only contrastive inverse-dynamics task to LeWM’s forward latent prediction. It matches SIGReg on standard pixel-control tasks, substantially outperforms it on OGBench Visual Scene, and leaves test-time planning unchanged.

  • Inference: The inverse branch is discarded after training, so encoding, forward prediction, planning, and computation at test time remain identical to LeWM.Evaluation uses only the unchanged encoder and forward predictor; the planner never calls the inverse head.
  • Motivation: AC-MTM avoids prescribing a global isotropic-Gaussian latent geometry, instead deriving anti-collapse pressure from transition data.The motivation is that SIGReg’s prescribed geometry may bottleneck performance in low-intrinsic-dimensionality environments.
  • Method: AC-MTM adds Action-NCE, requiring each latent transition to identify its observed action among other batch actions while retaining LeWM’s forward latent prediction.The inverse head is used only during training, and a constant encoder produces identical queries that cannot identify actions.
  • Standard-task results: AC-MTM never collapses across tasks and seeds, matches SIGReg on four standard tasks, wins two, ties one, and loses PushT.The comparison uses LeWM’s one-stage training recipe and evaluates with the same encoder, autoregressive predictor, latent distance, and CEM planner.
  • Scene stress test: 80.0% versus 58.0%: AC-MTM outperforms SIGReg on OGBench Visual Scene under the matched trajectory-goal MPC protocol.The gains are 24, 20, and 22 points across the three training seeds; a 50-episode random-policy run gives a 52% baseline estimate.
  • Inverse-signal reliability: Action-NCE removes inverse-regression bimodality on Reacher, while both inverse variants trail SIGReg on PushT.MTM-MSE collapses on two of three Reacher seeds, Action-NCE costs 3.8 points on a long-horizon stress test, and weakly controlled object state limits action identification on PushT.

2 BACKGROUND AND POSITIONING

The paper studies reward-free latent planning from offline pixel trajectories, positioning AC-MTM among distribution-free alternatives to Gaussian-matching anti-collapse objectives. LeWM provides the key comparison point by simplifying prior collapse prevention to prediction plus SIGReg.

  • Reward-free latent planning: Reward-free planning uses offline pixel trajectories without rewards or optimality assumptions, while test-time CEM searches action sequences by minimizing latent distance to an encoded goal.Recorded continuous commands connect observations during training, but action labels are unavailable at test time.
  • From PLDM to LeWM: LeWM replaces PLDM’s multi-term VICReg-derived objective with prediction plus SIGReg, which matches every one-dimensional latent marginal to a Gaussian.The Gaussian matching is motivated by the Cramér–Wold theorem.
  • A family of distribution-free signals: MTM-MSE regresses the observed action from an adjacent latent pair, while SMWM tunes its inverse-MSE loss weight per environment.The paper’s controlled baseline instead uses one coefficient selected on Reacher and held fixed across tasks.

3 ACTION-CONTRASTIVE MASKED TRANSITION MODELING

AC-MTM preserves LeWM’s forward latent-prediction task and adds a training-only contrastive inverse-dynamics head that identifies each transition’s action among in-batch alternatives. This signal opposes encoder collapse without SIGReg and leaves test-time prediction, planning, and computation unchanged.

  • Method: AC-MTM adds a training-only inverse-dynamics head that classifies each transition’s observed action against N in-batch action candidates.The candidates are raw actions, and negatives add only an N × N distance matrix.
  • Objective: The objective is L_AC-MTM = L_fwd + λL_NCE with λ = 0.30 and τ = 0.10, and it contains no SIGReg term.The coefficient was selected by a bounded Reacher stability sweep and then fixed across tasks.
  • Method: “Masked” means jointly predicting z_t+1 from (z_t, a_t) and a_t from (z_t, z_t+1), rather than masking image patches or sampling one task per batch.Both forward and inverse tasks are optimized on every batch.
  • Why Action-NCE opposes collapse: Under a constant encoder, forward prediction can achieve zero loss, but Action-NCE reduces only to the chance classifier, forcing transition information that makes actions more identifiable.Collapsed inputs produce identical classifier rows and cannot improve contrastive inverse loss below chance.
  • Ablation: MTM-MSE is strong on TwoRoom, PushT, and Cube, but two of three Reacher seeds collapse because regression provides only an action-variance floor under complete collapse.Its margin can be weak when multiple action blocks produce similar visual endpoints.
  • Controlled implementation: AC-MTM inherits LeWM’s encoder, predictor, conditioning, datasets, and CEM planner; its inverse head is removed after training, preserving identical test-time computation and planning.The head is never used during rollout or cost evaluation.

4 EXPERIMENTS

Across matched pixel-control experiments, AC-MTM prevents collapse through transition supervision and performs competitively with SIGReg using the same test-time planner. Its largest advantage appears on the multi-object OGBench Visual Scene task, while results also expose task-dependent tradeoffs and representation limitations.

  • Evaluation Setup: Experiments use end-to-end pixel training across five control environments evaluated with the same CEM/MPC planner.CEM samples 300 action sequences, retains 30 elites, and refits its distribution for 30 iterations.
  • Anti-Collapse Ablation: NoReg drives forward loss to ≈0, whereas MTM-MSE and AC-MTM prevent the trivial constant-latent solution through transition supervision.Both inverse-task variants preserve the same one-stage training recipe.
  • Controlled Comparison: AC-MTM improves over SIGReg on TwoRoom, Cube, and Scene, matches it on Reacher, and trails it on PushT under the shared test-time planner.Figure 3 additionally contextualizes results with external paper-reported baselines, while statistical claims come from matched SIGReg reruns.
  • OGBench Visual Scene: 80.0% success for AC-MTM versus 58.0% for SIGReg on OGBench Visual Scene shows the largest margin on the multi-object manipulation task.MTM-MSE reaches 75.3±2.3% with no collapsed seed and beats SIGReg in paired comparisons.
  • Stability and Long-Horizon Stress Test: 28.0±0.8% for the non-contrastive inverse-regression ablation versus 24.2±0.6% for AC-MTM on TwoRoom-long, while AC-MTM remains more reliable across Reacher seeds.MTM-MSE collapses in two Reacher runs, whereas all AC-MTM runs remain useful.
  • Representation Analysis: Both inverse variants underencode T-block orientation on PushT while preserving agent position and much of block position.The probes identify orientation as a weakly controlled variable determining contact geometry.

5 ANALYSIS

The analysis explains why Action-NCE can stabilize representations, identifies coverage limits and probe-planning mismatches, and shows that inverse dynamics changes training rather than test-time inference. Model-facing diagnostics further indicate that both objectives detect physically inconsistent transitions, with normalized ratios requiring cautious comparison.

  • Why Action-NCE stabilizes Reacher: Action-NCE changes the collapse geometry when inverse-MSE has a weak regression margin, because ambiguous multi-step action sequences can share similar Reacher endpoints.Forward latent MSE simultaneously rewards shrinking the online representation.
  • What Action-NCE does not solve: Action-NCE emphasizes action-separating state and therefore does not unconditionally preserve every task-relevant variable, unlike SIGReg’s global variance pressure.PushT requires block pose and orientation even when commanded actions have little visible effect before contact.
  • Why probes are diagnostic, not decisive: Linear probe decodability and planning success need not align: MTM-MSE and AC-MTM under-decode orientation yet plan substantially, whereas AC-CPC decodes orientation better but plans worse.Representation audits and planner behavior should therefore be reported together.
  • Training-only means test-time equivalence: The inverse head influences encoder optimization but is absent at test time, where CEM alone proposes and scores actions as in LeWM.The empirical differences therefore isolate training-time representation learning rather than stronger inference or extra runtime computation.
  • Model-facing surprise diagnostics: The diagnostic compares normal one-step latent prediction error with error after corrupting the same trajectory clip using action counterfactuals or physically inconsistent transitions.It is applied on PushT and OGBench-Cube.
  • Model-facing surprise diagnostics: Both objectives flag both corruption types with errors one to three orders of magnitude above normal prediction error on essentially every clip.Table 8 reports latent surprise ratios, where higher values indicate greater sensitivity to violations; ratios are normalized by each model’s own normal-transition error.

6 RELATED WORK

Prior work spans predictive world models, reconstruction-free latent planning, and JEPA anti-collapse objectives. AC-MTM most directly extends SMWM’s training-only inverse branch by replacing inverse-action MSE with contrastive Action-NCE while preserving LeWM’s deployed model.

  • Reconstruction-free latent planning: PLDM plans from reward-free offline data without pixel reconstruction but uses a VICReg-derived multi-term recipe and inverse-dynamics regularization.DINO-WM instead learns dynamics on frozen DINOv2 features, trading end-to-end simplicity for strong pretrained visual representations.
  • Reconstruction-free latent planning: LeWM is the closest predecessor, training end-to-end from pixels with next-latent prediction plus SIGReg while retaining the same deployed model.AC-MTM removes LeWM’s remaining global distributional assumption during training.
  • Inverse dynamics and probing: SMWM uses the LeWM encoder, forward predictor, and training-only inverse head, with inverse-action MSE as its sole anti-collapse term.The paper treats MTM-MSE as an SMWM-style baseline rather than a new mechanism.
  • Inverse dynamics and probing: SMWM finds that inverse MSE roughly matches SIGReg on the 2D tasks and improves on Cube.AC-MTM’s contribution is the contrastive Action-NCE form and its chance-level collapse result.
  • Benchmarks and goal-conditioned baselines: The evaluation aligns AC-MTM with the LeWM suite of top-down navigation, DMC-style reaching, PushT manipulation, and OGBench-Cube environments.External comparisons include PLDM, DINO-WM, goal-conditioned behavioral cloning, goal-conditioned offline RL, and random policies.

7 LIMITATIONS

AC-MTM’s anti-collapse signal depends on informative, varied action effects and offers no unconditional geometric guarantee. Its evaluation leaves open broader action-space scaling and whether latent diagnostics imply robust out-of-distribution physical reasoning.

  • 7 LIMITATIONS: AC-MTM provides no unconditional geometric guarantee and requires varied action candidates plus useful visual evidence about action effects.Unobserved actuators, stochastic dynamics, no-op-heavy datasets, and duplicate actions can weaken the contrastive task.
  • 7 LIMITATIONS: Normalized continuous controls were evaluated, while discrete, hybrid, structured, and very high-dimensional action spaces may require different scores, learned embeddings, or hard-negative sampling.Training also depends on batch composition, temperature τ, and coefficient λ; broader scaling evidence is needed.
  • 7 LIMITATIONS: Latent surprise diagnostics do not replace human-designed physical-event suites or prove robust out-of-distribution physical reasoning.OGBench Visual Scene used the trajectory-goal protocol rather than the official five-task fixed-goal protocol.

8 CONCLUSION

AC-MTM protects world-model representations from collapse without prescribing a latent distribution, replacing Gaussian matching with a contrastive signal grounded in transition dynamics. It remains a one-stage, end-to-end pixel JEPA without several auxiliary mechanisms.

  • Conclusion: AC-MTM replaces LeWM’s Gaussian matching with a contrastive, dynamics-native anti-collapse signal.The signal retains information identifying the action that drove each transition.
  • Conclusion: The method protects representations from collapse without requiring a prescribed latent distribution.The conclusion states that collapse protection need not take the form of matching a fixed distribution.
  • Conclusion: AC-MTM is a one-stage, end-to-end pixel JEPA without a target network, stop-gradient branch, frozen encoder, or reconstruction decoder.The passage also mentions a fixed global latent di..., but the supplied text is truncated at that phrase.

A EVALUATION PROTOCOL

The controlled standard-task comparisons use a fixed 200-episode planning protocol, exact epoch-10 checkpoints, and no validation-based checkpoint selection. Non-LeWM Figure 3 baselines are paper-reported values included for continuity.

  • Controlled comparisons: 200 episodes, evaluation seed 42, exact checkpoint stems, and unique output files define the controlled standard-task planning protocol.CEM uses 300 samples, 30 elites, and 30 iterations.
  • Controlled comparisons: Final epoch-10 checkpoints from training seeds {3072, 1, 2} are used for SIGReg, MTM-MSE, and AC-MTM.No validation-based checkpoint selection is used in the controlled tables.
  • Baseline reporting: Figure 3 includes paper-reported non-LeWM baseline values from Maes et al. (2026) for continuity with the original evaluation style.

B INTERROGATING THE OGBENCH SCENE RESULT

The Scene advantage survives the reported audit only within the matched trajectory-goal MPC protocol, while its likely mechanism involves controllable state information and lower forward-prediction error. The official fixed-goal protocol fails for both methods, and further sweeps and per-factor diagnostics are needed to test the interpretation.

  • Scope of the result: The Scene headline is evidence about the matched trajectory-goal MPC protocol, not the public OGBench fixed-goal benchmark.The audit preserves the result but narrows its causal interpretation.
  • Why the gap may be large: Scene couples drawer, window, two buttons, and cube state to one arm observation, making final-latent MPC brittle when any variable is dropped.A rollout may match agent pose while missing the target object configuration.
  • Why the gap may be large: AC-MTM favors latent variables whose changes explain executed controls, aligning its inverse-dynamics pressure with the CEM rollout interface.Action identification among in-batch alternatives requires adjacent latents to retain information about controllable changes.
  • Why the gap may be large: 0.002 versus 0.004: AC-MTM reaches a lower final one-step forward MSE on Scene than SIGReg in the last logged batches.The observed advantage may therefore combine better controllable-state retention with improved forward prediction.
  • Official fixed-goal protocol: 0/250: both SIGReg and AC-MTM score 0/250 on the official fixed-goal protocol, with every episode reaching the 750-step cap.These tasks require much longer-horizon manipulation than the 25-step trajectory goals, and the final-latent CEM planner does not solve them at this model scale.
  • What would falsify the interpretation?: An evaluation-seed sweep or larger n beyond the current 150 matched episodes, plus per-factor diagnostics, would test whether the paired advantage generalizes and identify lost Scene variables.Suggested probes target drawer, window, button, and cube factors or decompose success per target.

C CONTROLLED THREE-WAY ABLATION

The controlled ablation isolates Action-NCE’s contribution against SIGReg and non-contrastive inverse regression. AC-MTM preserves performance on most tasks while preventing Reacher collapse and leading on Scene.

  • Controlled three-way ablation: Action-NCE is isolated by comparing SIGReg, non-contrastive inverse regression, and AC-MTM under a shared autoregressive planner.Evaluation uses 200 episodes on standard tasks and 50 on OGBench-Scene, with means and standard deviations over three training seeds.
  • Controlled three-way ablation: Within 1.2 percentage points of MTM-MSE, AC-MTM remains competitive on TwoRoom, PushT, and Cube.The comparison covers the three standard tasks named in the ablation passage.
  • Controlled three-way ablation: Two of three seeds collapse for MTM-MSE on Reacher, whereas AC-MTM remains non-collapsed without changing the test-time planner.Figure 5 identifies removal of the non-contrastive inverse-regression ablation’s bimodal collapse as the motivation for using Action-NCE.
  • Controlled three-way ablation: 4.7 points: AC-MTM leads MTM-MSE on Scene, while both inverse variants decisively outperform SIGReg.The reported paired comparison is 16/9 with p≈0.23.

D EXPLORATORY SINGLE-SEED RESULTS ON ADDITIONAL OGBENCH FAMILIES

Single-seed exploratory comparisons across six additional OGBench families show method-dependent patterns rather than a uniform winner: SIGReg leads on button puzzles, AC-MTM modestly leads on stochastic-teleport maze and Powderworld, and they tie on stitching tasks. The puzzle disadvantage for AC-MTM is attributed to actions producing visually near-identical button-state changes, weakening inverse-dynamics pressure to encode the full configuration.

  • Experimental scope: Single-seed comparisons on six additional OGBench environments used the same trajectory-goal protocol: 50 episodes, evaluation seed 42, one training seed, and no tuning.The runs are exploratory scope evidence rather than controlled claims.
  • Results: SIGReg clearly outperformed AC-MTM on combinatorial button-puzzle tasks, while AC-MTM was modestly better on stochastic-teleport maze and Powderworld, and the methods tied on stitching tasks.These qualitative comparisons come from exploratory single-seed trajectory-goal success results.
  • Mechanistic interpretation: On button puzzles, visually near-identical state changes across actions provide little inverse-dynamics pressure to represent the full button configuration, whereas SIGReg’s global variance pressure preserves it.Frozen-latent probes support this mechanism: on Puzzle 4x4, SIGReg decoded variable button bits at R2 ≈0.98, while the supplied passage truncates the AC-MTM value.
Loading 2608.17542v1…