No Gaussian Required: Contrastive Inverse Dynamics for JEPA World Models

Jack Boylan; Chris Hokamp

cs.LGAug 18, 2026

No Gaussian Required: Contrastive Inverse Dynamics for JEPA World Models

Authors: Jack Boylan, Chris Hokamp

Organizations: Quantexa

Abstract

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

Figures & tables

Figure 1: Training and inference paths for AC-MTM. During training, the observed action is an input to the forward predictor and the positive label for the contrastive inverse task. At test time, CEM samples candidate actions and uses only the unchanged encoder and forward predictor; the inverse dynamics/Action-NCE branch is discarded.

Method	Anti-collapse signal	Latent geometry prescribed?
PLDM ( Sobal et al., 2025 )	VICReg variance/covariance (multi-term)	Yes (variance floor)
LeWM ( Maes et al., 2026 )	SIGReg isotropic-Gaussian matching	Yes (global Gaussian)
MTM-MSE ( Ivashkov et al., 2026 )	Inverse-action MSE $(\mathbf{z}_{t},\mathbf{z}_{t+1})\!\to\!\mathbf{a}_{t}$	No
AC-MTM (ours)	Contrastive inverse-action identification	No
AC-CPC	Contrastive future identification	Implicit (unit sphere)

Table 1: Anti-collapse mechanisms compared in this paper. AC-MTM uses negatives during training, but does not prescribe a global latent marginal and does not change the test-time planner.

Panel of Figure 2 — Figure 2: Evaluation suite. TwoRoom is low-dimensional navigation through a wall opening; Reacher matches a target arm configuration; PushT requires contact-rich pushing of a T-shaped block to a goal pose; OGBench-Cube is 3D robot-arm cube manipulation; OGBench-Scene adds a multi-object drawer/window/button/cube scene. These tasks stress different failure modes: low intrinsic dimension, articulated dynamics, contact, 3D manipulation, and multi-object visual complexity.

Model	TwoRoom success (%)
NoReg: no anti-collapse term	$28.0{\pm}2.0$
LeWM with SIGReg	$85.5{\pm}0.4$
MTM-MSE	$90.2{\pm}0.5$
AC-MTM	$\mathbf{90.7{\pm}0.6}$

Table 2: TwoRoom anti-collapse ablation with 200 evaluation episodes per seed. Values are mean

\pm

standard deviation over three training seeds. Plain next-latent prediction reaches a trivial low-loss solution and plans substantially worse.

Task	SIGReg (LeWM)	AC-MTM (ours)	$\Delta$
TwoRoom	$85.5{\pm}0.4$	$\mathbf{90.7{\pm}0.6}$	$+5.2$
Reacher	$68.8{\pm}0.2$	$68.3{\pm}3.1$	$-0.5$
PushT	$\mathbf{93.2{\pm}0.2}$	$86.7{\pm}1.5$	$-6.5$
OGB-Cube	$66.2{\pm}0.2$	$\mathbf{78.8{\pm}1.7}$	$+12.6$
OGB-Scene	$58.0{\pm}2.0$	$\mathbf{80.0{\pm}2.0}$	$+22.0$

Table 3: Planning success (%) with the shared autoregressive planner: 200 evaluation episodes per seed on the standard tasks and 50 on OGBench-Scene. Values are mean

\pm

standard deviation over three training seeds; bold marks differences larger than the cross-seed noise. Figure 3 places these numbers next to external paper-reported baselines.

Figure 3: Planning success across the five environments. Salmon and blue bars are paper-reported LeWM/external baselines from Maes et al. (2026) ; purple, gold, and green bars are our controlled SIGReg, MTM-MSE, and AC-MTM evaluations across three training seeds. No external baseline has published OGBench-Scene results under this protocol, so the Scene panel contains only our controlled models and the random policy. The external baselines place AC-MTM in the original evaluation frame; direct statistical claims are made only for the controlled comparisons in Table 3 .

Method	Seed 3072	Seed 1	Seed 2	Mean	Paired wins/losses vs. SIGReg
Random	–	–	–	$52.0$	–
SIGReg	56.0	58.0	60.0	$58.0{\pm}2.0$	–
MTM-MSE	78.0	74.0	74.0	$75.3{\pm}2.3$	39 / 13
AC-MTM	80.0	78.0	82.0	$\mathbf{80.0{\pm}2.0}$	40 / 7

Table 4: OGBench Visual Scene trajectory-goal success (%), with 50 evaluation episodes per training seed and evaluation seed 42. The random policy is a single 50-episode evaluation under the identical protocol and is a baseline estimate. Paired wins/losses pool all 150 outcomes against SIGReg: two-sided binomial/McNemar

p\!\approx\!1.1{\times}10^{-6}

for AC-MTM (40/7) and

p\!\approx\!4.1{\times}10^{-4}

for MTM-MSE (39/13); the AC-MTM–MTM-MSE pairing is 16/9 (

p\!\approx\!0.23

). Because the same 50 tasks recur across the three training seeds, these are episode-level descriptive tests, not tests with 150 independent experimental units. The seed-level gains for AC-MTM over SIGReg are 24, 20, and 22 points. This is not a public OGBench leaderboard claim; the official fixed-goal protocol is separate. Appendix B audits possible explanations for the large margin.

Method	Seed 3072	Seed 1	Seed 2	Mean
SIGReg	69.0	69.0	68.5	68.8
MTM-MSE	68.0	11.5	13.5	31.0
AC-MTM	70.5	70.5	64.0	68.3

Table 5: Reacher success (%) by training seed, with 200 evaluation episodes per checkpoint. MTM-MSE succeeds once but collapses twice; AC-MTM removes that bimodality.

TwoRoom-long ( $100/150$ )	SIGReg	MTM-MSE	AC-MTM
Success (%)	$17.0{\pm}0.4$	$\mathbf{28.0{\pm}0.8}$	$24.2{\pm}0.6$

Table 6: Long-horizon stress test with 200 evaluation episodes and three training seeds per method. AC-MTM retains most of the inverse-dynamics gain over SIGReg but trails MTM-MSE by 3.8 points.

Mechanism	Planner	Orient. $R^{2}{}$	Mean $R^{2}{}$	Success (%)
SIGReg (LeWM)	AR	0.791	0.701	93.2
MTM-MSE	AR	0.508	0.674	85.5
AC-MTM	AR	0.514	0.675	86.7
AC-CPC	AR	0.655	0.564	62.5

Table 7: PushT anti-collapse mechanisms: frozen-latent probes (

n{=}4000

\alpha{=}1

; orientation is

\mathrm{state}[4]

) and 200-episode planning. AC-MTM probe values are means over three training seeds; all rows use the same autoregressive planner.

Task	Model	Action counterfactual ( $\uparrow$ )	State discontinuity ( $\uparrow$ )
PushT	SIGReg	$151.7{\pm}4.7{\times}$	$1246.0{\pm}34.1{\times}$
PushT	AC-MTM	$9.6{\pm}0.3{\times}$	$40.0{\pm}1.4{\times}$
OGB-Cube	SIGReg	$274.2{\pm}12.6{\times}$	$835.4{\pm}35.2{\times}$
OGB-Cube	AC-MTM	$82.3{\pm}1.6{\times}$	$434.4{\pm}5.6{\times}$

Table 8: Latent surprise ratios: one-step prediction error on a corrupted (physically invalid) transition divided by the same clip’s error on the real transition.

1{\times}

would mean the model cannot tell invalid transitions from real ones; higher means the model is more surprised by violations. Each entry uses 4096 sampled clips per training seed (mean

\pm

standard deviation over seeds

\{3072,1,2\}

). In every row, the invalid transition has higher error than the matched real transition for at least 99.95% of clips. Ratios are normalized by each model’s own baseline error, so they should be compared within a model, not across models (see text).

Appendix figures & tables6 assets

Supplementary material from the paper’s appendix.

Appendix

Task	Controlled evaluation decision
TwoRoom	Use released/maintainer-confirmed goal offset $25$ , budget $50$ . Issue #38 records maintainer confirmation that the paper’s $100/150$ text was a typo; issue #67 documents the same mismatch and low success under $100/150$ . We keep $100/150$ only as the separately labeled TwoRoom-long stress test.
PushT	Use released config: goal offset $25$ , budget $50$ , CEM $300/30/30$ .
Reacher	Use released config: goal offset $25$ , budget $50$ , CEM $300/30/30$ . Issue #41 records that the paper text says $10$ CEM iterations for non-PushT tasks while the released config uses $30$ ; no maintainer correction is visible. Issues #37 and #62 show Reacher sensitivity to training/eval seed.
OGBench-Cube	Use released config: goal offset $25$ , budget $50$ , CEM $300/30/30$ , with the same CEM caveat as Reacher from issue #41.
OGBench-Scene	Added complex-task trajectory-goal protocol using visual-scene-play-v0 : goal offset $25$ , budget $50$ , CEM $300/30/30$ , 50 episodes per training seed, and matched train seeds $\{3072,1,2\}$ for both SIGReg and AC-MTM. This is separate from the official OGBench Visual Scene five-task fixed-goal protocol.

Appendix

Table 9: Protocol audit for the 200-episode controlled comparisons. Issue numbers refer to the public LeWM repository: https://github.com/lucas-maes/le-wm/issues : #38, #67, #41, #37, and #62.

Check	Finding	Interpretation
Episode pairing	All six result files use evaluation seed 42, 50 episodes, goal offset 25, budget 50, CEM $300/30/30$ , horizon 5, and action block 5. The raw success arrays give 40 AC-MTM-only wins and 7 SIGReg-only wins over 150 paired episodes.	The gap is not an artifact of comparing different episode sets.
Training comparability	The resolved training configs match on dataset ogbench/visual_scene_play , batch size 128, learning rate $5{\times}10^{-5}$ , 10 epochs, history size 3, encoder, action encoder, and forward predictor. They differ in the intended objective: SIGReg weight 0.09 versus Action-NCE with inverse weight 0.30 and no SIGReg.	The comparison isolates the anti-collapse/training signal up to the small training-only inverse head.
Checkpoint selection	Eval logs load the epoch-10 object checkpoint for both SIGReg and AC-MTM.	The result is not caused by accidentally evaluating an earlier or best-selected checkpoint.
Privileged-state leakage	The Scene evaluator uses privileged future state to reset simulator targets and define success, but the world-model cost path encodes goal pixels and drops goal action before computing final latent distance.	This matches the trajectory-goal setup: privileged state defines the task, not the planner input.
Collapse	Final training logs show nonzero latent scale for both methods (roughly 0.99 mean embedding standard deviation for SIGReg and 0.24 for AC-MTM over the last logged batches).	The SIGReg model is not constant-collapsed; the gap is a planner-usable-geometry or dynamics-quality failure, not trivial collapse.
Protocol scope	The goal is a future visual observation from visual-scene-play-v0 ; public OGBench Scene tables average five fixed goal tasks with a 750-step cap.	The 80% number should not be compared to public OGBench leaderboard scores.

Appendix

Table 10: Audit of the completed OGBench Visual Scene trajectory-goal comparison.

Figure 4: Action blocks and CEM rollouts (TwoRoom, $d{=}2$ ). Each coarse transition bundles $K{=}5$ raw simulator actions into one model input block. CEM samples $H{=}5$ -block candidate sequences, rolls them through the unchanged latent forward model, scores final latent-goal distance, refits to elites, and executes the final mean blocks as ordered low-level env.step actions before re-planning.

Task	SIGReg (LeWM)	MTM-MSE	AC-MTM (ours)
TwoRoom	$85.5{\pm}0.4$	$90.2{\pm}0.5$	$\mathbf{90.7{\pm}0.6}$
Reacher	$\mathbf{68.8{\pm}0.2}$	$31.0{\pm}26.2$	$68.3{\pm}3.1$
PushT	$\mathbf{93.2{\pm}0.2}$	$85.5{\pm}0.7$	$86.7{\pm}1.5$
OGB-Cube	$66.2{\pm}0.2$	$\mathbf{79.3{\pm}2.4}$	$78.8{\pm}1.7$
OGB-Scene	$58.0{\pm}2.0$	$75.3{\pm}2.3$	$\mathbf{80.0{\pm}2.0}$

Appendix

Table 11: Controlled planning success (%) with the shared autoregressive planner: 200 evaluation episodes per seed on the standard tasks and 50 on OGBench-Scene. Values are mean

\pm

standard deviation over three training seeds.

Task	SIGReg	AC-MTM	$\Delta$
Visual Puzzle 4x4 (play)	50.0	34.0	$-16$
Visual Puzzle 4x5 (play)	52.0	26.0	$-26$
Visual AntMaze teleport (navigate)	40.0	46.0	$+6$
Powderworld medium (play)	6.0	16.0	$+10$
Visual AntMaze large (stitch)	30.0	30.0	$0$
AntSoccer medium (stitch)	88.0	88.0	$0$

Appendix

Table 12: Exploratory single-seed (3072) trajectory-goal success (%) on additional OGBench families, 50 episodes each. SIGReg is clearly better on the combinatorial button-puzzle tasks, AC-MTM modestly better on the stochastic-teleport maze and Powderworld, and the two tie on the stitching tasks.

Joint-Embedding Predictive Architecture

Joint-Embedding Predictive Architecture (JEPA) learns representations by predicting latent features of hidden or future data from observed context.

See all 246 papers

Aug 6, 2026

PhyLatent: Learning Dynamics-Relevant Representations for JEPA World Models

We propose PhyLatent, a dynamics-relevant training objective for Joint-Embedding Predictive Architecture (JEPA) world models. Our key observation is that preventing global latent collapse does not necessarily ensure that the learned representation preserves physically meaningful state and action relationships. We identify three failure modes: sensitivity to appearance changes that leave the physical state unchanged, insufficient separation of distinct physical states, and insufficient separation of different action-conditioned futures. We refer to these as Physical Invariance Collapse, Physical Distinguishability Collapse, and Counterfactual Dynamics Collapse, respectively. PhyLatent targets these failures through three coordinated training pathways, implemented with static visual invariance, physical state grounding, future representation alignment, counterfactual branch separation, and latent denoising. On OGBench-Cube, PhyLatent reduces the three collapse rates by 43.9%, 27.9%, and 47.1%, respectively, while improving model predictive control (MPC) success by 12.0 percentage points (17.2% relative). Across four visual-control tasks, average planning success increases by 6.62 percentage points (8.3% relative). These results show that global non-collapse alone is insufficient for learning a reliable JEPA world-model state space, and that explicitly preserving dynamics-relevant structure can improve closed-loop planning.

Xi Zeng, Haojie Ren, Ziying Song +2

School of Mechanical and Aerospace Engineering, Nanyang Technological University, Singapore · The University of Sheﬃeld, Sheﬃeld, United Kingdom · School of Artiﬁcial Intelligence (School of Software), Yanshan University, Qinhuangdao, China

Representation Learning World Model Learning World Model-Based Planning Joint-Embedding Predictive Architecture State Representation Learning

Ask PDF

Mar 13, 2026

LeWorldModel: Stable End-to-End Joint-Embedding Predictive Architecture from Pixels

Joint Embedding Predictive Architectures (JEPAs) offer a compelling framework for learning world models in compact latent spaces, yet existing methods remain fragile, relying on complex multi-term losses, exponential moving averages, pre-trained encoders, or auxiliary supervision to avoid representation collapse. In this work, we introduce LeWorldModel (LeWM), the first JEPA that trains stably end-to-end from raw pixels using only two loss terms: a next-embedding prediction loss and a regularizer enforcing Gaussian-distributed latent embeddings. This reduces tunable loss hyperparameters from six to one compared to the only existing end-to-end alternative. With ~15M parameters trainable on a single GPU in a few hours, LeWM plans up to 48x faster than foundation-model-based world models while remaining competitive across diverse 2D and 3D control tasks. Beyond control, we show that LeWM's latent space encodes meaningful physical structure through probing of physical quantities. Surprise evaluation confirms that the model reliably detects physically implausible events.

Lucas Maes, Quentin Le Lidec, Damien Scieur +2

Mila & Université de Montréal · New York University · Samsung SAIL +1

World Model Learning Joint-Embedding Predictive Architecture Latent World Models

Ask PDF

Aug 7, 2026

UniJEPA: A Unified Joint-Embedding Predictive Architecture for Task-Agnostic Visual World Modeling

Joint-Embedding Predictive Architectures (JEPAs) have emerged as a principled framework for self-supervised learning of world models in compact latent spaces, yet existing methods are fragmented: some predict masked parts of a single image in latent space (I-JEPA), others learn to predict global photometric transformations (Image World Models), while video-scale JEPAs predict future temporal states and are post-trained for action-conditioned planning (V-JEPA~2, DINO-World, DINO-WM). These objectives are treated as distinct recipes with separate encoders, predictors, and anti-collapse regularizers, hindering a single model from unifying image-level and video-level world modeling. We present UniJEPA, a unified JEPA that jointly learns photometric prediction (image-level transformations) and temporal prediction (video-level next-state dynamics) in one shared latent space. A single end-to-end objective, composed of a next-embedding prediction loss and a Gaussian regularizer, yields a provably anti-collapse encoder-predictor pair trainable from raw pixels without EMA, stop-gradient, or pre-trained encoders. We show that the same latent space supports controllable abstraction: photometric prediction learns invariant structure while temporal prediction learns equivariant dynamics. After action-conditioned post-training on offline trajectories, UniJEPA enables zero-shot planning by treating goal features as prediction targets. On image, video, and control benchmarks, UniJEPA matches or surpasses task-specific JEPAs while requiring a single loss hyperparameter, and plans up to tens of times faster than generative world models at comparable accuracy.

An Lanji, Dawei Liu, Jin Li +3

University of Electronic Science and Technology of China, Chengdu, China.

World Model Learning Joint-Embedding Predictive Architecture Model-Based Planning Predictive Representation Learning Self-Supervised Visual Representation Learning Video World Models

Ask PDF

Abstract

Joint-Embedding Predictive Architecture

Explore similar work

PhyLatent: Learning Dynamics-Relevant Representations for JEPA World Models

LeWorldModel: Stable End-to-End Joint-Embedding Predictive Architecture from Pixels

UniJEPA: A Unified Joint-Embedding Predictive Architecture for Task-Agnostic Visual World Modeling