Predictive and Generative Objectives › Reconstruction and Likelihood Objectives

The most natural way to train a world model is to ask it to reproduce observations. If the model sees a frame of video, it should be able to output a frame of video that matches what will actually happen next. If the observations are a robot's joint-angle measurements, the model should output numbers close to the next measurement. Reconstruction and likelihood objectives formalize this intuition.

This family turns learning how the world works into a differentiable regression or classification problem with dense targets. It avoids designing a separate pretext task, but still requires choices about the likelihood, decoder, latent bottleneck, and reduction across observation dimensions. Those choices determine which accurate reconstructions are cheap or expensive for the model.

Suppose the world model produces a distribution $p_\theta(o_{t+1} \mid o_{\le t}, a_{\le t})$ over the next observation. If we assume a Gaussian form with fixed variance $\sigma^2 I$, then the negative log-likelihood of a target $o_{t+1}$ is

$$-\log p_\theta(o_{t+1} \mid \cdot) = \frac{1}{2\sigma^2} \lVert o_{t+1} - \hat{o}_\theta(\cdot) \rVert_2^2 + \text{const}.$$

where:

• $o_{t+1}\in\mathbb{R}^{D}$: the observed target
• $\hat{o}_\theta(\cdot)\in\mathbb{R}^{D}$: the predicted Gaussian mean conditioned on the history
• $\sigma^2 I$: the fixed isotropic covariance
• $\theta$: model parameters
• $\text{const}$: the Gaussian normalization constant, which is independent of $\theta$ when $\sigma$ is fixed
• $-\log p_\theta(o_{t+1}\mid\cdot)$: the negative log-likelihood of the target under the model
• $D$: dimensionality of the observation vector
• $I$: identity matrix
• $\cdot$: shorthand for the conditioning history $o_{\le t},a_{\le t}$
• $\lVert\cdot\rVert_2^2$: squared Euclidean norm

Because $\sigma$ is fixed, maximizing the Gaussian likelihood over $\theta$ is equivalent to minimizing the negative log-likelihood. The factor $1/(2\sigma^2)$ is positive and independent of $\theta$, and the additive constant also does not depend on $\theta$. Therefore the maximum-likelihood problem reduces to minimizing $\lVert o_{t+1}-\hat{o}_\theta(\cdot)\rVert_2^2$, which is the mean-squared error (MSE). Intuitively, the negative log-likelihood penalizes a prediction in proportion to how far it falls from the target, scaled by how confident the fixed-variance Gaussian is; because the confidence scale $1/(2\sigma^2)$ is constant, it cannot change which parameters are optimal, only the numeric size of the loss. This is why MSE is a common reconstruction loss in world-model work: it is the log-likelihood of a specific fixed-variance probabilistic model. If instead the model outputs a distribution over discrete tokens (for example, a vocabulary of image codes from a VQ-VAE), the loss is the categorical cross-entropy

$$-\log p_\theta(k \mid \cdot) = -\log \mathrm{softmax}(z_\theta)_k.$$

Here $\mathrm{softmax}$ converts the logits into a probability vector, so $-\log$ of the entry assigned to the observed token measures how surprised the model is by the true token; minimizing it pushes probability mass onto the correct token. The variables are:

• $k$: the observed token index
• $z_\theta\in\mathbb{R}^{V}$: logits for a vocabulary of size $V$
• $\mathrm{softmax}(z_\theta)_k$: the probability assigned to the observed token
• $V$: the vocabulary size
• $\theta$: model parameters
• $-\log p_\theta(k\mid\cdot)$: the categorical cross-entropy loss
• $\cdot$: shorthand for the conditioning history

This is again the log-likelihood of the assumed generative model. Autoregressive token world models, discussed in Part V on world-model architectures, train with exactly this cross-entropy over token sequences. The two cases share a common logic: pick a parametric family for the observation, write down its likelihood, and let maximum likelihood pick the parameters. The Gaussian case yields a squared error and the categorical case yields a cross-entropy, but both are the same principle applied to a different choice of distribution.

The distinction in the callout above matters more than it first appears. When the variance is learned rather than fixed, the model can lower the negative log-likelihood either by predicting a better mean or by inflating its uncertainty, and the two are entangled in a way that a plain MSE cannot express. A learned-variance model can therefore achieve a lower likelihood while producing worse point predictions, which is a concrete example of the split between predictive fidelity and uncertainty quality. If you care about calibrated uncertainty, you have to look past the aggregate likelihood and inspect whether the model's stated confidence matches its realized error.

Many world models infer a latent state $z_t$ and reconstruct observations from it. If the observation at each time step is generated by a latent that evolves according to an action-conditioned transition model, the marginal likelihood $\log p_\theta(o_{1:T} \mid a_{1:T-1})$ requires integrating over the entire latent trajectory. That integral is intractable in general because the latent trajectory is continuous and high-dimensional, so we cannot simply maximize the marginal likelihood directly. Variational models sidestep the integral by introducing an approximate posterior over latent trajectories and optimizing a tractable lower bound on the marginal likelihood, the evidence lower bound (ELBO):

$$\log p_\theta(o_{1:T} \mid a_{1:T-1}) \ge \mathbb{E}_{q_\phi(z_{1:T} \mid o_{1:T}, a_{1:T-1})} \left[ \log p_\theta(o_{1:T} \mid z_{1:T}) \right] - D_{\mathrm{KL}}\!\left(q_\phi(z_{1:T} \mid o_{1:T}, a_{1:T-1})\,\|\,p_\theta(z_{1:T} \mid a_{1:T-1})\right).$$

where:

• $o_{1:T}$: observation trajectory
• $z_{1:T}$: latent trajectory
• $a_{1:T-1}$: actions between consecutive states
• $q_\phi$: approximate posterior
• $p_\theta$: the generative model (with parameters $\theta$); here it provides the observation likelihood $p_\theta(o_{1:T}\mid z_{1:T})$ and the action-conditioned latent prior $p_\theta(z_{1:T}\mid a_{1:T-1})$
• $D_{\mathrm{KL}}(\cdot\|\cdot)$: Kullback-Leibler divergence
• $\log p_\theta(o_{1:T}\mid a_{1:T-1})$: marginal log-likelihood of the observation trajectory
• $T$: number of time steps
• $\phi$: variational parameters of the approximate posterior
• $\mathbb{E}$: expectation under the distribution indicated by the subscript
• $p_\theta(o_{1:T}\mid z_{1:T})$: observation likelihood (decoder)

The expectation is taken under $q_\phi$. The first term is expected reconstruction log-likelihood. The KL direction $D_{\mathrm{KL}}(q_\phi\|p_\theta)$ penalizes posterior trajectories that the learned prior cannot explain. This is the variational RNN structure that PlaNet and Dreamer use, with practical modifications. As we discussed in Recurrent State-Space Models, RSSMs use ELBO-related reconstruction and latent regularization terms together with reward prediction.

The two terms of the ELBO play complementary roles that are worth separating. The reconstruction term rewards the latent trajectory for containing enough information to explain the observations, which pushes the posterior to encode everything the decoder needs. The KL term rewards the prior for being able to generate the posterior's latent trajectories on its own, which pushes the prior to capture the action-conditioned dynamics without peeking at future observations. During open-loop imagination, after conditioning on observations available so far, future posterior updates are unavailable because those observations have not happened yet. The rollout is therefore only as good as its learned prior dynamics. Training the two terms together encourages the posterior states to remain predictable by that prior. Without the KL term, the posterior could encode information that the prior cannot reproduce, and rollouts could fail even when reconstruction looks excellent.

A useful measure when comparing variants is the negative log-likelihood in bits per dimension. This normalizes the likelihood both by the number of scalar dimensions $D$ and by the factor $\log 2$ needed to convert nats into bits, so that models over observations of different dimensionality can be compared on a common scale:

$$\mathrm{BPD} = \frac{-\log p_\theta(o \mid \cdot)}{D \log 2}.$$

where:

• $\mathrm{BPD}$: bits per dimension
• $D$: the number of modeled scalar dimensions
• $\log 2$: the conversion factor from nats to bits; the logarithm in the numerator is natural
• $p_\theta(o\mid\cdot)$: the model likelihood of observation $o$
• $\cdot$: shorthand for the conditioning history

Dividing by $D$ and $\log 2$ reports the average negative log-likelihood per dimension in bits. Comparisons are meaningful only when the observation representation, quantization or density convention, and likelihood family are held fixed; BPD is not directly comparable to an unnormalized MSE. BPD measures predictive fidelity, not representation quality or decision usefulness, so a model can win on BPD while carrying a latent state that is useless for control. The number is a useful sanity check on the generative side of the model, not a certificate that the model can be planned with.

Reconstruction has an appealing data-efficiency property: every pixel of every frame provides a training signal. Yet many pixel variations can be irrelevant to the dynamics or decision. In a driving scene, road texture, illumination, or a distant cloud may consume representational capacity without improving the variables needed for control. Reconstruction still rewards those details when they are predictable. This creates a potential mismatch between generative fidelity and predictive sufficiency.

There are three consequences worth internalizing:

• Capacity dilution. A large fraction of the model's parameters may be dedicated to modeling high-frequency texture that a downstream controller will never look at.
• Over-sensitivity to nuisance factors. Small perturbations in lighting or rendering that a reconstruction objective treats as important can dominate the loss and drive the latent state in ways that have no bearing on dynamics.
• Misleading leaderboards. A model that achieves beautiful reconstructions may have a latent state that is nearly useless for planning, while a model with ugly reconstructions may have the state a controller actually needs.

Each consequence reflects the same underlying issue: matching the observation distribution is not identical to learning a state sufficient for dynamics and decisions. Unpredictable flicker contributes irreducible observation uncertainty; predictable but task-irrelevant appearance can still earn likelihood. Whether either one harms the learned state depends on model capacity, inductive bias, and how the representation is shared.

A second problem is posterior collapse in latent-variable models. If the decoder can model observations without using the latent, the reconstruction objective gives the encoder little reason to transmit information. The approximate posterior then approaches the prior, the KL term approaches zero, and the latent state becomes uninformative. KL warm-up, free bits, decoder restrictions, and related training choices can mitigate this failure [@kingma2016improved]. Dreamer's KL balancing, discussed in the last technical section, controls a related temporal optimization tradeoff.

An aggregate loss curve may not expose posterior collapse. A flexible decoder can reduce reconstruction error while ignoring the latent; the encoder then gains little from encoding the observation, the posterior approaches the prior, and the KL term goes to zero. The result reconstructs reasonably yet carries little observation-specific information in its latent state. KL warm-up initially reduces pressure to match the prior, while free bits stops charging the encoder for a small per-group KL budget. Neither mechanism guarantees informative latents. Simply increasing the KL weight, as in the usual $\beta>1$ interpretation of $\beta$-VAE, can strengthen pressure toward the prior rather than cure collapse [@higgins2017beta; @kingma2016improved].

It is worth separating two ways of using a likelihood objective:

• Explicit generative models learn $p_\theta(o)$ or $p_\theta(o \mid z)$. Autoregressive token models, diffusion and flow models, and VAEs fall in this family. They can be sampled from to produce synthetic observations.
• Reconstruction as a proxy for representation learning uses a likelihood or MSE head to shape the latent state, while downstream planning may operate entirely in latent space. PlaNet, Dreamer, and many world-model variants use observation prediction alongside latent dynamics and task heads.

The distinction matters because the second family can use signals other than full observation likelihood to shape the latent state. Contrastive and masked objectives provide two alternatives.

In the first family, the generative model itself is the deliverable: you might want to sample plausible futures, generate training data, or evaluate how likely an observation is under the model. In the second, an observation decoder may be used mainly as a training head while planning consumes the encoder and transition state; some systems still retain the decoder for diagnostics or auxiliary uses. This is why two papers can use the same nominal loss and get different downstream behavior: one evaluates the generative output directly, while another uses reconstruction chiefly as a representation-shaping signal. When reading a world-model paper, check which components are used at deployment and which metrics correspond to that use.