No Gaussian Required: Contrastive Inverse Dynamics for JEPA World Models

# No Gaussian Required: Contrastive Inverse Dynamics for JEPA World Models
Source: [https://arxiv.org/html/2608.17542](https://arxiv.org/html/2608.17542)
###### 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\([13](https://arxiv.org/html/2608.17542#bib.bib1);[2](https://arxiv.org/html/2608.17542#bib.bib3);[5](https://arxiv.org/html/2608.17542#bib.bib2);[4](https://arxiv.org/html/2608.17542#bib.bib4)\)\. 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 reaches80\.0±2\.0%80\.0\{\\pm\}2\.0\\%success versus58\.0±2\.0%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](https://github.com/jackboyla/action-contrastive-jepa)\.

## 1Introduction

A useful world model should learn from experience, predict the consequences of actions, and support planning without reward labels or hand\-designed state\. Joint\-Embedding Predictive Architectures \(JEPAs\) offer an appealing route: an encoder maps sensory observations to a compact latent, and a predictor models future latents conditioned on actions\([13](https://arxiv.org/html/2608.17542#bib.bib1);[24](https://arxiv.org/html/2608.17542#bib.bib5);[19](https://arxiv.org/html/2608.17542#bib.bib6)\)\. Because prediction happens in representation space, the model need not reconstruct pixels\. Our experiments use images, but the objective itself only requires paired observations and actions\.

![Refer to caption](https://arxiv.org/html/2608.17542v1/images/hero_figure_nce.png)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\.The central difficulty is*collapse*: if every observation maps to the same latent, next\-latent prediction is trivially perfect but the representation is useless for planning\. Existing systems prevent collapse by constraining the representation with pretrained encoders, stop\-gradient and teacher–student heuristics, or explicit distributional regularizers\([2](https://arxiv.org/html/2608.17542#bib.bib3);[5](https://arxiv.org/html/2608.17542#bib.bib2);[4](https://arxiv.org/html/2608.17542#bib.bib4)\)\. LeWM\([14](https://arxiv.org/html/2608.17542#bib.bib8)\)distills this line to its simplest form: an end\-to\-end pixel JEPA trained with next\-embedding prediction plus SIGReg, a regularizer from LeJEPA\([3](https://arxiv.org/html/2608.17542#bib.bib7)\)that projects embeddings onto random directions and pushes every one\-dimensional marginal toward a Gaussian\. This reduces the multi\-term PLDM recipe\([19](https://arxiv.org/html/2608.17542#bib.bib6)\)to a single effective coefficient\.

SIGReg is principled, but it stabilizes the world model by prescribing a*global, isotropic\-Gaussian*latent geometry that the encoder must satisfy regardless of the environment being modeled\.[14](https://arxiv.org/html/2608.17542#bib.bib8)themselves flag this prescription as a possible cause of weak LeWM performance in low\-intrinsic\-dimensionality environments\. This motivates our question:*rather than imposing a distribution the dynamics never asked for, can anti\-collapse be derived from the transition data itself?*

Sensorimotor World Models \(SMWM\) answers this question with a standalone inverse\-action MSE regularizer in the same LeWM setting\([11](https://arxiv.org/html/2608.17542#bib.bib9)\)\. Our focus is the contrastive formulation: a chance\-level collapse bound, its reliability relative to inverse MSE, one loss weight held fixed across tasks, and a harder Scene stress test\.

We study Action\-Contrastive Masked Transition Modeling \(AC\-MTM\)\. As in LeWM, the model predicts the future latent from the current latent and action\. During training only, an inverse head additionally maps each adjacent latent pair to an action query that must identify the observed action among the other actions in the batch, an InfoNCE\-style objective\([7](https://arxiv.org/html/2608.17542#bib.bib13);[20](https://arxiv.org/html/2608.17542#bib.bib14)\)we call Action\-NCE\. A constant encoder maps every transition to the same query, so action identification cannot beat chance \(Section[3](https://arxiv.org/html/2608.17542#S3)\): collapse turns the forward objective’s global optimum into a representation that cannot outperform chance on the inverse task\. The resulting anti\-collapse pressure is dynamics\-native; the encoder is forced to keep exactly the information that distinguishes the effects of actions\. The entire inverse branch is discarded after training, leaving the deployed model identical to LeWM\. We evaluate AC\-MTM as a controlled LeWM modification against SIGReg, a non\-contrastive inverse\-regression ablation \(MTM\-MSE\), and a contrastive forward\-prediction alternative \(AC\-CPC\)\. Three findings emerge:

\(1\) Dynamics\-derived anti\-collapse works, with no distributional prior and no test\-time change\.Trained from scratch under LeWM’s one\-stage recipe, AC\-MTM never collapses across tasks and seeds and matches SIGReg on the standard four\-task suite \(TwoRoom, Reacher, PushT, OGBench\-Cube\), winning two tasks, tying one, and losing PushT\. The planner never calls the inverse head: evaluation uses the same encoder, autoregressive predictor, latent distance, and CEM planner as LeWM\.

\(2\) On the harder Scene task, results are consistent with the prescribed geometry becoming a bottleneck\.On OGBench Visual Scene, where a single arm controls a drawer, window, buttons, and a movable cube, SIGReg drops to a three\-seed mean of 58\.0% under the matched trajectory\-goal MPC protocol, while AC\-MTM remains at 80\.0%, with gains of 24, 20, and 22 points across the three training seeds\. A single 50\-episode random\-policy run gives a 52% baseline estimate\. The inverse\-regression ablation also clears SIGReg decisively \(75\.3%\), so the advantage belongs to the dynamics\-derived signal itself, not to one loss form\. Both models avoid collapse; the result may combine more useful transition geometry with more accurate short\-horizon latent dynamics\. Appendix[B](https://arxiv.org/html/2608.17542#A2)audits this result\.

\(3\) The contrastive form of the inverse signal is what makes it reliable\.Plain inverse regression \(MTM\-MSE\) is strong on TwoRoom, Cube, and Scene, but collapses on two of three Reacher seeds; Action\-NCE removes this bimodality at the cost of 3\.8 points on a long\-horizon stress test\. Both inverse variants trail SIGReg on PushT, where identifying the action underweights weakly controlled object state, a limitation we analyze in Section[4\.5](https://arxiv.org/html/2608.17542#S4.SS5)\.

Our contributions are: \(i\) AC\-MTM, a distribution\-free, dynamics\-native anti\-collapse mechanism for end\-to\-end JEPA world models, with a simple chance\-level lower bound showing why collapse cannot pay; \(ii\) a controlled three\-seed comparison on five pixel\-control tasks, matched down to the planner, isolating the training signal as the only moving part; and \(iii\) probe, stress\-test, and counterfactual\-surprise analyses that separate “does not collapse” from “supports planning,” including an explicit account of when inverse\-dynamics anti\-collapse should and should not be expected to hold\.

## 2Background and Positioning

#### Reward\-free latent planning\.

We consider offline trajectories of pixel observations and actionsτ=\(𝐨0,𝐚0,𝐨1,…\)\\tau=\(\\mathbf\{o\}\_\{0\},\\mathbf\{a\}\_\{0\},\\mathbf\{o\}\_\{1\},\\dots\), with no rewards or optimality assumptions\. During training,𝐚t\\mathbf\{a\}\_\{t\}is the recorded continuous command carrying𝐨t\\mathbf\{o\}\_\{t\}to𝐨t\+1\\mathbf\{o\}\_\{t\+1\}\. At test time, no action label is available: a model\-predictive\-control solver based on the cross\-entropy method \(CEM\) samples, scores, and refits a distribution over candidate action sequences to minimize latent distance to an encoded goal\([19](https://arxiv.org/html/2608.17542#bib.bib6);[14](https://arxiv.org/html/2608.17542#bib.bib8)\)\.

#### From PLDM to LeWM\.

PLDM prevents collapse with a VICReg\-derived objective of several interacting terms\([19](https://arxiv.org/html/2608.17542#bib.bib6);[5](https://arxiv.org/html/2608.17542#bib.bib2)\)\. LeWM replaces it with prediction plus SIGReg, matching all one\-dimensional latent marginals to a Gaussian \(motivated by the Cramér–Wold theorem\)\([3](https://arxiv.org/html/2608.17542#bib.bib7);[14](https://arxiv.org/html/2608.17542#bib.bib8)\)\. This simplification is the right comparison point for our work: the contribution of LeWM is not only performance, but the fact that stable end\-to\-end pixel JEPA world models can be made dramatically simpler than earlier recipes\. We keep that agenda while replacing the remaining global distribution\-matching term with a local, contrastive transition signal\.

#### A family of distribution\-free signals\.

MTM\-MSE is our SMWM\-style inverse\-regression baseline\([11](https://arxiv.org/html/2608.17542#bib.bib9)\): it regresses the observed action from an adjacent latent pair\. SMWM establishes this basic inverse\-MSE mechanism and tunes its loss weight per environment; our controlled baseline uses one coefficient chosen on Reacher and held fixed across tasks\. To separate “does not collapse” from “plans well,” we also evaluate AC\-CPC, an action\-conditioned contrastive*forward*objective that identifies the true future latent among batch negatives rather than identifying the action\. All variants are training signals plugged into the same LeWM encoder, predictor, and planner; Table[1](https://arxiv.org/html/2608.17542#S2.T1)summarizes what each signal asks of the representation\.

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\.

## 3Action\-Contrastive Masked Transition Modeling

Let latent embedding of observation𝐨t\\mathbf\{o\}\_\{t\}at timestepttbe𝐳t=encθ​\(𝐨t\)\\mathbf\{z\}\_\{t\}=\\mathrm\{enc\}\_\{\\theta\}\(\\mathbf\{o\}\_\{t\}\)\. We write a one\-step predictor for compactness; the implementation inherits LeWM’s causal latent\-history predictor\. The forward \(planning\) task is

𝐳^t\+1=fwdϕ​\(𝐳t,𝐚t\),ℒfwd=‖𝐳^t\+1−𝐳t\+1‖22\.\\hat\{\\mathbf\{z\}\}\_\{t\+1\}=\\mathrm\{fwd\}\_\{\\phi\}\(\\mathbf\{z\}\_\{t\},\\mathbf\{a\}\_\{t\}\),\\qquad\\mathcal\{L\}\_\{\\mathrm\{fwd\}\}=\\big\\\|\\hat\{\\mathbf\{z\}\}\_\{t\+1\}\-\\mathbf\{z\}\_\{t\+1\}\\big\\\|\_\{2\}^\{2\}\.\(1\)AC\-MTM adds a training\-only inverse\-dynamics head\. For each of theNNtransitions in a flattened batch/window, it produces an action query

𝐚^i=invψ​\(𝐳i,𝐳i\+1\)\.\\hat\{\\mathbf\{a\}\}\_\{i\}=\\mathrm\{inv\}\_\{\\psi\}\(\\mathbf\{z\}\_\{i\},\\mathbf\{z\}\_\{i\+1\}\)\.\(2\)We score every observed action𝐚j\\mathbf\{a\}\_\{j\}as a candidate for transitioniiusing negative squared distance, giving an InfoNCE\-style classification loss\([7](https://arxiv.org/html/2608.17542#bib.bib13);[20](https://arxiv.org/html/2608.17542#bib.bib14)\)over in\-batch actions:

si​j=−‖𝐚^i−𝐚j‖22τ​da,ℒNCE=−1N∑i=1Nlogexp⁡\(si​i\)∑j=1Nexp⁡\(si​j\),s\_\{ij\}=\-\\frac\{\\\|\\hat\{\\mathbf\{a\}\}\_\{i\}\-\\mathbf\{a\}\_\{j\}\\\|\_\{2\}^\{2\}\}\{\\tau d\_\{a\}\},\\qquad\\mathcal\{L\}\_\{\\mathrm\{NCE\}\}=\-\\frac\{1\}\{N\}\\sum\_\{i=1\}^\{N\}\\log\\frac\{\\exp\(s\_\{ii\}\)\}\{\\sum\_\{j=1\}^\{N\}\\exp\(s\_\{ij\}\)\},\(3\)wheredad\_\{a\}is the action dimension\. The total objective is

ℒA​C−M​T​M=ℒfwd\+λ​ℒNCE,λ=0\.30,τ=0\.10\.\\mathcal\{L\}\_\{AC\-MTM\}=\\mathcal\{L\}\_\{\\mathrm\{fwd\}\}\+\\lambda\\mathcal\{L\}\_\{\\mathrm\{NCE\}\},\\qquad\\lambda=0\.30,\\quad\\tau=0\.10\.\(4\)There is no SIGReg term\. The fixed coefficient was selected by a bounded Reacher stability sweep and then used unchanged across all tasks\.

“Masked” refers to factor prediction within a transition tuple, not image patch masking\. The forward task withholds𝐳t\+1\\mathbf\{z\}\_\{t\+1\}and predicts it from\(𝐳t,𝐚t\)\(\\mathbf\{z\}\_\{t\},\\mathbf\{a\}\_\{t\}\); the inverse task withholds𝐚t\\mathbf\{a\}\_\{t\}and predicts it from\(𝐳t,𝐳t\+1\)\(\\mathbf\{z\}\_\{t\},\\mathbf\{z\}\_\{t\+1\}\)\. We optimize both tasks on every batch rather than sample one mask at a time\.

The candidate set in Equation[3](https://arxiv.org/html/2608.17542#S3.E3)is theN=B⁡\(T−1\)N\{=\}B\(T\{\-\}1\)observed action blocks already present in the batch; any dataset action or a memory bank could serve instead, but we use in\-batch actions for several reasons\. First, the blocks are already resident on the GPU, so the negatives add only anN×NN\\times Ndistance matrix\. Second, the candidates𝐚j\\mathbf\{a\}\_\{j\}are*raw*actions, not encoder outputs, so they carry no gradient toencθ\\mathrm\{enc\}\_\{\\theta\}; enlarging the pool changes only the difficulty of the discrimination, not the gradient path\. Third, the objective is anti\-collapse, which the in\-batch pool already enforces through thelog⁡N\\log Nfloor of Equation[5](https://arxiv.org/html/2608.17542#S3.E5)\. Sampling globally would mostly add easy negatives \(blocks far from the prediction that contribute negligible gradient\) while raising the rate of false negatives, since control actions repeat \(near\-zero or saturated blocks\) and a duplicated “negative” penalizes a correct prediction\. A larger candidate set therefore adds compute and label noise to sharpen an action\-retrieval property the method does not require\.

#### Why Action\-NCE opposes collapse\.

Forward prediction alone cannot distinguish a useful representation from a constant one: ifencθ​\(𝐨\)=c\\mathrm\{enc\}\_\{\\theta\}\(\\mathbf\{o\}\)=cfor every observation andfwdϕ​\(c,𝐚\)=c\\mathrm\{fwd\}\_\{\\phi\}\(c,\\mathbf\{a\}\)=c, thenℒfwd=0\\mathcal\{L\}\_\{\\mathrm\{fwd\}\}=0\. Action\-NCE turns this degenerate solution into a failed classification problem\. Under collapse, every transition gives the inverse head the same input pair\(c,c\)\(c,c\), so every row of the action classifier is identical\. The model is then forced to assign one fixed probability vectorppto allNNpositives in the batch\. Since each candidate action is the correct label exactly once, the average loss satisfies

−1N∑i=1Nlogpi≥logN,\-\\frac\{1\}\{N\}\\sum\_\{i=1\}^\{N\}\\log p\_\{i\}\\geq\\log N,\(5\)with equality only at the chance classifier\. Thus collapse can drive the forward loss to zero but cannot drive the contrastive inverse loss below chance\. To improve Action\-NCE, the encoder must preserve transition information that makes the observed action more identifiable than the in\-batch alternatives\.

#### Why not use non\-contrastive inverse regression?

MTM\-MSE is the SMWM\-style distribution\-free replacement for SIGReg\([11](https://arxiv.org/html/2608.17542#bib.bib9)\): it replaces Equation[3](https://arxiv.org/html/2608.17542#S3.E3)withℒinv=‖𝐚^t−𝐚t‖22\\mathcal\{L\}\_\{\\mathrm\{inv\}\}=\\\|\\hat\{\\mathbf\{a\}\}\_\{t\}\-\\mathbf\{a\}\_\{t\}\\\|\_\{2\}^\{2\}\. It is an important ablation because it asks whether any inverse\-dynamics signal is enough\. MTM\-MSE is strong on TwoRoom, PushT, and Cube, but the instability appears on Reacher, where two of three seeds collapse\. The reason is that inverse regression supplies only a variance\-scale floor under complete collapse: its best collapsed prediction is the mean action, and the loss is the action variance\. When multiple action blocks can produce similar visual endpoints, that margin can be weak, while forward MSE still rewards shrinking the latent scale\. Action\-NCE keeps the same training\-only inverse head but changes the failure geometry to chance\-level identification, which gives a sharper barrier against constant latents\. Section[4\.4](https://arxiv.org/html/2608.17542#S4.SS4)analyzes this failure mode\.

#### Controlled implementation\.

We inherit the LeWM vision encoder, latent projector, causal forward predictor, action conditioning, offline datasets, and CEM planner\. The only added component is a small MLP inverse head used during training\. After training, the head is never called by model rollout or cost evaluation and can be removed without changing predictions\. Consequently AC\-MTM and LeWM use the same test\-time computation and planner\.

## 4Experiments

#### Evaluation Setup\.

We evaluate the LeWM continuous\-control suite\. This includes TwoRoom, Reacher, PushT, and OGBench\-Cube\. We add OGBench Visual Scene as a harder multi\-object manipulation study \(Figure[2](https://arxiv.org/html/2608.17542#S4.F2)\)\. All models train end\-to\-end from pixels in one stage\. We evaluate each policy with the same CEM/MPC planner used by[14](https://arxiv.org/html/2608.17542#bib.bib8)\. CEM samples 300 candidate action sequences, retains 30 elites, and refits the sampling distribution for 30 iterations\. In Figure[3](https://arxiv.org/html/2608.17542#S4.F3)we report AC\-MTM and MTM\-MSE against our matched SIGReg reproduction of LeWM, alongside paper\-reported PLDM, DINO\-WM, goal\-conditioned policy, and random baselines from[14](https://arxiv.org/html/2608.17542#bib.bib8); on OGBench\-Scene, where no external numbers exist, we evaluate the random policy ourselves under the identical protocol\. Our controlled claims use three training seeds\{3072,1,2\}\\\{3072,1,2\\\}, 200 evaluation episodes, evaluation seed 42, goal offset 25, and interaction budget 50\. TwoRoom\-long changes only the goal offset and budget to100/150100/150\. The OGBench\-Scene task uses the same trajectory\-goal MPC protocol as the other visual tasks, with 50 evaluation episodes per training seed\. Appendix[A](https://arxiv.org/html/2608.17542#A1)records the full protocol\.

![Refer to caption](https://arxiv.org/html/2608.17542v1/images/tworoom_still.png)![Refer to caption](https://arxiv.org/html/2608.17542v1/images/reacher_still.png)![Refer to caption](https://arxiv.org/html/2608.17542v1/images/pusht_still.png)![Refer to caption](https://arxiv.org/html/2608.17542v1/images/cube_still.png)![Refer to caption](https://arxiv.org/html/2608.17542v1/images/scene_still.png)TwoRoomReacherPushTOGBench\-CubeOGBench\-SceneFigure 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\.
### 4\.1Collapse sanity check

On TwoRoom, removing SIGReg without replacement \(NoReg\) drives forward loss to≈0\\approx 0in training diagnostics, the trivial constant\-latent solution described above\. Both MTM\-MSE and AC\-MTM inverse tasks prevent this failure while preserving the same one\-stage training recipe proposed by[14](https://arxiv.org/html/2608.17542#bib.bib8)\. This confirms that transition supervision, not merely removal of SIGReg, provides the anti\-collapse signal \(Table[2](https://arxiv.org/html/2608.17542#S4.T2)\)\.

Table 2:TwoRoom anti\-collapse ablation with 200 evaluation episodes per seed\. Values are mean±\\pmstandard deviation over three training seeds\. Plain next\-latent prediction reaches a trivial low\-loss solution and plans substantially worse\.
### 4\.2Planning results in the LeWM evaluation frame

Table[3](https://arxiv.org/html/2608.17542#S4.T3)gives the controlled AC\-MTM versus SIGReg comparison, while Figure[3](https://arxiv.org/html/2608.17542#S4.F3)restores the original evaluation style by including the external paper\-reported baselines\. The figure is contextual: PLDM, DINO\-WM, GCBC, GCiQL, and GCIVL are taken from the LeWM paper, while the statistical claims come from our matched SIGReg reruns\. AC\-MTM uses the same test\-time planner as LeWM and improves over SIGReg on TwoRoom, Cube, and Scene, matches it on Reacher, and trails it on PushT\. The PushT deficit is analyzed in Section[4\.5](https://arxiv.org/html/2608.17542#S4.SS5); the full MTM\-MSE ablation is reported in Appendix[C](https://arxiv.org/html/2608.17542#A3)\.

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±\\pmstandard deviation over three training seeds; bold marks differences larger than the cross\-seed noise\. Figure[3](https://arxiv.org/html/2608.17542#S4.F3)places these numbers next to external paper\-reported baselines\.![Refer to caption](https://arxiv.org/html/2608.17542v1/images/results_lewm_style.png)Figure 3:Planning success across the five environments\. Salmon and blue bars are paper\-reported LeWM/external baselines from[14](https://arxiv.org/html/2608.17542#bib.bib8); 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](https://arxiv.org/html/2608.17542#S4.T3)\.
### 4\.3Harder OGBench Scene separates the methods

The largest margin appears on the more complex OGBench Visual Scene task \(Table[4](https://arxiv.org/html/2608.17542#S4.T4)\)\. Scene is qualitatively different from OGBench\-Cube: the same arm must control a multi\-object visual state involving a drawer, window, buttons, and movable object, so a useful latent must preserve several slow\-changing task variables at once\. Under the trajectory\-goal MPC protocol, SIGReg falls to a three\-seed mean of 58\.0% success while AC\-MTM achieves 80\.0%\. The random\-policy run helps calibrate these numbers: like OGBench\-Cube \(48% random\), the 25\-step trajectory\-goal protocol leaves many episodes near\-solved at reset\. One 50\-episode random\-policy run scores 52\.0% \(Wilson 95% interval, 38\.5 to 65\.2%\), so we treat it as a baseline estimate\. The paired comparison pooled over the three repeated sets of 50 start–goal tasks gives 40 AC\-MTM\-only successes versus 7 SIGReg\-only successes\. Notably, SIGReg does not collapse on Scene \(its final latent scale is healthy; Appendix[B](https://arxiv.org/html/2608.17542#A2)\); the gap is between avoiding collapse and learning a geometry the planner can exploit\.

The MTM\-MSE ablation localizes the credit\. Trained on Scene under the identical protocol, it reaches75\.3±2\.3%75\.3\{\\pm\}2\.3\\%with no collapsed seed and beats SIGReg in the paired comparison \(39 wins vs\. 13 losses,p≈4×10−4p\\\!\\approx\\\!4\{\\times\}10^\{\-4\}in an episode\-level descriptive test\)\. The Scene advantage therefore belongs to the inverse\-dynamics*family*: deriving anti\-collapse from the transitions, in either form, beats prescribing the latent distribution once several controllable factors must coexist in the latent\. The contrastive form adds 4\.7 points over inverse regression on Scene \(16 wins vs\. 9 losses,p≈0\.23p\\\!\\approx\\\!0\.23in the same descriptive test\) on top of the Reacher reliability that motivates it \(Section[4\.4](https://arxiv.org/html/2608.17542#S4.SS4)\)\.

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/McNemarp≈1\.1×10−6p\\\!\\approx\\\!1\.1\{\\times\}10^\{\-6\}for AC\-MTM \(40/7\) andp≈4\.1×10−4p\\\!\\approx\\\!4\.1\{\\times\}10^\{\-4\}for MTM\-MSE \(39/13\); the AC\-MTM–MTM\-MSE pairing is 16/9 \(p≈0\.23p\\\!\\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](https://arxiv.org/html/2608.17542#A2)audits possible explanations for the large margin\.
### 4\.4Reliability and stress\-test trade\-offs

The Reacher seed breakdown makes the stability difference explicit \(Table[5](https://arxiv.org/html/2608.17542#S4.T5)\)\. One MTM\-MSE run is competitive, but two enter the constant\-latent state\. In contrast, all AC\-MTM runs remain useful\. On TwoRoom\-long the non\-contrastive inverse\-regression ablation is stronger:28\.0±0\.8%28\.0\{\\pm\}0\.8\\%versus24\.2±0\.6%24\.2\{\\pm\}0\.6\\%for AC\-MTM \(Table[6](https://arxiv.org/html/2608.17542#S4.T6)\)\. Both MTM\-MSE and Action\-NCE remain above SIGReg on this long\-goal stress test, but Action\-NCE trades 3\.8 points of the pure\-MSE long\-horizon gain for much better Reacher reliability\.

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\.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\.
### 4\.5PushT probes diagnose the remaining family\-level gap

Linear probes freeze the encoder, embed 4000 frames, and fit ridge regression \(α=1\\alpha\{=\}1\) from the latent to each privileged simulator\-state coordinate\. Both inverse variants trail SIGReg on PushT, and the probes show the same family\-level failure mode for MTM\-MSE and AC\-MTM \(Table[7](https://arxiv.org/html/2608.17542#S4.T7)\): both preserve agent position, retain much of block position, and underencode the T\-block orientation, a weakly controlled variable that determines contact geometry\.

Table 7:PushT anti\-collapse mechanisms: frozen\-latent probes \(n=4000n\{=\}4000,α=1\\alpha\{=\}1; orientation isstate⁡\[4\]\\mathrm\{state\}\[4\]\) and 200\-episode planning\. AC\-MTM probe values are means over three training seeds; all rows use the same autoregressive planner\.

## 5Analysis

#### Why Action\-NCE stabilizes Reacher\.

Inverse\-MSE only resists collapse to the extent that the observed action can be regressed below the mean\-action baseline\. Reacher applies a multi\-step action block to a visually observed two\-link arm; many action sequences can produce similar endpoints, so the regression margin can be weak\. At the same time, forward latent MSE rewards shrinking the online representation\. Action\-NCE changes the failure geometry: even when absolute action regression is ambiguous, identical transition queries cannot assign different positives to different rows and remain at chance by Equation[5](https://arxiv.org/html/2608.17542#S3.E5)\. This does not make the inverse problem fully identifiable, but it provides a stronger early barrier against the constant\-latent attractor\.

#### What Action\-NCE does not solve\.

The contrastive loss still emphasizes state that separates actions\. It does not unconditionally preserve every task\-relevant variable, unlike SIGReg’s global variance pressure\. PushT requires accurate block pose and orientation even when the same commanded action has little visible effect before contact\. This is a plausible reason both inverse variants remain below SIGReg\. The remaining problem is representation coverage, not test\-time action selection: every method uses the same CEM planner\.

#### Why probes are diagnostic, not decisive\.

Linear probes measure whether a quantity is linearly decodable from the latent; planning measures whether the forward model and latent metric place goal\-reaching trajectories where CEM can find them\. These need not align\. MTM\-MSE and AC\-MTM under\-decode block orientation yet still support substantial planning success, while AC\-CPC decodes orientation better but plans worse\. Representation audits and planner behavior should therefore be reported together: a probe can identify a missing state factor, but it is not a substitute for closed\-loop planning\.

#### Training\-only means test\-time equivalence\.

The inverse head affects the encoder during optimization, but it is not queried by CEM and does not propose or score actions at test time\. Candidate actions come entirely from CEM, as in LeWM\. Thus the empirical differences isolate training\-time representation learning; they do not arise from a stronger inference procedure or extra runtime computation\.

### 5\.1Does the model notice impossible transitions?

Goal\-reaching success does not by itself show that a latent world model assigns meaningful surprise to physically inconsistent events\. We therefore add a model\-facing diagnostic on PushT and OGBench\-Cube\. For each sampled trajectory clip we first record the model’s*normal*one\-step prediction error: the latent\-space MSE between its prediction and the encoding of the frame that actually followed\. We then corrupt the same clip in one of two ways and measure the error again: \(i\) an*action counterfactual*keeps the latent history but swaps in an action block from a different clip; the model should no longer expect the observed outcome; and \(ii\) a*state discontinuity*keeps the model’s prediction but replaces the target future with latents from a different clip; an impossible jump in the environment state\. The reported*surprise ratio*is the per\-clip corrupted error divided by the normal error\. A ratio of1×1\{\\times\}means the model finds impossible transitions exactly as predictable as real ones, i\.e\. its dynamics carry no physical content; higher is better\. Both tests use the trained encoder, action encoder, and autoregressive predictor used by MPC; no decoder or privileged state is introduced\.

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×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±\\pmstandard deviation over seeds\{3072,1,2\}\\\{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\)\.Table[8](https://arxiv.org/html/2608.17542#S5.T8)shows that both objectives learn action\-conditioned dynamics with strong internal violation signals: every model flags both corruption types with errors one to three orders of magnitude above its normal prediction error, and does so on essentially every clip\. Two caveats guide the reading\. First, because each ratio is normalized by that model’s own normal\-transition error, absolute magnitudes are not comparable across models: a model with near\-zero baseline error can post an enormous ratio without having better dynamics\. The meaningful statements are within\-model \(SIGReg and AC\-MTM both separate violations from real transitions by large margins\) and qualitative \(neither model is fooled\)\. Second, the smaller AC\-MTM margins on PushT are consistent with its weaker block\-state probes and lower PushT planning success \(Section[4\.5](https://arxiv.org/html/2608.17542#S4.SS5)\), so the diagnostic tracks the planning evidence rather than contradicting it\. These results support a narrower claim than human\-style physical reasoning: the learned latent dynamics are sensitive to impossible action\-conditioned transitions in the same representation used for planning\.

## 6Related Work

#### World models and planning\.

The idea of learning a predictive model and planning through it goes back at least to differentiable recurrent world models for control\([17](https://arxiv.org/html/2608.17542#bib.bib10)\)and was later popularized in compact latent form by[8](https://arxiv.org/html/2608.17542#bib.bib11)\. Reward\-driven latent world models such as Dreamer\([9](https://arxiv.org/html/2608.17542#bib.bib19)\)and TD\-MPC2\([10](https://arxiv.org/html/2608.17542#bib.bib20)\)shape their representations with reconstruction or value signals; our setting removes both\. We follow the reward\-free, goal\-conditioned variant: learn from offline state–action trajectories without task rewards, then plan to a goal image at test time, which leaves anti\-collapse as the central representational problem\.

#### Reconstruction\-free latent planning\.

PLDM shows that a JEPA\-style latent dynamics model can plan from reward\-free offline data without reconstructing pixels, but relies on a VICReg\-derived multi\-term anti\-collapse recipe and inverse\-dynamics regularization\([19](https://arxiv.org/html/2608.17542#bib.bib6);[5](https://arxiv.org/html/2608.17542#bib.bib2)\)\. DINO\-WM instead uses frozen DINOv2 features and learns a latent dynamics model on top, trading end\-to\-end simplicity for strong pretrained visual representations\([24](https://arxiv.org/html/2608.17542#bib.bib5)\)\. LeWM is the closest predecessor to our work: it trains from pixels end\-to\-end with a two\-term objective, replacing PLDM’s heavier recipe with next\-latent prediction plus SIGReg\([14](https://arxiv.org/html/2608.17542#bib.bib8);[3](https://arxiv.org/html/2608.17542#bib.bib7)\)\. We share LeWM’s simplification agenda but remove its remaining global distributional assumption during training while leaving its deployed model unchanged\.

SMWM is the most direct prior work: it uses the same basic LeWM encoder, forward latent predictor, and training\-only inverse head, with inverse\-action MSE as the sole anti\-collapse term\([11](https://arxiv.org/html/2608.17542#bib.bib9)\)\. It evaluates the same four standard environments and finds that inverse MSE roughly matches SIGReg on the 2D tasks and improves on Cube\. We therefore treat MTM\-MSE as an SMWM\-style baseline, not a new mechanism\. Our contribution is the contrastive Action\-NCE form, its chance\-level collapse bound and Reacher reliability, a single coefficient held fixed across tasks rather than tuned per environment, and the Scene and stress\-test analysis\.

#### Anti\-collapse objectives for JEPAs\.

I\-JEPA and V\-JEPA stabilize prediction with target encoders, stop\-gradient, and EMA updates\([2](https://arxiv.org/html/2608.17542#bib.bib3);[4](https://arxiv.org/html/2608.17542#bib.bib4);[6](https://arxiv.org/html/2608.17542#bib.bib21)\)\. VICReg prevents collapse by variance and covariance constraints\([5](https://arxiv.org/html/2608.17542#bib.bib2)\)\. LeJEPA derives a leaner alternative, SIGReg, by regularizing embeddings toward an isotropic Gaussian through random one\-dimensional projections\([3](https://arxiv.org/html/2608.17542#bib.bib7)\)\. These are representation\-distribution constraints: they state what the latent cloud should look like\. MTM\-MSE and AC\-MTM instead state what adjacent latents must*explain*: the action that connected them\. AC\-MTM sharpens this requirement from absolute regression to contrastive identification\. The signal is local, transition\-level, and dynamics\-native\.

#### Benchmarks and goal\-conditioned baselines\.

The LeWM suite combines top\-down navigation, DMC\-style reaching, PushT contact manipulation, and OGBench\-Cube\([15](https://arxiv.org/html/2608.17542#bib.bib12)\)\. The external baselines in Figure[3](https://arxiv.org/html/2608.17542#S4.F3)include PLDM, DINO\-WM, goal\-conditioned behavioral cloning, goal\-conditioned offline RL, and random policies as reported by[14](https://arxiv.org/html/2608.17542#bib.bib8)\. They place AC\-MTM in the same visual frame as the original LeWM evaluation; our controlled claims use the released LeWM datasets, architectures, and planner settings so that the comparison changes the training signal rather than the downstream policy class\.

#### Inverse dynamics and probing\.

Inverse dynamics has long been used to learn controllable representations: as a self\-supervised feature signal for manipulation\([1](https://arxiv.org/html/2608.17542#bib.bib15)\), to build the embedding underlying curiosity rewards\([16](https://arxiv.org/html/2608.17542#bib.bib16)\), as an auxiliary loss for RL\([18](https://arxiv.org/html/2608.17542#bib.bib17)\), and, in multi\-step form, with guarantees that it recovers exactly the control\-endogenous part of the state\([12](https://arxiv.org/html/2608.17542#bib.bib18)\)\. PLDM includes an inverse term inside a larger multi\-term objective\([19](https://arxiv.org/html/2608.17542#bib.bib6)\)\. Our use is narrower and sharper: a training\-only anti\-collapse head attached to an otherwise unchanged LeWM planner, with no privileged state and no planner\-time inverse model, and \(unlike ordinary inverse regression\) trained to*identify*the positive action among in\-batch continuous\-action negatives\. Linear probing of physical quantities is increasingly used to interpret latent world models\([14](https://arxiv.org/html/2608.17542#bib.bib8)\); our results caution against inferring planning quality from such probes alone\.

#### Action\-driven contrastive representation learning\.

CLOUD learns forward and inverse dynamics through contrastive estimation for planning and imitation\([21](https://arxiv.org/html/2608.17542#bib.bib22)\)\.[22](https://arxiv.org/html/2608.17542#bib.bib23)use InfoNCE between visual representations and observed actions to retain controllable factors, while also using reward prediction and temporal coherence for online visual RL\. TACO instead contrasts current\-state and action sequences against future\-state representations for online and offline visual RL\([23](https://arxiv.org/html/2608.17542#bib.bib24)\)\. These works establish action\-driven contrast as a useful representation signal\. We use observed continuous actions as the candidate labels of a training\-only inverse task and study that task as the sole anti\-collapse term in an end\-to\-end, reward\-free JEPA world model\.

## 7Limitations

AC\-MTM provides no unconditional geometric guarantee\. It requires variation in the action candidates and useful visual evidence about action effects; unobserved actuators, stochastic dynamics, no\-op\-heavy datasets, or many duplicate actions can weaken the contrastive task\. We evaluate normalized continuous controls only\. Discrete, hybrid, structured, and very high\-dimensional action spaces may require different scores, learned action embeddings, or hard\-negative sampling\. The loss also introduces training\-time dependence on batch composition, the temperatureτ\\tau, and one coefficientλ\\lambda; although one setting works across our suite, broader scaling evidence is needed\. PushT shows the flip side of the dynamics\-native pressure: a variable that actions barely move can be underweighted even when it matters for the task, and exploratory single\-seed runs on additional OGBench families \(Appendix[D](https://arxiv.org/html/2608.17542#A4)\) suggest the same trade\-off on combinatorial puzzle tasks\.

AC\-MTM improves several goal\-conditioned planning protocols and shows strong latent surprise on counterfactual\-action and state\-discontinuity probes, but these diagnostics remain internal to the trained latent dynamics\. They do not replace public violation\-of\-expectation suites with human\-designed physical events or prove robust out\-of\-distribution physical reasoning\. OGBench Visual Scene is also evaluated under our trajectory\-goal protocol rather than the official five\-task fixed\-goal OGBench leaderboard protocol, so its 80% success rate should be read as a matched stress test, not a public benchmark score; under the official 750\-step fixed\-goal protocol, neither method solves the tasks at this model scale \(Appendix[B](https://arxiv.org/html/2608.17542#A2)\)\.

## 8Conclusion

A world model’s representation has to be protected from collapse, but nothing about the problem says that protection must take the form of a prescribed latent distribution\. AC\-MTM replaces LeWM’s Gaussian matching with a contrastive, dynamics\-native signal: keep whatever information identifies the action that drove each transition\. The result is still a one\-stage, end\-to\-end pixel JEPA without target network, stop\-gradient branch, frozen encoder, reconstruction decoder, or fixed global latent distribution, and the deployed model is unchanged\. The empirical picture is not universal dominance: PushT exposes a real limitation of inverse\-action signals when task\-relevant state is weakly controlled\. Where the environment couples many controllable factors, our results are consistent with prescribed geometry becoming a bottleneck: on OGBench Visual Scene, SIGReg reaches 58\.0% while AC\-MTM reaches 80\.0% across matched seeds, though the Scene gain may also reflect lower one\-step forward error\. On TwoRoom\-long AC\-MTM retains most of the inverse\-dynamics gain over SIGReg\. These results show that transition\-derived anti\-collapse can replace a fixed global latent prior while preserving planning quality\.

## References

- Agrawalet al\.\(2016\)P\. Agrawal, A\. Nair, P\. Abbeel, J\. Malik, and S\. LevineLearning to poke by poking: experiential learning of intuitive physics\.InAdvances in Neural Information Processing Systems,Cited by:[§6](https://arxiv.org/html/2608.17542#S6.SS0.SSS0.Px5.p1.1)\.
- Assranet al\.\(2023\)M\. Assran, Q\. Duval, I\. Misra, P\. Bojanowski, P\. Vincent, M\. Rabbat, Y\. LeCun, and N\. BallasSelf\-supervised learning from images with a joint\-embedding predictive architecture\.InIEEE/CVF Conference on Computer Vision and Pattern Recognition \(CVPR\),Cited by:[§1](https://arxiv.org/html/2608.17542#S1.p2.1),[§6](https://arxiv.org/html/2608.17542#S6.SS0.SSS0.Px3.p1.1),[Abstract](https://arxiv.org/html/2608.17542#abstract1.1)\.
- Balestriero and LeCun \(2025\)R\. Balestriero and Y\. LeCunLeJEPA: provable and scalable self\-supervised learning without the heuristics\.arXiv preprint arXiv:2511\.08544\.Cited by:[§1](https://arxiv.org/html/2608.17542#S1.p2.1),[§2](https://arxiv.org/html/2608.17542#S2.SS0.SSS0.Px2.p1.1),[§6](https://arxiv.org/html/2608.17542#S6.SS0.SSS0.Px2.p1.1),[§6](https://arxiv.org/html/2608.17542#S6.SS0.SSS0.Px3.p1.1)\.
- Bardeset al\.\(2024\)A\. Bardes, Q\. Garrido, J\. Ponce, X\. Chen, M\. Rabbat, Y\. LeCun, M\. Assran, and N\. BallasRevisiting feature prediction for learning visual representations from video\.Transactions on Machine Learning Research\.Cited by:[§1](https://arxiv.org/html/2608.17542#S1.p2.1),[§6](https://arxiv.org/html/2608.17542#S6.SS0.SSS0.Px3.p1.1),[Abstract](https://arxiv.org/html/2608.17542#abstract1.1)\.
- Bardeset al\.\(2022\)A\. Bardes, J\. Ponce, and Y\. LeCunVICReg: variance\-invariance\-covariance regularization for self\-supervised learning\.InInternational Conference on Learning Representations \(ICLR\),Cited by:[§1](https://arxiv.org/html/2608.17542#S1.p2.1),[§2](https://arxiv.org/html/2608.17542#S2.SS0.SSS0.Px2.p1.1),[§6](https://arxiv.org/html/2608.17542#S6.SS0.SSS0.Px2.p1.1),[§6](https://arxiv.org/html/2608.17542#S6.SS0.SSS0.Px3.p1.1),[Abstract](https://arxiv.org/html/2608.17542#abstract1.1)\.
- Grillet al\.\(2020\)J\. Grill, F\. Strub, F\. Altché, C\. Tallec, P\. H\. Richemond, E\. Buchatskaya, C\. Doersch, B\. A\. Pires, Z\. D\. Guo, M\. G\. Azar, B\. Piot, K\. Kavukcuoglu, R\. Munos, and M\. ValkoBootstrap your own latent: a new approach to self\-supervised learning\.InAdvances in Neural Information Processing Systems,Cited by:[§6](https://arxiv.org/html/2608.17542#S6.SS0.SSS0.Px3.p1.1)\.
- Gutmann and Hyvärinen \(2010\)M\. Gutmann and A\. HyvärinenNoise\-contrastive estimation: a new estimation principle for unnormalized statistical models\.InInternational Conference on Artificial Intelligence and Statistics \(AISTATS\),Cited by:[§1](https://arxiv.org/html/2608.17542#S1.p5.1),[§3](https://arxiv.org/html/2608.17542#S3.p1.3)\.
- Ha and Schmidhuber \(2018\)D\. Ha and J\. SchmidhuberRecurrent world models facilitate policy evolution\.InAdvances in Neural Information Processing Systems,Cited by:[§6](https://arxiv.org/html/2608.17542#S6.SS0.SSS0.Px1.p1.1)\.
- Hafneret al\.\(2020\)D\. Hafner, T\. Lillicrap, J\. Ba, and M\. NorouziDream to control: learning behaviors by latent imagination\.InInternational Conference on Learning Representations \(ICLR\),Cited by:[§6](https://arxiv.org/html/2608.17542#S6.SS0.SSS0.Px1.p1.1)\.
- Hansenet al\.\(2024\)N\. Hansen, H\. Su, and X\. WangTD\-MPC2: scalable, robust world models for continuous control\.InInternational Conference on Learning Representations \(ICLR\),Cited by:[§6](https://arxiv.org/html/2608.17542#S6.SS0.SSS0.Px1.p1.1)\.
- Ivashkovet al\.\(2026\)P\. Ivashkov, R\. Balestriero, and B\. SchölkopfSensorimotor world models: perception for action via inverse dynamics\.arXiv preprint arXiv:2606\.20104\.Cited by:[§1](https://arxiv.org/html/2608.17542#S1.p4.1),[§2](https://arxiv.org/html/2608.17542#S2.SS0.SSS0.Px3.p1.1),[Table 1](https://arxiv.org/html/2608.17542#S2.T1.2.4.1),[§3](https://arxiv.org/html/2608.17542#S3.SS0.SSS0.Px2.p1.1),[§6](https://arxiv.org/html/2608.17542#S6.SS0.SSS0.Px2.p2.1)\.
- Lambet al\.\(2023\)A\. Lamb, R\. Islam, Y\. Efroni, A\. Didolkar, D\. Misra, D\. Foster, L\. Molu, R\. Chari, A\. Krishnamurthy, and J\. LangfordGuaranteed discovery of control\-endogenous latent states with multi\-step inverse models\.Transactions on Machine Learning Research\.Cited by:[§6](https://arxiv.org/html/2608.17542#S6.SS0.SSS0.Px5.p1.1)\.
- LeCun \(2022\)Y\. LeCunA path towards autonomous machine intelligence\.OpenReview\.Cited by:[§1](https://arxiv.org/html/2608.17542#S1.p1.1),[Abstract](https://arxiv.org/html/2608.17542#abstract1.1)\.
- Maeset al\.\(2026\)L\. Maes, Q\. Le Lidec, D\. Scieur, Y\. LeCun, and R\. BalestrieroLeWorldModel: stable end\-to\-end joint\-embedding predictive architecture from pixels\.arXiv preprint arXiv:2603\.19312\.Cited by:[Appendix A](https://arxiv.org/html/2608.17542#A1.p1.1),[§1](https://arxiv.org/html/2608.17542#S1.p2.1),[§1](https://arxiv.org/html/2608.17542#S1.p3.1),[§2](https://arxiv.org/html/2608.17542#S2.SS0.SSS0.Px1.p1.1),[§2](https://arxiv.org/html/2608.17542#S2.SS0.SSS0.Px2.p1.1),[Table 1](https://arxiv.org/html/2608.17542#S2.T1.2.3.1),[Figure 3](https://arxiv.org/html/2608.17542#S4.F3),[§4](https://arxiv.org/html/2608.17542#S4.SS0.SSS0.Px1.p1.1),[§4\.1](https://arxiv.org/html/2608.17542#S4.SS1.p1.1),[§6](https://arxiv.org/html/2608.17542#S6.SS0.SSS0.Px2.p1.1),[§6](https://arxiv.org/html/2608.17542#S6.SS0.SSS0.Px4.p1.1),[§6](https://arxiv.org/html/2608.17542#S6.SS0.SSS0.Px5.p1.1)\.
- Parket al\.\(2025\)S\. Park, K\. Frans, B\. Eysenbach, and S\. LevineOGBench: benchmarking offline goal\-conditioned reinforcement learning\.InInternational Conference on Learning Representations,Cited by:[§6](https://arxiv.org/html/2608.17542#S6.SS0.SSS0.Px4.p1.1)\.
- Pathaket al\.\(2017\)D\. Pathak, P\. Agrawal, A\. A\. Efros, and T\. DarrellCuriosity\-driven exploration by self\-supervised prediction\.InInternational Conference on Machine Learning \(ICML\),Cited by:[§6](https://arxiv.org/html/2608.17542#S6.SS0.SSS0.Px5.p1.1)\.
- Schmidhuber \(1990\)J\. SchmidhuberMaking the world differentiable: on using self\-supervised fully recurrent neural networks for dynamic reinforcement learning and planning in non\-stationary environments\.Technical reportTechnical ReportFKI\-126\-90,Institut für Informatik, Technische Universität München\.Cited by:[§6](https://arxiv.org/html/2608.17542#S6.SS0.SSS0.Px1.p1.1)\.
- Shelhameret al\.\(2016\)E\. Shelhamer, P\. Mahmoudieh, M\. Argus, and T\. DarrellLoss is its own reward: self\-supervision for reinforcement learning\.arXiv preprint arXiv:1612\.07307\.Cited by:[§6](https://arxiv.org/html/2608.17542#S6.SS0.SSS0.Px5.p1.1)\.
- Sobalet al\.\(2025\)U\. Sobal, W\. Zhang, K\. Cho, R\. Balestriero, T\. G\. J\. Rudner, and Y\. LeCunLearning from reward\-free offline data: a case for planning with latent dynamics models\.InAdvances in Neural Information Processing Systems,Vol\.38\.Cited by:[§1](https://arxiv.org/html/2608.17542#S1.p1.1),[§1](https://arxiv.org/html/2608.17542#S1.p2.1),[§2](https://arxiv.org/html/2608.17542#S2.SS0.SSS0.Px1.p1.1),[§2](https://arxiv.org/html/2608.17542#S2.SS0.SSS0.Px2.p1.1),[Table 1](https://arxiv.org/html/2608.17542#S2.T1.2.2.1),[§6](https://arxiv.org/html/2608.17542#S6.SS0.SSS0.Px2.p1.1),[§6](https://arxiv.org/html/2608.17542#S6.SS0.SSS0.Px5.p1.1)\.
- van den Oordet al\.\(2018\)A\. van den Oord, Y\. Li, and O\. VinyalsRepresentation learning with contrastive predictive coding\.arXiv preprint arXiv:1807\.03748\.Cited by:[§1](https://arxiv.org/html/2608.17542#S1.p5.1),[§3](https://arxiv.org/html/2608.17542#S3.p1.3)\.
- Wanget al\.\(2021\)J\. Wang, Y\. Lu, and H\. ZhaoCLOUD: contrastive learning of unsupervised dynamics\.InProceedings of the 2020 Conference on Robot Learning,Proceedings of Machine Learning Research, Vol\.155,pp\. 365–376\.Cited by:[§6](https://arxiv.org/html/2608.17542#S6.SS0.SSS0.Px6.p1.1)\.
- Yuanet al\.\(2024\)L\. Yuan, X\. Lu, and Y\. LiuLearning task\-relevant representations via rewards and real actions for reinforcement learning\.Knowledge\-Based Systems294,pp\. 111788\.External Links:[Document](https://dx.doi.org/10.1016/j.knosys.2024.111788)Cited by:[§6](https://arxiv.org/html/2608.17542#S6.SS0.SSS0.Px6.p1.1)\.
- Zhenget al\.\(2023\)R\. Zheng, X\. Wang, Y\. Sun, S\. Ma, J\. Zhao, H\. Xu, H\. Daumé III, and F\. HuangTACO: temporal latent action\-driven contrastive loss for visual reinforcement learning\.InAdvances in Neural Information Processing Systems,Vol\.36\.Cited by:[§6](https://arxiv.org/html/2608.17542#S6.SS0.SSS0.Px6.p1.1)\.
- Zhouet al\.\(2025\)G\. Zhou, H\. Pan, Y\. LeCun, and L\. PintoDINO\-WM: world models on pre\-trained visual features enable zero\-shot planning\.InProceedings of the 42nd International Conference on Machine Learning,Proceedings of Machine Learning Research, Vol\.267,pp\. 79115–79135\.Cited by:[§1](https://arxiv.org/html/2608.17542#S1.p1.1),[§6](https://arxiv.org/html/2608.17542#S6.SS0.SSS0.Px2.p1.1)\.

## Appendix AEvaluation Protocol

Our controlled standard\-task planning tables use 200 episodes, evaluation seed 42, exact checkpoint stems, CEM 300 samples / 30 elites / 30 iterations, and unique output files\. SIGReg, MTM\-MSE, and AC\-MTM use final epoch\-10 checkpoints from training seeds\{3072,1,2\}\\\{3072,1,2\\\}\. No validation\-based checkpoint selection is used in the controlled tables\. The non\-LeWM baselines in Figure[3](https://arxiv.org/html/2608.17542#S4.F3)are paper\-reported values from[14](https://arxiv.org/html/2608.17542#bib.bib8); they are included for continuity with the original evaluation style, not as reruns\.

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](https://github.com/lucas-maes/le-wm/issues): \#38, \#67, \#41, \#37, and \#62\.
## Appendix BInterrogating the OGBench Scene Result

The Scene margin is large enough that it deserves a separate audit\. Table[10](https://arxiv.org/html/2608.17542#A2.T10)lists the checks we performed on the completed trajectory\-goal runs\. The headline survives these checks, but the causal interpretation should remain narrow: this is evidence about our matched trajectory\-goal MPC protocol, not the public OGBench fixed\-goal benchmark\.

Table 10:Audit of the completed OGBench Visual Scene trajectory\-goal comparison\.#### Why can the gap be this large?

Scene couples several slow variables \(drawer, window, two buttons, and a cube state\) to the same arm observation\. A final\-latent MPC cost is brittle to dropping any one of these variables: a rollout can be close in agent pose while still missing the target object configuration\. SIGReg prevents constant collapse by enforcing a global isotropic latent distribution, but that guarantee does not say which state factors are preserved with a geometry useful for final latent distance\. In a multi\-object scene, part of the latent budget can be spent matching the global marginal or representing visually variable but planning\-irrelevant factors\.

AC\-MTM supplies a different pressure\. To identify the observed action among in\-batch alternatives, adjacent latents must retain information about controllable changes\. This favors variables whose changes explain the executed control and therefore aligns more directly with the CEM rollout interface\. The training logs also show that AC\-MTM reaches a lower final one\-step forward MSE on Scene \(about0\.0020\.002versus0\.0040\.004for SIGReg in the last logged batches\), so the observed advantage may combine better transition geometry with more accurate short\-horizon latent dynamics\. Consistent with this account, the non\-contrastive MTM\-MSE ablation trained under the identical Scene protocol also clears SIGReg \(75\.3% vs\. 58\.0%, paired 39/13,p≈4×10−4p\\\!\\approx\\\!4\{\\times\}10^\{\-4\}\), so the win belongs to inverse\-dynamics supervision as a family, with the contrastive form contributing a further \(noise\-level\) 4\.7 points\. The result should therefore be interpreted as a downstream planning win for dynamics\-derived training signals, not as a pure proof that latent marginal normality alone is harmful\.

#### The official fixed\-goal protocol\.

For completeness we also ran the official OGBench Visual Scene protocol \(five fixed goal tasks, 750\-step cap\) with the matched seed\-3072 checkpoints: both SIGReg and AC\-MTM score0/2500/250, with every episode reaching the step cap\. The official tasks demand much longer\-horizon manipulation than the2525\-step trajectory goals used in training\-matched evaluation, and the final\-latent CEM planner does not solve them at this model scale for either objective\. This is why Table[4](https://arxiv.org/html/2608.17542#S4.T4)reports the matched trajectory\-goal protocol: it is the regime in which these world models operate, and it keeps the comparison between training signals rather than between planners and task horizons\.

#### What would falsify the interpretation?

Two further checks would make the Scene story tighter\. First, an evaluation\-seed sweep or largernnwould test whether the same paired advantage holds beyond the current 150 matched episodes\. Second, per\-factor diagnostics \(drawer, window, button, and cube target probes or per\-target success decomposition\) would show which Scene variables SIGReg loses\.

Probe details:n=4000n\{=\}4000frames, ridgeα=1\\alpha\{=\}1, a70/3070/30split, and per\-dimensionR2R^\{2\}\{\}on privileged state\. PushT probes agentx,yx,y, blockx,yx,y, block orientation, and two velocity components\. All probe rows shown use epoch\-10 checkpoints\.

#### CEM action\-block semantics\.

At test time, no observed action is supplied\. The planner samples actions in the same normalized action space used by the world model\. With raw action dimensiondd, planning horizonH=5H\{=\}5, and action blockK=5K\{=\}5, CEM samples candidate sequencesA∈ℝB×300×H×K​dA\\in\\mathbb\{R\}^\{B\\times 300\\times H\\times Kd\}\. Each of theHHcoarse planner slots is therefore a flattened block ofKKlow\-level simulator actions\. CEM initializes a diagonal Gaussian over these blocks, samples 300 candidate sequences, inserts the current mean as one candidate, rolls the world model forward under every candidate, keeps the 30 lowest final latent\-goal MSE sequences, and refits the Gaussian mean and per\-component scale to those elites\. This sample\-score\-refit loop is repeated for 30 iterations, after which the policy returns the final mean sequence\.

Figure[4](https://arxiv.org/html/2608.17542#A2.F4)summarizes how the three time scales fit together\. They coincide at the value five in this configuration and should not be conflated\. The*action block*K=5K\{=\}5is the number of raw simulator steps bundled into one coarse transition \(equal to the dataset frameskip\); the*planning horizon*H=5H\{=\}5is the number of coarse transitions the planner looks ahead, i\.e\.H​K=25HK\{=\}25simulator steps; and the*receding horizon*R=5R\{=\}5is the number of coarse transitions executed before the policy re\-plans, i\.e\.R​K=25RK\{=\}25environment steps per replan\. HereR=HR\{=\}H, so each plan is executed in full before re\-planning; smallerRRyields tighter closed\-loop control and warm\-starts the next solve from the unused tail of the current plan\. The world model predicts the next coarse*latent state*, not actions and not theK−1K\{\-\}1skipped frames: one forward step maps the latent of coarse framettto that of coarse framet\+1t\{\+\}1conditioned on the entireKK\-action block, and only the final predicted latent is scored against the goal\. The action block is thus the optimization variable supplied by CEM at test time and by the dataset at training time\. It is never an output of the model\.

![Refer to caption](https://arxiv.org/html/2608.17542v1/images/action-blocks-and-cem-rollout-explainer.png)Figure 4:Action blocks and CEM rollouts \(TwoRoom,d=2d\{=\}2\)\. Each coarse transition bundlesK=5K\{=\}5raw simulator actions into one model input block\. CEM samplesH=5H\{=\}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\-levelenv\.stepactions before re\-planning\.For example, on a two\-dimensional action task withK=5K\{=\}5, one coarse action slot is

\[a0x,a0y,a1x,a1y,a2x,a2y,a3x,a3y,a4x,a4y\]\.\[a^\{x\}\_\{0\},a^\{y\}\_\{0\},a^\{x\}\_\{1\},a^\{y\}\_\{1\},a^\{x\}\_\{2\},a^\{y\}\_\{2\},a^\{x\}\_\{3\},a^\{y\}\_\{3\},a^\{x\}\_\{4\},a^\{y\}\_\{4\}\]\.This slot connects adjacent model frames separated by five simulator steps\. The policy reshapes it into the ordered low\-level sequence\(a0,a1,a2,a3,a4\)\(a\_\{0\},a\_\{1\},a\_\{2\},a\_\{3\},a\_\{4\}\), inverse\-normalizes those five actions, and executes them with five consecutive environment steps\. Thus the action associated with a coarse transition is not one command repeated for five frames; it is the whole intervening action block\. The HDF5 training loader uses the same convention: pixels and other non\-action columns are subsampled everyKKsteps, while the action column is kept dense and reshaped into one flattened block per coarse transition\.

#### Cold start and self\-conditioning\.

The rollout never receives an empty context\. On episode reset the frame\-history buffer is pre\-filled by repeating the initial observation, so the policy always supplies a fullhistory\_size\-length context; for TwoRoom evaluationhistory\_size=1\{=\}1, that context is the single current frame\. The predictor is a Transformer, so a length\-one sequence is valid input: it encodes the initial frame to𝐳0\\mathbf\{z\}\_\{0\}and predicts𝐳1\\mathbf\{z\}\_\{1\}from𝐳0\\mathbf\{z\}\_\{0\}and the first action block\. Thereafter the autoregressive rollout conditions on its*own*predicted latents until the window reaches the trained length and then slides, so beyond the first step the history is imagined rather than observed, which is precisely why a stable latent geometry matters, as an early error propagates through the remaining coarse steps\. The autoregressive model trains on three real context frames but TwoRoom evaluation begins from one, a mild context\-length difference the variable\-length predictor tolerates\.

#### CEM distribution updates\.

The sampling distribution is a*diagonal*Gaussian over the flattened plan: for TwoRoom the plan isH×K​d=5×10H\{\\times\}Kd=5\{\\times\}10, so the mean and scale are each5×105\{\\times\}10tensors; fifty independent one\-dimensional Gaussians per environment, with no cross\-coordinate covariance\. Each iteration draws300300plansAs=μ\+σ⊙εsA\_\{s\}=\\mu\+\\sigma\\odot\\varepsilon\_\{s\}withεs∼𝒩⁡\(0,I\)\\varepsilon\_\{s\}\\sim\\mathcal\{N\}\(0,I\), forcesA0=μA\_\{0\}\{=\}\\mu, scores each by final\-latent goal distance, and refitsμ\\muandσ\\sigmato the per\-coordinate mean and standard deviation of the3030elite plans\. Selection is*joint*\(whole plans are ranked\) but the refit is*marginal*\(each coordinate’s new moments are those of the elites at that coordinate\), so a coordinate that drives the cost is sharpened while an irrelevant one stays diffuse\. For instance, a coordinate whose useful value is≈\+0\.7\{\\approx\}\{\+\}0\.7might moveμ:→→→0\.69\\mu:0\.00\\\!\\to\\\!0\.42\\\!\\to\\\!0\.61\\\!\\to\\\!0\.69withσ\\sigmacontracting→→→0\.161\.00\\\!\\to\\\!0\.55\\\!\\to\\\!0\.30\\\!\\to\\\!0\.16over the first iterations and freezing nearμ≈0\.71\\mu\{\\approx\}0\.71,σ≈0\.04\\sigma\{\\approx\}0\.04, whereas an unleveraged coordinate keepsμ≈0\\mu\{\\approx\}0,σ≈1\\sigma\{\\approx\}1\. The refit is*hard*\. Moments are replaced, with no step size or variance floor, so the Gaussian collapses to nearly a point within roughly fifteen iterations\.

#### Training\-only action contrast\.

For eachB×TB\\times Tbatch, the inverse task flattensN=B⁡\(T−1\)N=B\(T\-1\)adjacent transitions and uses all observed actions in that flattened batch as candidates in Equation[3](https://arxiv.org/html/2608.17542#S3.E3)\. Each query𝐚^i=invψ​\(𝐳i,𝐳i\+1\)\\hat\{\\mathbf\{a\}\}\_\{i\}=\\mathrm\{inv\}\_\{\\psi\}\(\\mathbf\{z\}\_\{i\},\\mathbf\{z\}\_\{i\+1\}\)and candidate𝐚j\\mathbf\{a\}\_\{j\}is a full coarse action block of dimensionK​dKd\(allKKraw actions, e\.g\.1010on TwoRoom\), regressed in a single pass from the latent pair rather than decoded intoKKordered actions; the horizon\-conditioned variant instead predicts the first block𝐚t\\mathbf\{a\}\_\{t\}from akk\-step\-apart pair\(𝐳t,𝐳t\+k,ek\)\(\\mathbf\{z\}\_\{t\},\\mathbf\{z\}\_\{t\+k\},e\_\{k\}\)\. The inverse head is optimized and checkpointed for reproducibility, but rollout and cost functions never call it; removing it after training leaves predictions and CEM action selection unchanged\.

## Appendix CControlled Three\-Way Ablation

Table[11](https://arxiv.org/html/2608.17542#A3.T11)and Figure[5](https://arxiv.org/html/2608.17542#A3.F5)give the matched comparison between SIGReg, the non\-contrastive inverse\-regression ablation, and AC\-MTM\. This diagnostic isolates the effect of Action\-NCE: on TwoRoom, PushT, and Cube, AC\-MTM remains within 1\.2 percentage points of MTM\-MSE; on Reacher, MTM\-MSE collapses for two of three seeds whereas AC\-MTM remains non\-collapsed; on Scene, both inverse variants beat SIGReg decisively and AC\-MTM leads MTM\-MSE by 4\.7 points \(paired 16/9,p≈0\.23p\\\!\\approx\\\!0\.23\)\.

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±\\pmstandard deviation over three training seeds\.![Refer to caption](https://arxiv.org/html/2608.17542v1/images/results_ablation_threeway.png)Figure 5:Controlled three\-way ablation with 200 evaluation episodes\. Error bars show standard deviation over three training seeds\. The Reacher panel shows why Action\-NCE is used as the main method: it removes the bimodal collapse of the non\-contrastive inverse\-regression ablation without changing the test\-time planner\.
## Appendix DExploratory Single\-Seed Results on Additional OGBench Families

To probe how the two training signals behave outside the main suite, we ran single\-seed \(3072\) comparisons on six additional OGBench environments under the same trajectory\-goal protocol \(50 episodes, evaluation seed 42\)\. These runs use one training seed and no tuning, so they are reported as exploratory scope evidence, not controlled claims\.

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\.The pattern is consistent with the PushT analysis in Section[4\.5](https://arxiv.org/html/2608.17542#S4.SS5)\. The puzzle tasks are dominated by a grid of buttons whose visual state changes discretely and near\-identically regardless of which action toggled them; identifying the executed action then provides little pressure to represent the full button configuration, whereas SIGReg’s global variance pressure preserves it\. Frozen\-latent probes confirm this mechanism: on Puzzle 4x4, SIGReg decodes the variable button bits atR2≈0\.98R^\{2\}\{\}\\\!\\approx\\\!0\.98while AC\-MTM decodes them at chance \(R2≤0R^\{2\}\{\}\\\!\\leq\\\!0\), even though both models decode the continuous arm pose atR2≈0\.99R^\{2\}\{\}\\\!\\approx\\\!0\.99\. Conversely, where the environment adds stochastic or diffuse dynamics \(teleporting maze, falling\-particle Powderworld\), the dynamics\-native signal is at least as good\. We report these runs to delimit the method’s scope: the claim is not that contrastive inverse dynamics dominates distributional regularization everywhere, but that it is a robust distribution\-free alternative whose advantage grows with the number of coupled controllable factors, as in Scene\.