Next Embedding Prediction Makes World Models Stronger

The pixel decoder is a tax

A world model earns its keep by carrying a compact state the agent can plan in. The Dreamer line of agents learns that state with a pixel decoder: reconstruct the current frame, and trust that whatever survives compression is useful. Reconstruction is dense, stabilizing supervision, but it bills the model for every texture and background detail whether or not the task cares.

Decoder-free variants drop the generative head, yet most still anchor the representation with a same-timestep signal: agree with the current frame's embedding, under augmentations or redundancy reduction. Under partial observability that is not the property we need. The state at time $t$ must be predictive of what happens at $t+1$, and nothing in a same-step objective asks for it.

Predict the next embedding instead

NE-Dreamer keeps the standard machinery. An encoder produces $e_t=f_{\mathrm{enc}}(x_t)$, a recurrent state-space model carries a deterministic state $h_t=f_{\mathrm{rec}}(h_{t-1},z_{t-1},a_{t-1})$ alongside a stochastic latent $z_t$, and the reward head, continuation head, and prior-posterior KL stay untouched. What we remove is the decoder; what we add is a small causal transformer that reads the history and predicts the next embedding, $\hat{e}_{t+1}=T_{\theta}(h_{\le t},z_{\le t},a_{\le t})$.

The prediction is pulled toward a frozen target $\mathrm{sg}(e_{t+1})$ with a Barlow Twins loss, $\mathcal{L}_{\mathrm{NE}}=\sum_i(1-C_{ii})^2+\lambda_{\mathrm{BT}}\sum_{i\neq j}C_{ij}^2$, where $C$ cross-correlates predicted and target embeddings over the batch. Driving the diagonal to one makes the prediction informative; suppressing the off-diagonal keeps the embedding from collapsing onto a few redundant directions. No pixels, no augmentations, no extra regularizers.

The world-model objective is then $\mathcal{L}_{\mathrm{wm}}=\mathcal{L}_{\mathrm{rew}}+\mathcal{L}_{\mathrm{cont}}+\beta_{\mathrm{kl}}\mathcal{L}_{\mathrm{kl}}+\beta_{\mathrm{ne}}\mathcal{L}_{\mathrm{NE}}$, and the transformer is genuinely small: two layers, four heads, hidden size 256. Everything else, including every actor-critic hyperparameter, is DreamerV3's, so any gain has to come from the representation objective.

One time index does the work

The design has two ingredients that look minor and are not: the sequence model and the target's time index. Aligning to the current embedding, as R2-Dreamer does, checks that the state describes the present. Shifting the target one step forward turns the same loss into a forecast, and a forecast is exactly the property a world model is supposed to have.

Our ablations on DMLab make the attribution clean. Remove the causal transformer and performance collapses across all four tasks. Keep the transformer but train against the same-step target and nearly the entire gain disappears. Remove the lightweight projector in front of the transformer and little changes beyond optimization speed. The mechanism is the pair: causal sequence modeling plus the next-step shift.

Where the objective pays

DeepMind Lab's Rooms tasks are the regime this objective was built for: first-person 3D navigation with sparse rewards, where success means remembering what you saw many steps ago. We train every agent for 50 million environment steps at matched capacity, about 12 million parameters, five seeds each, in one unified codebase.

The aggregate actually understates how one-sided the per-task story is. On Select Nonmatching Object our return reaches roughly 61 while every baseline stays near one; on Watermaze roughly 12 against about one. These are tasks where the baselines do not underperform so much as fail to leave the floor.

On the DeepMind Control Suite the same agent stays in the leading pack: our task-mean return lands around 800 after one million steps, alongside DreamerV3, R2-Dreamer, and DreamerPro, while Dreamer without any representation objective and the model-free DrQ-v2 trail well behind. Dropping reconstruction costs nothing on standard continuous control.

A diagnostic makes the mechanism visible. We train a post-hoc decoder on frozen latents, never used during agent training, and reconstruct what the state remembers. Our latents keep task-relevant objects and room layout stable across time; Dreamer's and R2-Dreamer's show them appearing for a moment and then fading, even when the underlying scene has not changed.

What we did not test

This is a workshop-scale study, and its boundary is visual fidelity: our environments reward long-term structure rather than fine detail, and whether next-embedding prediction can fully replace reconstruction when pixels genuinely matter remains open. Barlow Twins is also just our instantiation; any collapse-resistant alignment loss could slot into the same objective.

The takeaway we will defend is narrower and sturdier: when the world is partially observable, train the state to predict the next embedding rather than to redraw the present. One index in the target, plus a two-layer transformer, is what separated flatlining agents from ones that solve the task.