Contrastive World Models
Summary
Contrastive World Models propose a new approach for learning latent dynamics without pixel reconstruction, using a contrastive objective to improve robustness and efficiency in visually complex environments for model-based reinforcement learning.
View Cached Full Text
Cached at: 09/22/26, 09:16 AM
# Contrastive World Models
Source: [https://arxiv.org/html/2609.22175](https://arxiv.org/html/2609.22175)
###### Abstract
World models trained via pixel reconstruction can struggle in visually complex environments, where irrelevant information dominates the objective and distract the model from information relevant to planning and control\. We present Contrastive World Models, an approach for learning latent dynamics models without pixel reconstruction\. Building on Dreamer, we replace observation reconstruction in the standard world model objective with a Deep InfoMax\-like lower bound that maximizes the mutual information between state\-action sequences and local patch features of future observations, encouraging state representations to retain information that is predictive of the future without requiring the model to reconstruct visually irrelevant details\. We evaluate our approach in small\-scale experiments across three settings of increasing visual complexity\. Our method matches Dreamer and a momentum prediction baseline in the default setting, and substantially outperforms both once distractors or natural video backgrounds are introduced, while also training more efficiently by removing the pixel decoder entirely\. Our approach is general and makes minimal assumptions beyond access to state\-action sequences and future observations\. These results suggest that contrastive, infomax\-based objectives are a principled and promising direction for building world models that are robust to visual nuisance factors, a property particularly relevant for transferring model\-based RL agents to the real world\.
Figure 1:Pixel\-prediction\-based World Models \(left\) and Contrastive World Models \(right\)\. Standard world models encode each observation into a latent state \(green\) via an encoder \(blue\) and reconstruct the input image through a pixel decoder \(gray\)\. Contrastive World Models remove the decoder entirely and instead train the same encoder and latent dynamics model by maximizing mutual information between latent states and local features of future observations\.## 1Introduction
Learning a model of the environment’s dynamics – aworld model– has emerged as a central paradigm for building agents that can plan, imagine, and act efficiently from high\-dimensional observations\.[Ha and Schmidhuber \[2018\]](https://arxiv.org/html/2609.22175#bib.bib12)showed that an agent could be decomposed into a vision component \(a variational autoencoder\), a memory component \(a recurrent network\), and a small controller, and that the controller could be trained entirely "inside" imagined rollouts of the learned model before being transferred back to the real environment\. World models have now become a practical recipe and a broader direction that has been adopted well beyond simple control tasks, such as for interactive world generation\[[Parker\-Holder et al\., 2024](https://arxiv.org/html/2609.22175#bib.bib16),[Ball et al\., 2025](https://arxiv.org/html/2609.22175#bib.bib17)\], autonomous driving\[[Zhang et al\., 2024](https://arxiv.org/html/2609.22175#bib.bib18),[Russell et al\., 2025](https://arxiv.org/html/2609.22175#bib.bib19)\], and robotics\[[Yang et al\., 2024](https://arxiv.org/html/2609.22175#bib.bib20),[Ye et al\., 2026](https://arxiv.org/html/2609.22175#bib.bib21)\]\.
The Dreamer and Genie series has driven much of the progress along this direction\.[Hafner et al\. \[2019\]](https://arxiv.org/html/2609.22175#bib.bib1)introduced the Recurrent State Space Model \(RSSM\), combining deterministic and stochastic recurrent latents for planning via model\-predictive control;[Hafner et al\. \[2020\]](https://arxiv.org/html/2609.22175#bib.bib2)built on this by learning a value function and actor from imagined rollouts, replacing planning with amortized policy learning; and DreamerV2\[[Hafner et al\., 2021](https://arxiv.org/html/2609.22175#bib.bib13)\]and DreamerV3\[[Hafner et al\., 2025](https://arxiv.org/html/2609.22175#bib.bib14)\]progressively improved robustness on top of the RSSM backbone, with DreamerV3 mastering over 150 tasks under a single fixed hyperparameter set\. In parallel, the Genie series has pushed toward large\-scale, action\-controllable video generation\. Genie\[[Bruce et al\., 2024](https://arxiv.org/html/2609.22175#bib.bib15)\]learns to generate interactive 2D environments purely from unlabeled video, Genie 2\[[Parker\-Holder et al\., 2024](https://arxiv.org/html/2609.22175#bib.bib16)\]extends this to diverse 3D environments, and Genie 3\[[Ball et al\., 2025](https://arxiv.org/html/2609.22175#bib.bib17)\]generates real\-time, minutes\-long interactive worlds from text alone\. More recently, SIMA 2\[[Bolton et al\., 2025](https://arxiv.org/html/2609.22175#bib.bib22)\]demonstrated an embodied agent operating directly within Genie 3\-generated worlds\. All these works use reconstruction or pixel\-prediction objectives to learn latent state representations\.
Contrastive learning offers a natural alternative to reconstruction\. Rather than regenerating its input, contrastive objectives train an encoder to distinguish related \("positive"\) pairs of views or timesteps from unrelated \("negative"\) ones, typically by optimizing the InfoNCE lower bound on mutual information\[[van den Oord et al\., 2019](https://arxiv.org/html/2609.22175#bib.bib3)\]\. This principle underlies much of the progress in self\-supervised visual representation learning: Contrastive Predictive Coding\[[van den Oord et al\., 2019](https://arxiv.org/html/2609.22175#bib.bib3)\]and Deep InfoMax\[[Hjelm et al\., 2019](https://arxiv.org/html/2609.22175#bib.bib4)\]maximize mutual information between global and local features of an input; SimCLR\[[Chen et al\., 2020](https://arxiv.org/html/2609.22175#bib.bib23)\]and MoCo\[[He et al\., 2020](https://arxiv.org/html/2609.22175#bib.bib24)\]scale instance discrimination with large or memory\-bank negative sets; and BYOL\[[Grill et al\., 2020](https://arxiv.org/html/2609.22175#bib.bib25)\]replaces explicit negatives with a momentum target network and prediction task\. For world models, these objectives never require pixels to be reconstructed: the encoder only needs to preserve information that discriminates the correct future from alternatives\. Similar motivations have driven representation learning in model\-free RL, where DrQ\[[Yarats et al\., 2021](https://arxiv.org/html/2609.22175#bib.bib7)\], ATC\[[Stooke et al\., 2021](https://arxiv.org/html/2609.22175#bib.bib8)\], ST\-DIM\[[Anand et al\., 2020](https://arxiv.org/html/2609.22175#bib.bib9)\], DRIML\[[Mazoure et al\., 2020](https://arxiv.org/html/2609.22175#bib.bib11)\], and MPR\[[Schwarzer et al\., 2021](https://arxiv.org/html/2609.22175#bib.bib26)\]use augmentation\-based, contrastive, or momentum\-based auxiliary objectives to improve sample efficiency, though without an explicit dynamics model for planning\.
There has recently been a resurgence of interest in learning world models directly in latent space, without pixel\-level generation\. Notably,[Hafner et al\. \[2020\]](https://arxiv.org/html/2609.22175#bib.bib2)evaluated a contrastive variant of Dreamer and found that it underperformed pixel reconstruction on the majority of tasks at the time\. Another early work is CSWM\[[Kipf et al\., 2020](https://arxiv.org/html/2609.22175#bib.bib10)\], which encodes each observation into a set of object slots and models transitions with a graph neural network over these slots, trained with a TransE\-style energy\-based hinge loss with negative sampling; however, it is evaluated only on latent ranking metrics in simple deterministic domains, without rewards or downstream control\. The Joint Embedding Predictive Architecture \(JEPA\)\[[LeCun, 2022](https://arxiv.org/html/2609.22175#bib.bib27)\]pursues the same principle at scale, predicting the latent embedding of a target view from a context view so that unpredictable low\-level detail is discarded, with instantiations for images\[[Assran et al\., 2023](https://arxiv.org/html/2609.22175#bib.bib28)\]and video\[[Bardes et al\., 2024](https://arxiv.org/html/2609.22175#bib.bib29),[Assran et al\., 2025](https://arxiv.org/html/2609.22175#bib.bib30)\]\. Closest to our setting,[Maes et al\. \[2026\]](https://arxiv.org/html/2609.22175#bib.bib31)introduce LeWorldModel \(LeWM\), a JEPA\-based world model that trains stably end\-to\-end from raw pixels using only a next\-embedding prediction loss and a Gaussian regularizer, dispensing with the stop\-gradients, momentum encoders, and multi\-term losses earlier variants required to prevent collapse\. However, it has so far been validated only at small scale in visually clean domains—LeWM is a∼15\{\\sim\}15M\-parameter model evaluated on four low\-distraction environments with stationary backgrounds \(2D maze navigation, Reacher, Push\-T, and OGBench\-Cube\)—and has not been shown to be robust to complex visual nuisance factors, which our work addresses\.
In this work, we combine these two directions – contrastive representation learning and Dreamer world model – by replacing the reconstruction\-based observation model in Dreamer with a mutual information maximization objective similar to Deep InfoMax, removing the pixel decoder entirely\. We hypothesize that this approach should have little effect in visually simple settings but should yield substantially more robust representations and more efficient training, and therefore stronger downstream control performance, as visual complexity and distraction increase\.
## 2Contrastive World Models
We consider a partially observable Markov decision process \(POMDP\)\. We define a discrete time step t, image observationsoto\_\{t\}, hidden statessts\_\{t\}, continuous action vectorsata\_\{t\}, and scalar rewardsrtr\_\{t\}\. We seek to build a latent dynamic model of the environment, of which an agent can use to plan and to maximize its expected cumulative reward\. To do this, we build on top of a strong model\-based RL algorithm, Dreamer\[[Hafner et al\., 2020](https://arxiv.org/html/2609.22175#bib.bib2)\], and propose an alternative InfoMax objective to learn robust state representations\. We now describe our approach in detail\.
### 2\.1Latent Dynamics Model
We consider learning a latent dynamics model for a POMDP with image observationsoto\_\{t\}, continuous actionsata\_\{t\}, scalar rewardsrtr\_\{t\}, and latent statessts\_\{t\}\. The generative model, or “world model,” factorizes as
transition modelst∼p\(st∣st−1,at−1\)\\displaystyle\\quad s\_\{t\}\\sim p\(s\_\{t\}\\mid s\_\{t\-1\},a\_\{t\-1\}\)observation modelot∼p\(ot∣st\)\\displaystyle\\quad o\_\{t\}\\sim p\(o\_\{t\}\\mid s\_\{t\}\)reward modelrt∼p\(rt∣st\)\\displaystyle\\quad r\_\{t\}\\sim p\(r\_\{t\}\\mid s\_\{t\}\)state representation model \(encoder\)st∼q\(st∣st−1,at−1,ot\)\\displaystyle\\quad s\_\{t\}\\sim q\(s\_\{t\}\\mid s\_\{t\-1\},a\_\{t\-1\},o\_\{t\}\)
We would like to fit this model by maximizing the marginal log\-likelihood of the observed data,lnp\(o1:T,r1:T∣a1:T\)\\ln p\(o\_\{1:T\},r\_\{1:T\}\\mid a\_\{1:T\}\), but this requires marginalizing out the full latent trajectorys1:Ts\_\{1:T\}and is intractable\. We instead derive a tractable lower bound by introducing the variational posterior above and applying Jensen’s inequality\.
First, we rewrite the marginal likelihood as an expectation underqqby multiplying and dividing byq\(s1:T∣o1:T,a1:T\)=∏tq\(st∣st−1,at−1,ot\)q\(s\_\{1:T\}\\mid o\_\{1:T\},a\_\{1:T\}\)=\\prod\_\{t\}q\(s\_\{t\}\\mid s\_\{t\-1\},a\_\{t\-1\},o\_\{t\}\):
lnp\(o1:T,r1:T∣a1:T\)=ln𝔼q\[p\(o1:T,r1:T,s1:T∣a1:T\)q\(s1:T∣o1:T,a1:T\)\]\\ln p\(o\_\{1:T\},r\_\{1:T\}\\mid a\_\{1:T\}\)=\\ln\\mathbb\{E\}\_\{q\}\\left\[\\frac\{p\(o\_\{1:T\},r\_\{1:T\},s\_\{1:T\}\\mid a\_\{1:T\}\)\}\{q\(s\_\{1:T\}\\mid o\_\{1:T\},a\_\{1:T\}\)\}\\right\]Sinceln\(⋅\)\\ln\(\\cdot\)is concave, Jensen’s inequality givesln𝔼\[X\]≥𝔼\[lnX\]\\ln\\mathbb\{E\}\[X\]\\geq\\mathbb\{E\}\[\\ln X\], so
≥𝔼q\[lnp\(o1:T,r1:T,s1:T∣a1:T\)−lnq\(s1:T∣o1:T,a1:T\)\]\\geq\\;\\mathbb\{E\}\_\{q\}\\left\[\\ln p\(o\_\{1:T\},r\_\{1:T\},s\_\{1:T\}\\mid a\_\{1:T\}\)\-\\ln q\(s\_\{1:T\}\\mid o\_\{1:T\},a\_\{1:T\}\)\\right\]Expanding both terms using the factorizations of the generative model and the posterior,
=𝔼q\[∑t\(lnp\(ot∣st\)\+lnp\(rt∣st\)\+lnp\(st∣st−1,at−1\)−lnq\(st∣st−1,at−1,ot\)\)\]=\\;\\mathbb\{E\}\_\{q\}\\left\[\\sum\_\{t\}\\bigg\(\\ln p\(o\_\{t\}\\mid s\_\{t\}\)\+\\ln p\(r\_\{t\}\\mid s\_\{t\}\)\+\\ln p\(s\_\{t\}\\mid s\_\{t\-1\},a\_\{t\-1\}\)\-\\ln q\(s\_\{t\}\\mid s\_\{t\-1\},a\_\{t\-1\},o\_\{t\}\)\\bigg\)\\right\]and grouping the last two terms per timestep into a KL divergence yields our final bound,
=𝔼q\[∑t\(lnp\(ot∣st\)\+lnp\(rt∣st\)−KL\(q\(st∣st−1,at−1,ot\)∥p\(st∣st−1,at−1\)\)\)\]=\\;\\mathbb\{E\}\_\{q\}\\left\[\\sum\_\{t\}\\bigg\(\\ln p\(o\_\{t\}\\mid s\_\{t\}\)\+\\ln p\(r\_\{t\}\\mid s\_\{t\}\)\-KL\\Big\(q\(s\_\{t\}\\mid s\_\{t\-1\},a\_\{t\-1\},o\_\{t\}\)\\parallel p\(s\_\{t\}\\mid s\_\{t\-1\},a\_\{t\-1\}\)\\Big\)\\bigg\)\\right\]


Figure 2:Architecture of Recurrent State Space Model \(left\) and Contrastive World Model \(right\)Specifically, we build on top of Recurrent State Space Model\[[Hafner et al\., 2019](https://arxiv.org/html/2609.22175#bib.bib1)\]\. We now describe the architecture in detail, as shown in figure[2](https://arxiv.org/html/2609.22175#S2.F2)\.
Given observationoto\_\{t\}from timesteptt, encodereeencodes the high\-dimensional observation to a lower dimension vector, denoted byztz\_\{t\}\. The update functionggtakes encodingztz\_\{t\}and the predicted prior state features and outputs the posterior state features, denotedsts\_\{t\}\. The transition functionffconcatenates actionata\_\{t\}and the posterior state featuressts\_\{t\}and predicts the prior state features for the next state, denoted ass^t\+1\\hat\{s\}\_\{t\+1\}\. The prior prediction is Gaussian with mean and variance parameterized by the model output\. Furthermore, the reward functionRRmaps the posterior state feature to rewards\. The reward prediction is Gaussian with mean parameterized by the model output and unit variance\.
For the observation likelihood term, RSSM uses a decoderOOthat reconstructs the observationoto\_\{t\}from the posterior state featuressts\_\{t\}\.
The training objective becomes observation reconstruction, reward prediction, and a KL regularizer\. All the components of RSSM are optimized jointly\.
JREC=𝔼q\(∑t\(JOt\+JRt\+JDt\)\)J\_\{REC\}=\\mathbb\{E\}\_\{q\}\\left\(\\sum\_\{t\}\(J\_\{O\}^\{t\}\+J\_\{R\}^\{t\}\+J\_\{D\}^\{t\}\)\\right\)
where
JOt=ln\(O\(ot∣st\)\)J\_\{O\}^\{t\}=\\ln\\left\(O\(o\_\{t\}\\mid s\_\{t\}\)\\right\)JRt=ln\(R\(rt∣st\)\)J\_\{R\}^\{t\}=\\ln\\left\(R\(r\_\{t\}\\mid s\_\{t\}\)\\right\)JDt=−KL\(q\(st\+1∣st,at,ot\+1\)∥f\(s^t\+1∣st,at\)\)J\_\{D\}^\{t\}=\-KL\\big\(q\(s\_\{t\+1\}\\mid s\_\{t\},a\_\{t\},o\_\{t\+1\}\)\\parallel f\(\\hat\{s\}\_\{t\+1\}\\mid s\_\{t\},a\_\{t\}\)\\big\)
### 2\.2InfoMax Representations
Consider the observation likelihood term in the lower bound\. We can subtract the observation marginal, which does not depend onsts\_\{t\}and so does not shift the optimum with respect to the encoder, and apply Bayes’ rule\.
𝔼\[ln\(p\(ot∣st\)\)\]≐𝔼\[lnp\(ot∣st\)−lnp\(ot\)\]=𝔼\[lnp\(ot,st\)p\(ot\)p\(st\)\]\\begin\{split\}\\mathbb\{E\}\\left\[\\ln\{\\left\(p\(o\_\{t\}\\mid s\_\{t\}\)\\right\)\}\\right\]&\\doteq\\mathbb\{E\}\\left\[\\ln\{p\(o\_\{t\}\\mid s\_\{t\}\)\}\-\\ln\{p\(o\_\{t\}\)\}\\right\]\\\\ &=\\mathbb\{E\}\\left\[\\ln\{\\frac\{p\(o\_\{t\},s\_\{t\}\)\}\{p\(o\_\{t\}\)\\,p\(s\_\{t\}\)\}\}\\right\]\\end\{split\}\(1\)Remark that the resulting density ratio is the mutual information between the observationoto\_\{t\}and posterior state featuressts\_\{t\}\. This suggests mutual information maximization as an alternative objective for learning state representations in world models, without requiring an explicit pixel decoder\.
To learn temporally dependent representations, we want to maximize mutual information between state\-action sequences and future observations in the trajectory\.
maxθI\(\[st,at:t\+h\];ot\+h\)\\max\_\{\\theta\}I\(\[s\_\{t\},a\_\{t:t\+h\}\];o\_\{t\+h\}\)which is bounded by the InfoNCE lower bound up to a constant
𝔼\[lnfθ\(\[st,at:t\+h\],ot\+h\)∑t′fθ\(\[st,at:t\+h\],ot′\)\]\\mathbb\{E\}\\left\[\\ln\\frac\{f\_\{\\theta\}\(\[s\_\{t\},a\_\{t:t\+h\}\],o\_\{t\+h\}\)\}\{\\sum\_\{t^\{\\prime\}\}f\_\{\\theta\}\(\[s\_\{t\},a\_\{t:t\+h\}\],o\_\{t^\{\\prime\}\}\)\}\\right\]wherefθf\_\{\\theta\}is a score function which preserves mutual information, i\.e\.fθ\(xt\+h,ct\)∝p\(xt\+h∣ct\)p\(xt\+h\)f\_\{\\theta\}\(x\_\{t\+h\},c\_\{t\}\)\\propto\\frac\{p\(x\_\{t\+h\}\\mid c\_\{t\}\)\}\{p\(x\_\{t\+h\}\)\}when optimizing the above bound\[[van den Oord et al\., 2019](https://arxiv.org/html/2609.22175#bib.bib3)\]\. While any score function can be used, a simple log\-bilinear model usually suffices:
fθ\(xt\+h,ct\)=exp\(ct⊤Wθxt\+h\)f\_\{\\theta\}\(x\_\{t\+h\},c\_\{t\}\)=\\exp\(c\_\{t\}^\{\\top\}W\_\{\\theta\}x\_\{t\+h\}\)
Following Deep InfoMax\[[Hjelm et al\., 2019](https://arxiv.org/html/2609.22175#bib.bib4)\], in our setting we maximize mutual information between the global state posteriorsts\_\{t\}and local features of future states, i\.e\., patches of the encodingem,n\(ot\+h\)e\_\{m,n\}\(o\_\{t\+h\}\)\. Here we concatenate the global state posteriorsts\_\{t\}with the action sequenceat:t\+ha\_\{t:t\+h\}leading up to the future state\.
maxθI\(\[st,at:t\+h\];ot\+h\)≥𝔼T,B\(1M1N∑m∑nlnexp\(\[st,at:t\+h\]⊤Wθem,n\(ot\+h\)\)∑b∈Bexp\(\[st,at:t\+h\]⊤Wθem,n\(ob\)\)\)\\max\_\{\\theta\}I\(\[s\_\{t\},a\_\{t:t\+h\}\];o\_\{t\+h\}\)\\geq\\mathbb\{E\}\_\{T,B\}\\left\(\\frac\{1\}\{M\}\\frac\{1\}\{N\}\\sum\_\{m\}\\sum\_\{n\}\\ln\\frac\{\\exp\(\[s\_\{t\},a\_\{t:t\+h\}\]^\{\\top\}W\_\{\\theta\}e\_\{m,n\}\(o\_\{t\+h\}\)\)\}\{\\sum\_\{b\\in B\}\\exp\(\[s\_\{t\},a\_\{t:t\+h\}\]^\{\\top\}W\_\{\\theta\}e\_\{m,n\}\(o\_\{b\}\)\)\}\\right\)\(2\)
whereT,BT,Bdenote the time and batch dimension of the sampled sequence respectively, andm,nm,ndenote themmth andnnth patch of the spatial features of the encoding\. For the training objective, we have
JIt=1M1N∑m∑nlnexp\(\[st,at:t\+h\]⊤Wθem,n\(ot\+h\)\)∑b∈Bexp\(\[st,at:t\+h\]⊤Wθem,n\(ob\)\)J\_\{I\}^\{t\}=\\frac\{1\}\{M\}\\frac\{1\}\{N\}\\sum\_\{m\}\\sum\_\{n\}\\ln\\frac\{\\exp\(\[s\_\{t\},a\_\{t:t\+h\}\]^\{\\top\}W\_\{\\theta\}e\_\{m,n\}\(o\_\{t\+h\}\)\)\}\{\\sum\_\{b\\in B\}\\exp\(\[s\_\{t\},a\_\{t:t\+h\}\]^\{\\top\}W\_\{\\theta\}e\_\{m,n\}\(o\_\{b\}\)\)\}JRt=ln\(R\(rt∣st\)\)J\_\{R\}^\{t\}=\\ln\\left\(R\(r\_\{t\}\\mid s\_\{t\}\)\\right\)JDt=−KL\(q\(st\+1∣st,at,ot\+1\)∥f\(s^t\+1∣st,at\)\)J\_\{D\}^\{t\}=\-KL\\big\(q\(s\_\{t\+1\}\\mid s\_\{t\},a\_\{t\},o\_\{t\+1\}\)\\parallel f\(\\hat\{s\}\_\{t\+1\}\\mid s\_\{t\},a\_\{t\}\)\\big\)JDIM=𝔼\(1T∑t\(JIt\+λ1JDt\+λ2JRt\)\)J\_\{DIM\}=\\mathbb\{E\}\\left\(\\frac\{1\}\{T\}\\sum\_\{t\}\(J\_\{I\}^\{t\}\+\\lambda\_\{1\}J\_\{D\}^\{t\}\+\\lambda\_\{2\}J\_\{R\}^\{t\}\)\\right\)
whereJItJ\_\{I\}^\{t\}is the Deep InfoMax representation loss,JRtJ\_\{R\}^\{t\}is the reward prediction objective, andJDtJ\_\{D\}^\{t\}is the KL transition objective\.λ1,λ2\\lambda\_\{1\},\\lambda\_\{2\}scale the transition loss and reward loss respectively\.
## 3Experiments
Our hypothesis is that our infomax approach should learn substantially more robust representations in visually diverse environments\. To this end, we evaluate our approach in three settings shown in Figure[3](https://arxiv.org/html/2609.22175#S3.F3), results are averaged over 3 seeds\.Default Setting: this is the default setting from DeepMind Control Suite with stationary background\.Simple Distractor Setting: we incorporate simple moving distractors in the background, consisted of colored balls moving and bouncing off the frames\[[Yarats et al\., 2020](https://arxiv.org/html/2609.22175#bib.bib6),[Zhang et al\., 2021](https://arxiv.org/html/2609.22175#bib.bib5)\]\.Natural Video Setting: We incorporate natural video from the Kinetics dataset as background, as done in\[[Zhang et al\., 2021](https://arxiv.org/html/2609.22175#bib.bib5)\]\. Note that in DBC\[[Zhang et al\., 2021](https://arxiv.org/html/2609.22175#bib.bib5)\], the natural background are grayscale, here we use full color channels without grayscale, which could be more challenging for the agents\.









Figure 3:Evaluation settings\. Upper left three: default DMC setting; Upper right three: simple distractor setting; Lower three: natural video setting\.
Figure 4:Default DMC setting\.
Figure 5:Simple distractor setting\.
Figure 6:Natural video setting\.### 3\.1Baselines and Experimental Setup
We compare our Contrastive World Model \(InfoMax, ours\) against two baselines:Dreamer\[[Hafner et al\., 2020](https://arxiv.org/html/2609.22175#bib.bib2)\], which learns state representations via pixel reconstruction, andMomentum Prediction, a baseline that replaces the reconstruction objective with a momentum\-encoder\-based target but otherwise shares the same RSSM backbone and actor\-critic as our method\.
#### Momentum Prediction baseline\.
As a second baseline, we consider a self\-predictive representation objective similar to BYOL\[[Grill et al\., 2020](https://arxiv.org/html/2609.22175#bib.bib25)\]and MPR\[[Schwarzer et al\., 2021](https://arxiv.org/html/2609.22175#bib.bib26)\], which we refer to as Momentum Prediction\. Rather than reconstructing pixels or maximizing an InfoNCE bound, this baseline maintains a momentum\-averaged target encodereξe\_\{\\xi\}, whose parametersξ\\xiare an exponential moving average of the online encoderee’s parametersθ\\theta:
ξ←τξ\+\(1−τ\)θ\\xi\\leftarrow\\tau\\xi\+\(1\-\\tau\)\\thetaGiven the RSSM statests\_\{t\}\(the concatenation of the deterministic belief and stochastic posterior components\), a projection headϕθ\\phi\_\{\\theta\}predicts the embedding of a future observationot\+ho\_\{t\+h\}, and the online prediction is trained to match the stop\-gradient target embedding produced by the momentum encoder:
JMPt=cos\(ϕθ\(st\),sg\(eξ\(ot\+h\)\)\)J\_\{MP\}^\{t\}=\\cos\\Big\(\\phi\_\{\\theta\}\(s\_\{t\}\),\\ \\mathrm\{sg\}\\big\(e\_\{\\xi\}\(o\_\{t\+h\}\)\\big\)\\Big\)wheresg\(⋅\)\\mathrm\{sg\}\(\\cdot\)denotes the stop\-gradient operator andcos\(⋅,⋅\)\\cos\(\\cdot,\\cdot\)is cosine similarity\. Unlike our InfoMax objective,JMPtJ\_\{MP\}^\{t\}contains no explicit negative samples; representational collapse is instead avoided through the asymmetry between the online projector and the momentum target, following[Grill et al\. \[2020\]](https://arxiv.org/html/2609.22175#bib.bib25)\. As with our method, this loss is combined with the same reward and transition losses used elsewhere:
JMP=𝔼\(1T∑t\(JMPt\+λ1JDt\+λ2JRt\)\)J\_\{MP\}=\\E\\left\(\\frac\{1\}\{T\}\\sum\_\{t\}\(J\_\{MP\}^\{t\}\+\\lambda\_\{1\}J\_\{D\}^\{t\}\+\\lambda\_\{2\}J\_\{R\}^\{t\}\)\\right\)
All three agents use an identical actor\-critic planner on top of their respective learned latent dynamics, and differ only in the world model representation\-learning objective\. Each agent is trained for1×1061\\times 10^\{6\}environment steps on three DeepMind Control Suite tasks \(finger\-spin,cheetah\-run,walker\-walk\), and all results are averaged over 3 seeds\.
### 3\.2Results
Default DMC setting\.In the default, distraction\-free setting \(Figure[6](https://arxiv.org/html/2609.22175#S3.F6)\), all three methods eventually reach comparable asymptotic performance, indicating that when the observation is uncluttered, pixel reconstruction provides a sufficient learning signal and our InfoMax objective provides near\-equivalent performance\. Momentum Prediction is competitive with the other two methods in this setting, confirming that BYOL objectives are viable state representation learners when the observation is simple\.
Simple distractor setting\.When simple moving distractors are added\. \(Figure[6](https://arxiv.org/html/2609.22175#S3.F6)\), a clear separation emerges\. InfoMax matches or exceeds Dreamer on two tasks, most notably oncheetah\-run, where Dreamer’s reconstruction objective struggles to disentangle the task\-relevant cheetah dynamics from the moving distractors and plateaus at a substantially lower reward\. Momentum Prediction, in contrast, collapses in this setting across all three tasks, staying near its initial performance throughout training\. We hypothesize that without an explicit local\-global mutual information structure, the prediction target is easily satisfied by encoding distractor motion rather than task\-relevant dynamics, whereas our patch\-based Deep InfoMax objective encourages the state representation to retain information that is jointly predictive of future local image regions, which distractors do not consistently provide\.
Natural video setting\.The gap widens further when natural video from Kinetics is used as the background \(Figure[6](https://arxiv.org/html/2609.22175#S3.F6)\), the most visually challenging of our three settings since the background is neither stationary nor synthetically simple\. Here InfoMax substantially outperforms both baselines on every task, where Dreamer’s reconstruction loss is dominated by the high\-entropy natural video background and the agent fails to make significant task progress\. Momentum Prediction likewise plateaus early and fails to escape near\-random performance\. These results support our central hypothesis: because our objective only requires the state representation to be predictive of future observations, rather than requiring the full observation \(background included\) to be reconstructed, it is far more robust to visually complex and non\-stationary environments that carry less task\-relevant signal\.
Training efficiency\.Beyond final performance, removing the pixel decoder also reduces the per\-step compute cost of our method relative to Dreamer, since no pixel decoding network needs to be trained; we found InfoMax to be consistently substantially faster in wall\-clock time per training step across all three settings while using equal number of parameters in the encoder and dynamics model\.
## 4Conclusion
We present Contrastive World Models, an approach for training world models without pixel reconstruction\. In place of the observation reconstruction term in the RSSM objective, we derive a Deep InfoMax\-like lower bound that maximizes the mutual information between state\-action sequences and local patch features of future observations\. We evaluate our approach on the DeepMind Control Suite across three settings of increasing visual complexity\. Our method matches both the reconstruction\-based and momentum\-based baselines in the simple, distraction\-free setting, and significantly outperforms them once simple distractors or natural video backgrounds are introduced\. We further observe more efficient training from removing the pixel decoder entirely\. Together, these results suggest that Contrastive World Models are a principled and promising direction for building world models that are robust to visual nuisance factors, a property that is particularly relevant for transferring RL agents from simulation to the real world\.
Limitations:This work is limited to small\-scale experiments built on top of an older architecture \(RSSM\) on a relatively established benchmark \(DeepMind Control Suite\)\. Extending our approach to more recent transformer\-based world model architectures, and scaling to large domains, would be a valuable direction for further analysis\.
Future work:Our objective is general and makes minimal assumptions beyond access to state\-action sequences and future observations, without requiring action\-conditioned reconstruction, which makes it straightforward to apply to large\-scale, unlabeled pretraining settings as well\. Scaling our approach to larger domains, leveraging action\-free video data \(e\.g\., from YouTube\) for pretraining, and incorporating language\-specified goals are all promising directions for future work\.
## References
- Anandet al\.\(2020\)A\. Anand, E\. Racah, S\. Ozair, Y\. Bengio, M\. Côté, and R\. D\. HjelmUnsupervised state representation learning in atari\.External Links:1906\.08226,[Link](https://arxiv.org/abs/1906.08226)Cited by:[§1](https://arxiv.org/html/2609.22175#S1.p3.1)\.
- Assranet al\.\(2025\)M\. Assran, A\. Bardes, D\. Fan, Q\. Garrido, R\. Howes, M\. Komeili, M\. Muckley, A\. Rizvi, C\. Roberts, K\. Sinha, A\. Zholus,et al\.V\-jepa 2: self\-supervised video models enable understanding, prediction and planning\.arXiv preprint arXiv:2506\.09985\.Cited by:[§1](https://arxiv.org/html/2609.22175#S1.p4.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\),pp\. 15619–15629\.Cited by:[§1](https://arxiv.org/html/2609.22175#S1.p4.1)\.
- Ballet al\.\(2025\)P\. J\. Ball, J\. Bauer, F\. Belletti, B\. Brownfield, A\. Ephrat, S\. Fruchter, A\. Gupta, K\. Holsheimer, A\. Holynski, J\. Hron, C\. Kaplanis, M\. Limont, M\. McGill, Y\. Oliveira, J\. Parker\-Holder, F\. Perbet, G\. Scully, J\. Shar, S\. Spencer, O\. Tov, R\. Villegas, E\. Wang, J\. Yung, C\. Baetu, J\. Berbel, D\. Bridson, J\. Bruce, G\. Buttimore, S\. Chakera, B\. Chandra, P\. Collins, A\. Cullum, B\. Damoc, V\. Dasagi, M\. Gazeau, C\. Gbadamosi, W\. Han, E\. Hirst, A\. Kachra, L\. Kerley, K\. Kjems, E\. Knoepfel, V\. Koriakin, J\. Lo, C\. Lu, Z\. Mehring, A\. Moufarek, H\. Nandwani, V\. Oliveira, F\. Pardo, J\. Park, A\. Pierson, B\. Poole, H\. Ran, T\. Salimans, M\. Sanchez, I\. Saprykin, A\. Shen, S\. Sidhwani, D\. Smith, J\. Stanton, H\. Tomlinson, D\. Vijaykumar, L\. Wang, P\. Wingfield, N\. Wong, K\. Xu, C\. Yew, N\. Young, V\. Zubov, D\. Eck, D\. Erhan, K\. Kavukcuoglu, D\. Hassabis, Z\. Gharamani, R\. Hadsell, A\. van den Oord, I\. Mosseri, A\. Bolton, S\. Singh, and T\. RocktäschelGenie 3: a new frontier for world models\.External Links:Cited by:[§1](https://arxiv.org/html/2609.22175#S1.p1.1),[§1](https://arxiv.org/html/2609.22175#S1.p2.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\.Note:arXiv:2404\.08471; introduces V\-JEPACited by:[§1](https://arxiv.org/html/2609.22175#S1.p4.1)\.
- Boltonet al\.\(2025\)A\. Bolton, A\. Lerchner, A\. Cordell, A\. Moufarek, A\. Bolt, A\. Lampinen, A\. Mitenkova, A\. O\. Hallingstad, B\. Vujatovic, B\. Li, C\. Lu, D\. Wierstra, D\. P\. Sawyer, D\. Slater, D\. Reichert, D\. Vercelli, D\. Hassabis, D\. A\. Hudson, D\. Williams, E\. Hirst, F\. Pardo, F\. Hill, F\. Besse, H\. Openshaw, H\. Chan, H\. Soyer, J\. X\. Wang, J\. Clune, J\. Agapiou, J\. Reid, J\. Marino, J\. Kim, K\. Gregor, K\. Sridhar, K\. McKinney, L\. Kampis, L\. M\. Zhang, L\. Matthey, L\. Wang, M\. A\. Raad, M\. Loks\-Thompson, M\. Engelcke, M\. Kecman, M\. Jackson, M\. Gazeau, O\. Purkiss, O\. Knagg, P\. Stys, P\. Mendolicchio, R\. Hadsell, R\. Ke, R\. Faulkner, S\. Chakera, S\. S\. Baveja, S\. Legg, S\. Kashem, T\. Terzi, T\. Keck, T\. Harley, T\. Scholtes, T\. Roberts, V\. Mnih, Y\. Liu, Z\. Wang, and Z\. GhahramaniSIMA 2: a generalist embodied agent for virtual worlds\.External Links:2512\.04797,[Link](https://arxiv.org/abs/2512.04797)Cited by:[§1](https://arxiv.org/html/2609.22175#S1.p2.1)\.
- Bruceet al\.\(2024\)J\. Bruce, M\. Dennis, A\. Edwards, J\. Parker\-Holder, Y\. Shi, E\. Hughes, M\. Lai, A\. Mavalankar, R\. Steigerwald, C\. Apps, Y\. Aytar, S\. Bechtle, F\. Behbahani, S\. Chan, N\. Heess, L\. Gonzalez, S\. Osindero, S\. Ozair, S\. Reed, J\. Zhang, K\. Zolna, J\. Clune, N\. de Freitas, S\. Singh, and T\. RocktäschelGenie: generative interactive environments\.Note:arXiv:2402\.15391Cited by:[§1](https://arxiv.org/html/2609.22175#S1.p2.1)\.
- Chenet al\.\(2020\)T\. Chen, S\. Kornblith, M\. Norouzi, and G\. HintonA simple framework for contrastive learning of visual representations\.InInternational Conference on Machine Learning \(ICML\),Note:arXiv:2002\.05709Cited by:[§1](https://arxiv.org/html/2609.22175#S1.p3.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 \(NeurIPS\),Note:arXiv:2006\.07733Cited by:[§1](https://arxiv.org/html/2609.22175#S1.p3.1),[§3\.1](https://arxiv.org/html/2609.22175#S3.SS1.SSS0.Px1.p1.1),[§3\.1](https://arxiv.org/html/2609.22175#S3.SS1.SSS0.Px1.p1.3)\.
- Ha and Schmidhuber \(2018\)D\. Ha and J\. SchmidhuberWorld models\.arXiv preprint arXiv:1803\.10122\.Cited by:[§1](https://arxiv.org/html/2609.22175#S1.p1.1)\.
- Hafneret al\.\(2020\)D\. Hafner, T\. Lillicrap, J\. Ba, and M\. NorouziDream to control: learning behaviors by latent imagination\.External Links:1912\.01603Cited by:[§1](https://arxiv.org/html/2609.22175#S1.p2.1),[§1](https://arxiv.org/html/2609.22175#S1.p4.1),[§2](https://arxiv.org/html/2609.22175#S2.p1.1),[§3\.1](https://arxiv.org/html/2609.22175#S3.SS1.p1.1)\.
- Hafneret al\.\(2019\)D\. Hafner, T\. Lillicrap, I\. Fischer, R\. Villegas, D\. Ha, H\. Lee, and J\. DavidsonLearning latent dynamics for planning from pixels\.External Links:1811\.04551Cited by:[§1](https://arxiv.org/html/2609.22175#S1.p2.1),[§2\.1](https://arxiv.org/html/2609.22175#S2.SS1.p4.1)\.
- Hafneret al\.\(2021\)D\. Hafner, T\. Lillicrap, M\. Norouzi, and J\. BaMastering atari with discrete world models\.InInternational Conference on Learning Representations \(ICLR\),Note:arXiv:2010\.02193Cited by:[§1](https://arxiv.org/html/2609.22175#S1.p2.1)\.
- Hafneret al\.\(2025\)D\. Hafner, J\. Pasukonis, J\. Ba, and T\. LillicrapMastering diverse control tasks through world models\.Nature640,pp\. 647–653\.Note:Also arXiv:2301\.04104 \(2023\), titled "Mastering Diverse Domains through World Models"Cited by:[§1](https://arxiv.org/html/2609.22175#S1.p2.1)\.
- Heet al\.\(2020\)K\. He, H\. Fan, Y\. Wu, S\. Xie, and R\. GirshickMomentum contrast for unsupervised visual representation learning\.InIEEE/CVF Conference on Computer Vision and Pattern Recognition \(CVPR\),Note:arXiv:1911\.05722Cited by:[§1](https://arxiv.org/html/2609.22175#S1.p3.1)\.
- Hjelmet al\.\(2019\)R\. D\. Hjelm, A\. Fedorov, S\. Lavoie\-Marchildon, K\. Grewal, P\. Bachman, A\. Trischler, and Y\. BengioLearning deep representations by mutual information estimation and maximization\.External Links:1808\.06670Cited by:[§1](https://arxiv.org/html/2609.22175#S1.p3.1),[§2\.2](https://arxiv.org/html/2609.22175#S2.SS2.p3.1)\.
- Kipfet al\.\(2020\)T\. Kipf, E\. van der Pol, and M\. WellingContrastive learning of structured world models\.External Links:1911\.12247,[Link](https://arxiv.org/abs/1911.12247)Cited by:[§1](https://arxiv.org/html/2609.22175#S1.p4.1)\.
- LeCun \(2022\)Y\. LeCunA path towards autonomous machine intelligence\.Note:Open Review technical report, Meta AIIntroduces the Joint Embedding Predictive Architecture \(JEPA\)Cited by:[§1](https://arxiv.org/html/2609.22175#S1.p4.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:[§1](https://arxiv.org/html/2609.22175#S1.p4.1)\.
- Mazoureet al\.\(2020\)B\. Mazoure, R\. T\. des Combes, T\. Doan, P\. Bachman, and R\. D\. HjelmDeep reinforcement and infomax learning\.External Links:2006\.07217,[Link](https://arxiv.org/abs/2006.07217)Cited by:[§1](https://arxiv.org/html/2609.22175#S1.p3.1)\.
- Parker\-Holderet al\.\(2024\)J\. Parker\-Holder, P\. Ball, J\. Bruce, V\. Dasagi, K\. Holsheimer, C\. Kaplanis, A\. Moufarek, G\. Scully, J\. Shar, J\. Shi, S\. Spencer, J\. Yung, M\. Dennis, S\. Kenjeyev, S\. Long, V\. Mnih, H\. Chan, M\. Gazeau, B\. Li, F\. Pardo, L\. Wang, L\. Zhang, F\. Besse, T\. Harley, A\. Mitenkova, J\. Wang, J\. Clune, D\. Hassabis, R\. Hadsell, A\. Bolton, S\. Singh, and T\. RocktäschelGenie 2: a large\-scale foundation world model\.External Links:[Link](https://deepmind.google/discover/blog/genie-2-a-large-scale-foundation-world-model/)Cited by:[§1](https://arxiv.org/html/2609.22175#S1.p1.1),[§1](https://arxiv.org/html/2609.22175#S1.p2.1)\.
- Russellet al\.\(2025\)L\. Russell, A\. Hu, L\. Bertoni, G\. Fedoseev, J\. Shotton, E\. Arani, and G\. CorradoGAIA\-2: a controllable multi\-view generative world model for autonomous driving\.External Links:2503\.20523,[Link](https://arxiv.org/abs/2503.20523)Cited by:[§1](https://arxiv.org/html/2609.22175#S1.p1.1)\.
- Schwarzeret al\.\(2021\)M\. Schwarzer, A\. Anand, R\. Goel, R\. D\. Hjelm, A\. Courville, and P\. BachmanData\-efficient reinforcement learning with self\-predictive representations\.External Links:2007\.05929,[Link](https://arxiv.org/abs/2007.05929)Cited by:[§1](https://arxiv.org/html/2609.22175#S1.p3.1),[§3\.1](https://arxiv.org/html/2609.22175#S3.SS1.SSS0.Px1.p1.1)\.
- Stookeet al\.\(2021\)A\. Stooke, K\. Lee, P\. Abbeel, and M\. LaskinDecoupling representation learning from reinforcement learning\.External Links:2009\.08319,[Link](https://arxiv.org/abs/2009.08319)Cited by:[§1](https://arxiv.org/html/2609.22175#S1.p3.1)\.
- van den Oordet al\.\(2019\)A\. van den Oord, Y\. Li, and O\. VinyalsRepresentation learning with contrastive predictive coding\.External Links:1807\.03748Cited by:[§1](https://arxiv.org/html/2609.22175#S1.p3.1),[§2\.2](https://arxiv.org/html/2609.22175#S2.SS2.p2.3)\.
- Yanget al\.\(2024\)S\. Yang, Y\. Du, K\. Ghasemipour, J\. Tompson, L\. Kaelbling, D\. Schuurmans, and P\. AbbeelLearning interactive real\-world simulators\.External Links:2310\.06114,[Link](https://arxiv.org/abs/2310.06114)Cited by:[§1](https://arxiv.org/html/2609.22175#S1.p1.1)\.
- Yaratset al\.\(2021\)D\. Yarats, R\. Fergus, A\. Lazaric, and L\. PintoMastering visual continuous control: improved data\-augmented reinforcement learning\.External Links:2107\.09645,[Link](https://arxiv.org/abs/2107.09645)Cited by:[§1](https://arxiv.org/html/2609.22175#S1.p3.1)\.
- Yaratset al\.\(2020\)D\. Yarats, A\. Zhang, I\. Kostrikov, B\. Amos, J\. Pineau, and R\. FergusImproving sample efficiency in model\-free reinforcement learning from images\.External Links:1910\.01741Cited by:[§3](https://arxiv.org/html/2609.22175#S3.p1.1)\.
- Yeet al\.\(2026\)S\. Ye, Y\. Ge, K\. Zheng, S\. Gao, S\. Yu, G\. Kurian, S\. Indupuru, Y\. L\. Tan, C\. Zhu, J\. Xiang, A\. Malik, K\. Lee, W\. Liang, N\. Ranawaka, J\. Gu, Y\. Xu, G\. Wang, F\. Hu, A\. Narayan, J\. Bjorck, J\. Wang, G\. Kim, D\. Niu, R\. Zheng, Y\. Xie, J\. Wu, Q\. Wang, R\. Julian, D\. Xu, Y\. Du, Y\. Chebotar, S\. Reed, J\. Kautz, Y\. Zhu, L\. "\. Fan, and J\. JangWorld action models are zero\-shot policies\.External Links:2602\.15922,[Link](https://arxiv.org/abs/2602.15922)Cited by:[§1](https://arxiv.org/html/2609.22175#S1.p1.1)\.
- Zhanget al\.\(2021\)A\. Zhang, R\. McAllister, R\. Calandra, Y\. Gal, and S\. LevineLearning invariant representations for reinforcement learning without reconstruction\.External Links:2006\.10742Cited by:[§3](https://arxiv.org/html/2609.22175#S3.p1.1)\.
- Zhanget al\.\(2024\)L\. Zhang, Y\. Xiong, Z\. Yang, S\. Casas, R\. Hu, and R\. UrtasunCopilot4D: learning unsupervised world models for autonomous driving via discrete diffusion\.External Links:2311\.01017,[Link](https://arxiv.org/abs/2311.01017)Cited by:[§1](https://arxiv.org/html/2609.22175#S1.p1.1)\.Similar Articles
WCM: A World Critic Model for Vision-Language-Action Reinforcement Learning
Introduces WCM, a World Critic Model that jointly predicts future latent states and estimates values to improve temporal modeling for Vision-Language-Action reinforcement learning, achieving state-of-the-art results across robotic manipulation benchmarks.
Learning How the World Evolves: Extrapolative Video World Models via Latent Dynamics Reasoning
Introduces Latent Dynamics Reasoning (LDR), a video world model that integrates kinematic dynamics in a structured latent space, enabling extrapolation of learned dynamics far beyond training distributions while using far fewer parameters and running much faster than video diffusion baselines.
Predicting Consequences and Reinforcing Navigation Policies with Latent World Models
This paper proposes a Latent World Model (LWM) for robot navigation that predicts action-conditioned latent feature compatibility, enabling policy learning from unlabeled video data and reinforcement learning without additional environment interaction, outperforming existing methods.
Learning Visual Feature-Based World Models via Residual Latent Action
This paper introduces RLA-WM, a visual feature-based world model that leverages residual latent actions and flow matching to efficiently predict future visual states. The method outperforms existing video-diffusion and feature-based approaches while enabling novel robot learning techniques from offline, actionless demonstration videos.
No Gaussian Required: Contrastive Inverse Dynamics for JEPA World Models
The paper introduces AC-MTM, a contrastive inverse dynamics method to prevent encoder collapse in JEPA world models, achieving improved performance on multi-object tasks without Gaussian constraints.