TL;DR: A Temporal Discriminative objective increases representational quality, avoids trivial collapse, and improves long-horizon planning performance.
The standard recipe for JEPA-style world models [1] is by now familiar: an encoder, a next-latent predictor co-trained against it, and a regularizer to keep the encoder from collapsing to a trivial representation. Planning is L2 cost-to-go in latent space. LeWorldModel (LeWM) [2] is a canonical instance with SIGReg as its anti-collapse regularizer. SIGReg constrains the geometry of the latent distribution by pushing embeddings towards an isotropic Gaussian, which is an elegant solution to trivial collapse. Yet when we tested it, it planned poorly, and the representation turned out to be the cause.
SIGReg prevents collapse, but does not enforce information retention. While it successfully avoids trivial collapse, it does not incentivize the encoder to retain physical information. This is exactly what we observe when we train and evaluate on OGBench’s Scene environment [3].
Example start and end of trajectory in OGBench Scene.
The highly dynamic cube’s \((x,y,z)\) position and the 7-DOF arm were weakly represented within embeddings. This is in contrast to the slower-moving objects (drawer, window, buttons) that were cleanly recoverable. From this, we hypothesized that although anti-collapse is necessary, it is not sufficient to produce informative latents.
Held-out \(R^2\) from probing a SIGReg-trained encoder: the cube and 7-DOF arm are only weakly recoverable, while the drawer, window, and buttons decode almost perfectly.
Motivated by these results, we trace the problem to the training objective. Next-embedding MSE is minimized by making the future predictable, and the surest way to achieve that is to stop representing the parts of it that are hard to anticipate. The small, fast-moving, control-relevant variables (the arm’s precise pose, the cube) are exactly what a regressor would rather forget, because forgetting them lowers the loss. Clearly, the training objective needs to change.
Broadly speaking, a good world model should encode agent-relevant dynamics and simulate future states conditioned on candidate actions. The appeal of a JEPA-style latent world model is its ability to learn abstractions of the world without having to represent irrelevant details of the scenery. In order to do so, the objective must specify what types of information should be retained.
In light of this, we propose a Temporal Discriminative training objective. We leverage the natural temporal structure of action-labeled video to contrast the current state against the future, forcing the encoder to pick out the change. Further, the action conditioning of the predictor biases the model towards learning agent-relevant dynamics.
With our Temporal Discriminative World Model (TDWM), we are able to train stably and avoid trivial collapse with only one training objective. The intrinsic anti-collapse from temporal discrimination is qualitatively different from variance-based regularization: our latent spaces have an effective rank of around 40, while SIGReg expands to around 150. For reference, the environment contains a total of 24 degrees of freedom.
Following the basic JEPA recipe, we co-train an encoder \(f_\theta\) and a predictor \(g_\phi\). The encoder embeds each frame \(o_i\) of a clip, \(z_i = f_\theta(o_i)\), and from a 3-frame history the predictor predicts the embedding \(k\) steps ahead using the intermediate actions,
\[\hat z_{t+k} = g_\phi\big(z_{t-2:t},\; a_{t:t+k-1},\; k\big).\]Both encoder and predictor outputs are norm-constrained, so only their direction is free. Writing \(c_{ij} = \text{cossim}(\hat{z}_i, z_j)\) for the cosine similarity between a prediction with target frame \(i\) and an encoded frame \(j\) of the same clip, we regress every such cosine similarity onto a target similarity that decays with the temporal gap \(\lvert i-j\rvert\):
\[\mathcal{L}=\mathbb{E}_{i,j\sim \text{clip}}\left[\left(c_{ij}-e^{-|i-j|/\sigma}\right)^2\right].\]Here, \(\sigma\) determines how fast target similarity decays with temporal distance. At \(i=j\) this reduces to a next-embedding prediction objective in a norm-constrained space, while the partial-similarity targets at \(i \neq j\) push temporally distant frames towards orthogonality (anti-collapse). Cosine similarity is a natural measure here as it allows us to specify a degree of similarity rather than an explicit target latent.
Importantly, with exponential attenuation these target similarities can all be satisfied at once. Write \(\tau_{ij} = e^{-\lvert i-j\rvert/\sigma}\) for the target similarity between frames \(i\) and \(j\), so the objective regresses each realized cosine \(c_{ij}\) onto \(\tau_{ij}\). For frames ordered \(i \le j \le k\),
\[\tau_{ik} = e^{-|i-k|/\sigma} = e^{-|i-j|/\sigma}\cdot e^{-|j-k|/\sigma} = \tau_{ij}\cdot \tau_{jk},\]so a long-gap target is just the product of the short-gap targets. More formal analysis to come when we find a competent mathematician.
Example trajectory in OGBench Scene.
By temporal distance, we mean how far two states are separated in time rather than in physical space. Latent planning uses the Cross-Entropy Method (CEM), a sampling-based planner that repeatedly draws candidate action sequences, scores each by its predicted latent cost-to-go, and refits the sampling distribution to the lowest-cost elite set until it converges on a plan [4]. The cost-to-go is determined by the choice of latent geometry, and here temporal distance is a more suitable measure than LeWorldModel’s embedding distance.
As in the figure above, the arm is in its goal state while the drawer needs to be opened. With a latent space like LeWorldModel’s, the start and goal are already low cost as most of the scene is already correctly positioned. Opening the drawer requires moving the arm out of position, which temporarily increases the cost. In TDWM’s latent space, states are ordered by their distance in time, which means that moving the arm towards the drawer smoothly decreases cost as this action decreases the time until the goal is achieved.
Noting that 97% of training FLOPs are allocated to the encoder and that we sample long contiguous blocks of frames, we employ multi-horizon prediction to increase gradient quality substantially while modestly increasing the compute cost. In addition, we find that using a fused Muon-AdamW optimizer substantially accelerates convergence.
Normalized predictor fidelity as a function of out-of-distribution action perturbation magnitude (\(\mathrm{RMS}(\Delta a / \sigma_a)\)). Trained on expert trajectories only, TDWM degrades sharply on OOD actions; training on the noised dataset largely closes the gap. We report perturbation size as the RMS of the per-dimension z-scored action deviation so that every action dimension contributes on equal, unit-free footing.
Interestingly, when training only on expert data, LeWM outperforms TDWM on planning. To disentangle this, we test both predictors with off-policy (noised) actions. TDWM’s prediction accuracy degrades sharply on actions outside of the expert distribution (we return to this in The off-manifold tradeoff). We hypothesize that LeWM’s robustness can be attributed to it encoding less of the actual scene. In this way, it has a simpler extrapolation task. Further, we note that in the expert trajectories, actions correlate linearly with state to an \(R^2\) of 0.71. TDWM may therefore be modeling agent intent rather than the underlying physics of the environment.
To address this, we supplement the dataset by re-rolling the trajectories with added action noise and re-rendering. This brings action correlation with state down to an \(R^2\) of 0.31. We argue that world models should be trained on data including both expert and unsuccessful trajectories, as this forces the world model to understand the underlying physics of the world. Importantly, this kind of data is readily available. Recent embodied datasets retain failed trajectories at meaningful scale: DROID [5] releases roughly 16k failed teleoperation attempts alongside its 76k successes; RoboMIND [6] contributes a further 5k real-world failure demonstrations, each labeled with its cause; and AgiBot World [7] preserves failed demonstrations within its 1M-trajectory corpus. Alongside these public datasets, as embodied intelligence becomes deployed at scale, failed trajectories will be abundant.
We trained every model on this dataset with the recipe otherwise unchanged. Both the TDWM and LeWM variants are trained for 0.8 exaFLOPs and the presented results come from the checkpoint with the highest O25 CEM score.
We evaluate TDWM along three axes: what its latents encode (probing), whether anti-collapse emerges without explicit regularization (SIGReg loss), and how it plans (CEM success rate at long goal horizons). All comparisons are against LeWM trained on the same noised dataset and we stick to pure 1-step auto-regressive prediction for a fair comparison.
We probe representation quality by fitting supervised regressors from embeddings to ground-truth object states (cube/arm/drawer/window/button positions from the simulator’s qpos). Here, we report held-out R² on an 80/20 train-test split. Both a linear ridge probe (\(\alpha\)=10) and a 2-layer MLP (512d and 256d hidden dimension, ReLU, early-stopped) are fit this way.
The probe results separate cleanly into two regimes. On the static scene elements (drawer, window, and buttons), both models score between 0.91 and 0.97 and are effectively tied. On the dynamic variables, the gap is large: TDWM reaches \(R^2 = 0.56\) on cube position versus LeWM’s 0.27 (2.1×), and 0.69 versus 0.39 on the 7-DOF arm (1.8×).
Information about static geometry is easily predictable and thus survives mere anti-collapse pressure. The arm and cube are harder to predict yet are the most reliable discriminant between adjacent frames, pushing TDWM to represent them more cleanly.
TDWM is trained with no SIGReg term, yet its final SIGReg loss falls from ~30 to 2.4. LeWorldModel achieves a SIGReg loss of 0.6 with direct optimization. Clearly, explicit regularization is not necessary for avoiding collapse. The discriminative objective enforces enough decorrelation through its orthogonality targets to keep the representation healthy on its own.
CEM success rate at goal offsets O25, O50, and O75. Note that O75 is 3Ă— the horizon evaluated in [2]. Error bars are standard error across seeds.
TDWM outperforms LeWM at every horizon: 63% vs. 50% at O25, 38% vs. 28% at O50, and 15% vs. 4% at O75. The absolute margin is broadly stable, but the relative gap widens with horizon. Long-horizon planning is the regime in which a missing cube position compounds worst, and it is also where the representational advantage pays out most.
Planning success at O25 across CEM ablations: LeWM (50%), TDWM expert + noise (63%), TDWM expert with a perfect (MuJoCo) predictor (66%), and a perfect encoder with a perfect predictor (77%, the absolute upper bound).
As discussed in Data, the richer representation degrades predictor robustness out of distribution. To better understand where the models fail, we perform ablations on each component of CEM planning. As shown in the figure above, taking the expert-trained encoder and using the ground-truth MuJoCo environment as the predictor module, we get a CEM score of 66% (isolating the perfect-predictor upper bound). By training on noised action trajectories, the trained predictor narrows this gap to 63%. As context, performing CEM directly in the simulator’s representation space (qpos) with its ground-truth physics simulator, we achieve a maximal CEM score of 77%, an absolute upper bound for performance under our eval configuration.
This is preliminary work and should be read as such. Our evaluation is on a single environment (the Scene manipulation task). We are actively expanding to more environments and evaluating scalability of these ideas.
The standard recipe for latent world models is misspecified: predicting the future and retaining information in a JEPA architecture are conflicting goals. Anti-collapse may stop the encoder from collapsing, but does not make it informative.
TDWM replaces that objective. Discriminating between temporally adjacent states forces a focus on dynamics, and the predictor’s action conditioning points the gradient at agent-relevant dynamics. The result is a representation that selectively preserves the agent-controlled variables (cube position, arm pose) and outperforms previous world models at long-horizon planning.
If you’re in industry and find this work interesting, feel free to reach out! We’re open to collaborations and are looking for resources to scale.
We would like to thank Tiange Xiang, McNair Shah, Eshaan Barkataki and Yash Ranjith for insightful discussions and proofreading.