Predictive and Generative Objectives

Imagine building a robot that has to catch a bouncing ball. Before it can move its arm, it needs some internal sense of where the ball is going. That internal sense, whatever form it takes, is a world model. The whole point of a world model is to anticipate what happens next, and the way we teach a machine to anticipate is to give it an objective: a scalar number that measures how well its predictions match the world. Everything else in the field, the architectures, the latent codes, the planning loops, ultimately serves that single scalar.

This chapter is about the objective function. It is easy to treat the choice of loss as a detail left to whoever writes the training code, but that choice silently determines what the model learns to care about, which aspects of the world it will represent, and which it will ignore. Two models with identical architectures can end up behaving very differently because one was trained to predict pixels and the other was trained to predict a compact latent state. One may produce gorgeous rollouts of a static scene while being useless for control, and the other may look ugly but drive a car successfully. The objective is the model's job description.

It helps to separate several meanings of a good world model. Predictive fidelity asks whether outputs match future observations, which reconstruction and likelihood measure directly. Representation quality asks whether the latent state captures task-relevant variables, independent of whether they can be rendered into pixels. Uncertainty quality asks whether probabilistic predictions are calibrated. Decision usefulness asks whether a planner or controller built on the model achieves good outcomes. Their priority depends on the application: control centers decision usefulness, while simulation or generation may center fidelity and calibration. A loss usually optimizes only part of this set directly, so performance on one axis does not establish performance on another.

Part V surveyed the families of world-model architectures: recurrent state-space models, transformers, autoregressive token models, diffusion and flow models, and predictive-embedding models like JEPA. Part VI turns to the question that those architectures leave open: what do we actually ask them to compute, and how do we prevent the training signal from building representations that look great on a leaderboard yet collapse the moment the agent has to plan? This chapter builds the vocabulary and the practical intuition for predictive and generative objectives. Along the way it reveals a recurring tension that will follow us through the rest of the book: the loss that is easiest to optimize need not produce a control-relevant state.

We start with the oldest and most direct family of objectives, reconstruction and likelihood, and see why they are so appealing and so treacherous. We then examine contrastive and masked prediction, which trade pixel fidelity for representation quality. From there we move to multi-step and overshooting losses, the machinery that turns a one-step predictor into a rollout engine. Finally, we look at how to balance representation, dynamics, and reward losses when a single network is asked to do several jobs at once.

Fitting a dynamical model by reducing prediction error predates modern machine learning. System identification developed statistical methods for choosing parameters from discrepancies between observed and predicted outputs [@ljung1999system]. A common criterion is one-step-ahead prediction error: condition on the available measured history, predict the next output, and penalize the residual. Its consistency and efficiency depend on assumptions about model structure, disturbances, excitation, and the estimator.

Related state-space tools made their modeling assumptions explicit. The Kalman filter recursively updates a Gaussian state estimate for linear dynamics with Gaussian noise, and the linear-quadratic regulator is optimal for its specified linear dynamics and quadratic cost. Learned world models replace those restricted maps with nonlinear, high-dimensional functions, but still face choices about state, uncertainty, and the gap between fitting and use.

Classical system identification also taught us a subtle warning. Under model misspecification, parameters that minimize a one-step prediction criterion need not minimize free-running simulation error, especially when the dynamics are weakly contracting or unstable. The one-step criterion evaluates predictions conditioned on measured history, whereas simulation repeatedly feeds the model's own outputs back into its state. Prediction-error and output-error methods make different assumptions about that conditioning and the disturbance model [@ljung1999system]. This distinction is an ancestor of modern multi-step and overshooting losses.

The reason one-step error can mislead is worth stating plainly. A teacher-forced one-step objective conditions each prediction on measured history, but the fitted transition may later be used to propagate a state for many steps without corrective observations. A parameter that slightly improves the conditional one-step residual while biasing the free-running dynamics can therefore win the one-step comparison and lose the rollout comparison. Output-error and simulation-error criteria instead evaluate trajectories generated by the fitted dynamics, bringing the training criterion closer to free-running use. Their assumptions and statistical properties differ, so neither is a universal fix; the important point is that the deployment mode belongs in the objective design.

Learned world models inherited this story when neural networks replaced linear and low-order parametric models. Different lines of work now make different objective choices: PlaNet and Dreamer combine observation, reward, and latent-dynamics terms; short model rollouts limit the region in which learned dynamics are trusted; and JEPA-style models predict learned target embeddings rather than pixels. These choices respond to recurring practical problems: observation losses can spend capacity on nuisance detail, teacher-forced accuracy may not survive recursive rollout, and a pixel metric need not align with the variables that matter for a decision.

As Part I: The World-Model Idea and Part V emphasized, an objective does more than fit outputs: it shapes the state available to the agent. The rest of this chapter asks when that state is sufficient for the intended predictions and decisions.

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

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

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):

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:

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.

Contrastive and masked objectives shape a representation through recognition or prediction tasks rather than requiring an unconditional full-image reconstruction. Contrastive and embedding-target methods need not reward pixel fidelity at all. Raw-target masked autoencoders do reconstruct masked pixels, so their pressure toward texture depends on the target, mask, and decoder design.

This is a deliberate trade. Instead of always defining the task as "produce the complete observation," these objectives ask the model to discriminate candidates or infer hidden targets. An embedding target can omit texture, whereas a raw-pixel target can retain it. The representation is therefore aligned with whatever information the chosen pretext task rewards. The price is that many such objectives lack a direct likelihood interpretation and are sensitive to how positives, negatives, masks, and targets are constructed.

Contrastive learning builds a representation by pulling together embeddings of similar examples and pushing apart different ones. The InfoNCE loss, introduced in the context of contrastive predictive coding, has the form

The numerator measures how similar the anchor is to its positive; the denominator sums the similarities to every candidate, including the positive. The ratio is therefore a softmax probability that the positive is the correct match, and the loss is the negative log of that probability. The temperature $\tau$ sharpens or flattens the distribution. The variables are:

• $\mathcal{L}_{\mathrm{InfoNCE}}$: the InfoNCE loss
• $x$: the anchor
• $x^+$: its positive
• $\{x_i\}_{i=1}^{N}$: candidate set including that positive and $N-1$ negatives
• $f_\theta$: the encoder
• $\mathrm{sim}$: a similarity such as cosine similarity
• $\tau>0$: the temperature [@oord2018representation]
• $N$: number of candidates

The loss is an $N$-way classification objective: identify the positive among all candidates.

Read as classification, the model must identify the positive future or view among candidates. A low temperature $\tau$ sharpens differences in their logits; a high temperature flattens them. Temperature therefore changes both the probability distribution and gradient scale, and should be selected together with similarity normalization and the negative set.

The temperature reshapes the InfoNCE softmax without changing the underlying similarities, and fixing one similarity profile makes that effect concrete.

In the temporal setting, positives may be future samples paired with their preceding context, while other samples serve as negatives. Contrastive predictive coding (CPC) trains a context representation to identify future latents under this loss [@oord2018representation]. The objective biases the representation toward features that help distinguish predictable future content; it does not directly require pixel-level reconstruction.

Why can this favor predictable structure? A feature of the future that can be inferred from context helps identify the correct candidate. Unpredictable variation is less useful for that discrimination, although finite data and correlations can still let the encoder represent it. This is often a useful bias for dynamics modeling, but "predictable from the past" is not the same as "relevant to the task." A contrastive representation can preserve predictable nuisance variables or omit variables that matter for control but not for its pretext task.

Contrastive approaches have two practical drawbacks. First, they need negatives, and the choice of negatives controls what the representation learns. Too-easy negatives give a weak discrimination task; false or ambiguous negatives can give a misleading gradient. Second, small batches provide fewer in-batch negatives and can weaken the signal. MoCo's queue and SimCLR's large batches address the supply of negatives in different ways [@he2020momentum; @chen2020simple]. Non-contrastive alternatives such as BYOL instead train an online network to predict the output of a slowly updated target network [@grill2020bootstrap], a pattern related to later joint-embedding predictive methods.

Negative construction is therefore part of the modeling problem rather than a minor implementation detail. A semantically valid match treated as a negative creates a false-negative error; a trivially different negative supplies little pressure to represent fine structure. Queues, large batches, and sampling rules trade off diversity against false-negative risk. Non-contrastive methods remove explicit negative sampling, but introduce their own requirements for target updates, asymmetry, or regularization.

The naive matching objective admits a constant solution: if both branches emit the same vector for every input, their discrepancy is zero. Successful methods combine stop-gradient with architectural or update asymmetry, and some add explicit anti-collapse statistics. The exact mechanism differs across BYOL, SimSiam, and JEPA-family implementations, so the safeguards should be evaluated as a system rather than attributed to one operator.

Masked prediction asks the model to reconstruct a hidden portion of the input from the visible portion. BERT popularized the approach for text [@devlin2018bert], and MAE adapted it to images, achieving strong representation quality by masking a large fraction of patches and reconstructing only those patches [@he2021masked]. In the world-model context, masked prediction has become a natural fit because the inputs and targets can live in different spaces. You can mask out a future observation and ask the model to predict its embedding, without ever producing a pixel.

The design space is captured by three questions:

• What is masked? Tokens of text, patches of an image, frames of a video, or timesteps of a trajectory. In a world model, masking time steps forces the model to integrate information from the past and future to fill in what it has not seen.
• What is the target? The raw masked content (as in MAE) or an embedding of it produced by a target encoder (as in JEPA). The latter avoids spending capacity on reconstructing unpredictable detail.
• What is the loss? Typically MSE in embedding space, or cross-entropy over token identities for discrete data.

Masked prediction prevents direct access to a selected target region. In temporal data, a suitable mask can require the model to infer future content from earlier observations and actions rather than copy the target. The strength and relevance of that bottleneck depend on what is masked and what remains visible.

A reconstruction objective gives the model credit for every reproduced pixel, including pixels copied from visible context. A masked objective withholds its target region, so reducing its loss requires inferring that region from the remaining context. In a temporal world model, the analogous task is to infer a hidden future from the past and actions. Whether this matches deployment depends on how the mask, context, and prediction target are constructed.

Joint-Embedding Predictive Architectures (JEPAs) combine masked prediction with a learned target embedding [@lecun2022path]. A JEPA trains a context encoder and a target encoder, often updated by exponential moving average, then predicts the target embedding of a masked observation from the context. No pixel decoder or explicit negatives are required. In the expression below, $g_\theta(f_\phi(x_{\mathrm{ctx}}))$ is the prediction, $f_{\bar\phi}(x_{\mathrm{tgt}})$ is the target, and $\mathrm{sg}$ prevents gradients from flowing into the target encoder:

where:

• $\mathcal{L}_{\mathrm{JEPA}}$: the JEPA predictive loss
• $x_{\mathrm{ctx}}$: visible context
• $x_{\mathrm{tgt}}$: masked target region
• $f_\phi$: context encoder
• $f_{\bar\phi}$: target encoder
• $g_\theta$: predictor of the target representation
• $\mathrm{sg}$: stop-gradient operator, blocking gradients through the target branch
• $\bar\phi$: target encoder parameters

The target parameters $\bar\phi$ are typically updated by an exponential moving average of $\phi$. As we discussed in Predictive Embeddings, JEPA, and Energy-Based Models, bypassing reconstruction removes the pixel-level capacity trap, but direct pixel samples are unavailable and additional asymmetry or regularization is needed to avoid collapse.

The pieces serve different roles. Predicting an embedding rather than pixels lets the target representation omit observation detail that its encoder does not preserve. A moving-average target changes gradually as the online encoder learns. Stop-gradient makes that target a fixed value for each online update. None of these choices alone rules out a constant representation; collapse avoidance comes from their interaction with masking, predictor asymmetry, normalization, and any explicit regularization in the particular method.

Compared to reconstruction, these objectives offer three concrete benefits:

• A configurable inductive bias. Embedding targets can make it cheap to ignore texture and lighting; raw masked-pixel targets retain more appearance pressure.
• Prediction-sufficiency pressure. When the target is a future embedding, the encoder is pushed to preserve information needed by the chosen predictive task. This becomes control-sufficiency pressure only when that task is aligned with the downstream decision, as discussed in Sufficient State and State Abstraction.
• Interface flexibility. Masked objectives can be adapted to any modality that has a natural masking operation, which is useful when the agent sees video, proprioception, and language.

The second benefit connects most directly to the control-theoretic notion of a sufficient state. A state is sufficient for a task if knowing it is as good as knowing the full history for the required prediction or decision. Contrastive and masked objectives reward information useful to their pretext tasks. If those tasks expose every distinction needed downstream, the learned representation may approach task sufficiency; if not, it can be sufficient for the wrong prediction.

The cost is a different optimization and evaluation burden. Temperature, mask ratio, target updates, and anti-collapse mechanisms can materially change the result, and many objectives do not come with the direct likelihood interpretation of reconstruction. Numerical comparisons across papers are also harder because there is no universal "bits per dimension" scale.

A model with low teacher-forced one-step error can still produce poor recursive rollouts. In one-step training, each target is predicted from a measured or inferred state anchored to data; in open-loop use, each later transition consumes a model-generated state. Multi-step and overshooting losses expose the objective to that recursive mode.

In teacher-forced one-step training, an error at one prediction does not become the next input. In a free-running rollout it does, so the model can visit inputs that are uncommon under the data-conditioned training distribution. Multi-step training starts from a data state and feeds predictions through several transitions before comparing with later targets. This does not eliminate distribution shift, but it trains directly on a bounded version of the recursive computation used at deployment.

Suppose the transition map is locally $\rho$-Lipschitz and introduces at most $\epsilon$ of one-step model error over the states visited by both rollouts. Starting from the same state, the rollout error obeys

where:

• $e_H$: state error after $H$ model steps
• $\epsilon$: bound on local transition mismatch
• $\rho$: bound on how much the transition amplifies a state perturbation
• $H$: rollout horizon
• $k$: summation index over rollout steps

Intuitively, each rollout step can amplify the error carried from the previous step by up to a factor $\rho$ and then adds at most $\epsilon$ of fresh one-step error, so the horizon-$H$ error is the sum of $H$ geometric terms. The bound therefore grows exponentially when $\rho>1$ (perturbations are amplified faster than they are injected), linearly when $\rho=1$, and approaches the steady-state value $\epsilon/(1-\rho)$ when $0\le\rho<1$ (fresh error is injected but prior error is damped away). It applies only while the stated local bounds continue to hold; leaving the covered state region can make the model error uncontrolled.

The regimes show how stability changes the consequence of a fixed local error bound. When $\rho<1$, the geometric bound saturates at $\epsilon/(1-\rho)$. Near $\rho=1$, it grows approximately linearly over the plotted horizon. When $\rho>1$, the upper bound grows geometrically. Greater one-step accuracy still reduces the bound in every regime; zero local mismatch would make it zero. Expansive dynamics amplify any nonzero residual more rapidly. The guarantee also applies only while both rollouts remain in the region where the Lipschitz and mismatch bounds hold.

The three regimes behave very differently, and evaluating the same closed form under three gains makes the contrast visible at a glance.

There is also a subtler failure. Even if the model's one-step error is small over the data-supported state region, a predicted state may wander into regions rarely seen during training. Behavior there is weakly constrained by data and depends on inductive bias. This is a distribution-shift problem, and it motivates explicit rollout-consistency objectives and out-of-distribution evaluation.

The simplest fix is to train on multi-step rollouts. Given a trajectory of length $T$ and a horizon $K$, the k-step objective is

where:

• $\mathcal{L}_K$: the k-step objective
• $\mathcal{D}$: training distribution over trajectories and actions
• $t$: start time; the expectation is over valid start times
• $K$: rollout horizon
• $w_k$: nonnegative weight controlling horizon emphasis
• $\hat{o}_{t+k}$: model prediction after rolling forward $k$ steps from the observed state at $t$, using recorded actions and feeding predicted states back into later transitions
• $o_{t+k}$: target observation at time $t+k$
• $k$: rollout step index
• $o,a$: observation and action sequences in the training distribution

The prediction $\hat{o}_{t+k}$ is obtained by rolling the model forward from the observed state at $t$ for $k$ steps, using the recorded actions and feeding predicted states back into later transitions. The nonnegative weights $w_k$ control the horizon emphasis. Setting $w_1=1$ and $w_k=0$ for $k>1$ recovers one-step training; $w_k=1$ weights all included horizons equally; and $w_k=\lambda^{k-1}$ with $0<\lambda<1$ discounts distant errors while leaving the one-step weight at one.

These three settings are easier to compare when drawn on the same axis, so the next cell computes the weight vectors for a rollout horizon of 15 steps.

Two implementation details matter. First, a recursive $K$-step loss normally starts from a data-conditioned state and then feeds model predictions into the remaining transitions. Supplying a ground-truth state at every step reduces the computation to a sum of teacher-forced one-step terms rather than testing recursive behavior. Scheduled sampling [@bengio2015scheduled] mixes data and model-generated inputs according to a schedule, although the resulting training distribution and estimator require care. Second, differentiating through $K$ transitions generally increases compute and activation memory with the unroll length; checkpointing, truncation, and model structure change the exact factor. The useful $K$ is task- and budget-dependent.

The weight schedule states which horizons matter to the training criterion. If a planner will use long rollouts, discounting distant errors almost to zero leaves that use weakly supervised. Uniform weights give distant targets equal scalar weight, but those targets can have larger, noisier gradients and destabilize optimization. A geometric schedule is one compromise, not a universal default. Choose both $K$ and $w_k$ using validation at the deployment horizons, while accounting for the additional optimization cost.

The latent overshooting objective introduced with PlaNet asks a multi-step latent prior, rolled forward without intermediate observations, to match posterior latent-state distributions inferred at future observations [@hafner2019planet]. For a start time $t$ and overshooting distance $d$, a representative term is

where:

• $\mathcal{L}_{\mathrm{over}}^{(d)}$: overshooting loss term at distance $d$
• $p_\theta^{(d)}$: the $d$-step predictive prior obtained by composing the transition model
• $q_\phi$: observation-conditioned posterior used as the target distribution
• $D_{\mathrm{KL}}(\cdot\|\cdot)$: Kullback-Leibler divergence
• $d$: overshooting distance
• $t$: start time

The KL direction trains the open-loop prior to cover states supported by the posterior. An implementation can average this term over selected distances and start times, often with horizon-dependent weights. Later Dreamer work retained closely related latent consistency terms, although the exact objective and whether overshooting helps are implementation-dependent.

The distinctive feature of overshooting is that it supervises a multi-step prior against a posterior inferred at a future observation. The composed prior advances without intermediate posterior corrections. With $D_{\mathrm{KL}}(q\|p)$ and gradients allowed into $p$, missing probability mass where the posterior places mass is penalized strongly. Actual gradient routing matters: implementations may stop gradients through one side, balance directions, or update both distributions, so the written KL direction alone does not determine which network moves.

Overshooting has its own tradeoffs. A deterministic prediction trained with squared error can approach the conditional mean of a multimodal future, producing an average that may not be plausible. Representing aleatoric alternatives requires a conditional distribution, for example through stochastic latents, mixture outputs, autoregressive tokens, or another probabilistic parameterization. Stochasticity and Uncertainty develops this issue further. DreamerV2 and DreamerV3 use discrete stochastic latent variables, but that is one design among several.

The averaging problem follows from using squared error on a multimodal conditional target. With two equally likely branches, the conditional mean minimizes expected squared error even when that mean lies between plausible outcomes. A probabilistic model can instead assign mass to both branches. Its output parameterization and proper scoring rule must therefore be designed together with the rollout objective.

If the model has a target encoder and a latent state, the multi-step loss can be

where:

• $\mathcal{L}_K^{\mathrm{latent}}$: latent-space multi-step loss
• the expectation is over data trajectories and valid start times
• $\hat{z}_{t+k}$: model prediction after $k$ action-conditioned latent transitions
• $z_{t+k}$: target-encoder representation of the future observation
• $\mathrm{sg}$: stop-gradient operator, preventing gradients through that target
• $K$: rollout horizon
• $w_k$: horizon weight
• $\mathcal{D}$: training distribution over trajectories and actions
• $o,a$: observation and action sequences
• $t$: start time
• $k$: rollout step index

The target encoder may be frozen or updated slowly; if it moves during training, its update must be controlled. Latent predictions avoid solving pixel reconstruction at every horizon, but their meaning depends on the quality and stability of the target representation.

The stop-gradient makes the target representation fixed for this loss update, so the prediction branch moves toward it rather than both branches moving together. It does not by itself rule out constant targets. Target initialization, update rule, normalization, architectural asymmetry, and any variance regularization determine whether the representation remains informative. A rapidly changing target can also turn the objective into a moving-target problem.

Choosing $K$ involves a three-way tension:

• Short $K$ is cheaper and exposes the model to only a limited prefix of its own errors.
• Long $K$ exposes training to more of the recursive computation used by a long-horizon planner, but costs more and can yield noisy far-horizon gradients or averaged point predictions.
• A growing or sampled $K$ is one possible compromise, but it changes the training distribution and should be validated.

Whatever the choice, evaluate at the intended deployment horizon. A model trained only on three-step recursive losses can still work at thirty steps when its learned dynamics are accurate and stable, but the longer use carries a larger extrapolation risk that short-horizon metrics will not reveal.

A modern world model is usually trained with several loss terms at once:

• A reconstruction or embedding loss that shapes the representation.
• A dynamics loss that shapes the transition model.
• A reward loss that shapes the reward predictor, and sometimes a continuation or termination head.

These terms pull the shared parameters in different directions and have different scales. Choosing how to weight them is not a footnote; it can decide whether training converges, and if it does, whether the model is useful for planning.

The reason multi-term training is a systems problem rather than a tuning problem is that the terms are not independent. They share parameters, so the gradient from one term changes the parameters that the other terms see, and the terms respond to those changes differently. A term with a large scale can dominate the shared trunk and effectively decide what the trunk represents, which then determines what the smaller terms can achieve even if their own heads are perfectly tuned. The question is not "which term is most important" but "which term's gradient is actually steering the shared parameters," and those two questions can have very different answers.

A typical world-model loss looks like

where:

• $\mathcal{L}$: total world-model loss
• $\mathcal{L}_{\mathrm{repr}}$: reconstruction, contrastive, or JEPA-style representation loss
• $\mathcal{L}_{\mathrm{dyn}}$: KL or latent dynamics term
• $\mathcal{L}_{\mathrm{rew}}$: reward prediction loss
• $\mathcal{L}_{\mathrm{aux}}$: any auxiliary regularizer, such as a sparsity or KL-balancing term
• $\alpha, \beta, \gamma, \delta$: weights that are either hand-tuned (the common case), derived from a principled scheme such as uncertainty weighting, or adapted during training by a gradient-balancing algorithm

The terms can have very different scales and learning curves. A sparse reward may be difficult despite its low output dimension; a reconstruction term can be easy or hard depending on the observation model; and an unconstrained KL term can change abruptly. Equal scalar weights do not imply equal influence on shared parameters. Output dimension and head size alone also do not determine that influence: the relevant quantity is each weighted term's gradient on the shared parameter block. Measure those norms and their pairwise alignment rather than inferring them from loss values.

Consider two problems that recurrent world models encounter:

• Encoder-prior mismatch. If the encoder and prior diverge, latent states used by the dynamics model are not the states the encoder actually produced during training. Then rollout predictions quickly degrade even though training loss keeps improving.
• Reward starvation. If the reward head shares parameters with the representation, and the reward gradient is much weaker than the reconstruction gradient, then shared weights drift toward reconstruction-relevant features that may be useless for reward prediction.

Useful controls include:

• Separate small networks ("heads") for reward and continuation so that gradient scale on the shared trunk is managed by explicit multipliers.
• Stop-gradient barriers where one module should be updated only by one term.
• Non-uniform weighting, either static or adaptive.

The critical point is that "balancing" here does not mean making the loss values numerically equal. It means making the gradients on shared parameters informative rather than dominated by one term.

Equal loss values do not imply equal gradients. Gradient norms depend on local derivatives, reductions, targets, and the parameter block being measured. The shared parameters move according to the weighted sum of those gradients, so topology, stop-gradients, and per-term monitoring are more informative than matching scalar loss values.

Kendall, Gal, and Cipolla proposed weighting task losses by learned homoscedastic uncertainty [@kendall2018multi]. Under the corresponding likelihood assumptions, each regression term receives a learned scalar $\sigma_i$ and the effective loss becomes

where:

• $\mathcal{L}$: total weighted loss
• $\mathcal{L}_i$: a task loss
• $\sigma_i>0$: a learned homoscedastic noise scale for that task
• $i$: task index

The inverse-variance factor reduces the contribution of noisier regression tasks, while $\log\sigma_i$ prevents the trivial solution $\sigma_i\to\infty$. Intuitively, if a task's noise scale $\sigma_i$ grows, the first term shrinks (down-weighting that task's loss) but the second term grows to penalize the inflated uncertainty, so the optimum trades off task fit against claimed certainty. This derivation depends on a probabilistic observation model; classification tasks use a related but not identical likelihood expression, so the formula is not a universal plug-in rule for arbitrary auxiliary losses.

GradNorm [@chen2018gradnorm] adapts task weights using training rates and gradient norms, while PCGrad [@yu2020gradient] projects a task gradient when its dot product with another task gradient is negative. PCGrad can reduce pairwise destructive interference, but the projection also changes the optimization direction and does not decide the desired task tradeoff for you.

These schemes solve narrower problems than their names may suggest. Uncertainty weighting learns relative scales under specified likelihood assumptions; it cannot identify a missing target or repair an under-determined task. PCGrad changes conflicting gradient directions; it cannot make incompatible representation requirements simultaneously satisfiable in one trunk. Both methods presuppose that the component tasks and shared architecture are sensible.

Reward and continuation predictors are often small heads on the latent state, trained with a suitable reward likelihood and a Bernoulli-style continuation loss. Their gradient connections to the representation matter as much as their architecture.

Two patterns dominate:

• Shared trunk, separate heads. Reward, continuation, and reconstruction share the trunk. This is efficient but forces the trunk to be useful for all heads.
• Stop-gradient heads. A task head can read the latent state through a stop-gradient so that its loss updates the head but not the representation trunk. This protects the trunk from that task gradient, but it cannot recover task-relevant information the representation omitted.

These are design options rather than a universal prescription. Allowing reward gradients into the trunk can make the state more task-aware, while detaching the state isolates representation learning from a potentially noisy objective. The correct choice depends on the implementation and should not be inferred from a model-family name alone.

The tradeoff between these two patterns is a concrete example of the general principle that decision usefulness and representation quality can pull in different directions. Letting the reward gradient shape the trunk makes the representation more aligned with reward, which sounds desirable, but it also makes the representation vulnerable to the noise and sparsity of the reward signal. Cutting the reward gradient off from the trunk protects the representation's quality but means the reward head has to do its learning on top of a fixed representation, which can be slower and less accurate if the representation happens to omit features the reward depends on. There is no universally correct choice, and the right one depends on how noisy the reward signal is and how much the task depends on reward-specific features that the representation loss would otherwise ignore.

Three practical rules of thumb, all of which will resurface when we look at concrete training recipes:

• Start with the simplest possible composition. Reconstruction + dynamics + reward, with rough weights from the source paper, is a stronger baseline than a bespoke mixture.
• Monitor per-term losses and shared-trunk gradient norms. A loss that is numerically larger can still have a smaller gradient, so log both when diagnosing dominance.
• Tune the smallest number of controls that addresses a measured failure. If the encoder-prior KL is drifting, inspect KL routing and balancing. If the reward head is underfit, test its capacity, weight, learning rate, targets, and whether the representation contains the required information.

We will revisit these ideas concretely in Interaction Data and Passive Observation and Online, Offline, and Continual Learning, where data composition interacts with loss weight selection. For now, the main takeaway is that world-model losses are a system, not a scalar. Every term shapes the representation toward a different notion of what matters, and the useful move is to design the composition so that the terms are pulling in compatible directions.

To make the multi-step story concrete, let us train two models on the same synthetic controlled system and compare their rollout error as the horizon grows. The system is a two-dimensional linear dynamical system with a mild nonlinearity:

where:

• $x_t\in\mathbb{R}^2$: state at time $t$
• $x_{t,1}$ and $x_{t,2}$: the two state components
• $a_t\in\mathbb{R}$: scalar action at time $t$
• $A$: linear state-transition matrix (defined immediately below)
• $B$: action matrix (defined immediately below)
• $\eta:\mathbb{R}^2\to\mathbb{R}^2$: deterministic nonlinear residual
• $x_{t+1}$: next state
• $t$: time index
• $0.05$: amplitude of the nonlinear residual

With

where:

• $A\in\mathbb{R}^{2\times2}$: linear state-transition matrix
• $B\in\mathbb{R}^{2\times1}$: maps the action into state space

the linear part is stable and cross-couples the two state dimensions, while the residual makes the transition nonlinear. This system is small enough for a tiny MLP, but still lets one-step and recursive objectives produce measurably different rollout behavior.

The eigenvalues of $A$ have magnitude below one, so the linear part is discrete-time stable; this alone does not prove global stability of the nonlinear controlled system for arbitrary actions.

The chosen matrix has eigenvalues inside the unit circle, so its uncontrolled linear component is asymptotically stable. That does not guarantee that either learned model will be contractive, nor does it determine which training objective wins under a finite optimization budget. The nonlinear residual ensures that the MLP is fitting more than a linear map.

Each trajectory is a sequence of states and actions; the model will be trained to predict state transitions.

We fit two models. The first minimizes one-step MSE over all 4,000 training transitions. The second starts from each sampled trajectory's $x_0$, rolls recursively for $K=5$, and minimizes the average of its five per-step MSE values. The model architecture, initial weights, optimizer, learning rate, and number of optimizer updates match. The training information and compute do not: every one-step update uses all 4,000 transitions, while every five-step update samples 32 trajectories and uses only their first five transitions. This is an intentionally compact demonstration of two training procedures, not a controlled estimate of the causal effect of the objective. The coverage mismatch is a plausible reason the five-step model performs worse in this run.

We then evaluate both models by rolling them out for the full horizon $T = 20$ on held-out trajectories and recording mean squared error at each horizon. In this deterministic run, the five-step procedure is about $7.26\times$ worse at horizon 1 and $33.05\times$ worse at horizon 20. That result is evidence about the procedures as implemented, including their unequal training coverage, not evidence that multi-step objectives are generally inferior.

The example does not establish that either objective is generally better. It shows that their rollout behavior can diverge and that training-budget controls are necessary before attributing a result to the objective itself.

Let us walk through the example in code. The book-wide plotting setup and shared imports are loaded in the hidden setup cell at the top of the chapter.

We now generate synthetic trajectories from the controlled system described in the worked example. We hold out separate training and test sets.

In a full world model, each loss also updates head-specific parameters that the other losses do not see. Those private gradients do not tell you which loss controls the shared state. The practical rule is to inspect and deliberately configure gradient flow on shared modules instead of reasoning from scalar losses or head sizes alone.

The key parameters for this implementation are:

• hidden: Number of hidden units in the MLP layers (32). Larger values increase capacity but also training time and overfitting risk.
• lr: Learning rate for the Adam optimizer (1e-2). Controls the step size for both the one-step and multi-step objectives.
• K: Rollout horizon for the multi-step training loss (5). Longer horizons expose the objective to more recursively generated states but cost more compute and memory.
• batch_size: Number of trajectories sampled per multi-step update (32). Larger batches reduce gradient variance but use more memory.
• num_iterations: Number of gradient updates for each objective (2000). More iterations can improve fit but also increase the chance of overfitting to the training trajectories.
• T: Trajectory length used both for data generation and evaluation rollouts (20). It sets the maximum horizon over which held-out rollout error is measured.

The choices we have discussed are the difference between a model that looks good on a training curve and a model that survives being asked to plan. Every family of objectives has failure modes, and the professional skill is knowing which failure mode is most painful for the task at hand.

Reconstruction and likelihood objectives provide dense supervision and a direct probabilistic interpretation when the likelihood is well specified. They reward modeling observations, including predictable detail that may not affect decisions. Posterior collapse is a separate latent-variable failure: a flexible decoder can model observations without using the latent, leaving the posterior close to the prior. KL schedules, free bits, decoder design, and KL balancing can mitigate collapse but do not guarantee a task-sufficient state. More generally, maximizing likelihood is not the same as maximizing decision quality, a point that Stochasticity and Uncertainty makes from the probabilistic side and that will recur in Part XI.

Likelihood scores observations according to the chosen density and reduction. Its effective emphasis depends on parameterized variances, normalization, discretization, and correlations; it does not generically weight every dimension in proportion to empirical variance. A decision objective instead emphasizes distinctions that change action values. Rare but consequential states may contribute little to average log-likelihood while dominating safety or return. Agreement between likelihood and decision quality therefore requires an argument about coverage, calibration, and task sufficiency rather than an assumption.

Contrastive and embedding-target objectives can avoid pixel-fidelity pressure, while raw-target masked methods retain part of it. Their behavior depends on positive and negative construction, mask policy, target updates, and anti-collapse design. These sensitivities differ from those of likelihood training rather than being uniformly worse. Their loss values also lack a universal cross-method scale, so downstream probes and rollout evaluations remain important.

Multi-step and overshooting losses address a specific structural problem: mismatch between teacher-forced training and recursive use. They do not solve it in general. Longer unrolls cost more and can produce noisier gradients. Under a pointwise squared-error objective, multimodal futures can be averaged. Long-horizon consistency targets can also be unreliable while the representation and transition are still changing. Schedules can help in some settings, but add choices that must be validated per environment.

Balancing multiple losses is consequential because only gradients on shared parameters create direct interference. Naive summation uses fixed scalar weights, but a larger loss value does not necessarily produce the larger shared gradient. Uncertainty weighting, GradNorm, PCGrad, stop-gradient boundaries, and KL balancing encode different assumptions and introduce their own controls. Keep the composition simple, log per-term values and shared gradients, and add machinery only for a diagnosed failure.

These objectives matter because they shape whether a learned state supports the rollouts and predictions a planner needs. A strong observation metric does not certify decision usefulness, while a modest generative score need not rule it out. The next chapters in Part VI therefore treat data and evaluation as part of objective design rather than afterthoughts.

The objective function is the world model's job description. Reconstruction and likelihood reward observation prediction and may spend capacity on task-irrelevant detail. Variational objectives add latent regularization whose strength and gradient routing affect posterior use. Contrastive and masked objectives instead emphasize information selected by their positives, masks, and targets; embedding targets can discard texture, whereas raw-pixel targets need not. Multi-step and overshooting losses train bounded recursive consistency at additional compute and optimization cost. Across multi-term objectives, gradient topology and shared-parameter norms matter alongside scalar weights.

Three principles tie the chapter together. First, evaluate the metrics and horizons used at deployment: low one-step error alone does not establish accurate long rollouts. Second, control sufficiency depends on the representation, objective, data, and downstream task, not on architecture alone. Third, there is no universal loss composition; instrument the training loop, monitor each term and its shared gradients, and change the objective in response to a specific measured failure.

A fourth principle is to track distinct quality axes without treating them as statistically independent. Predictive fidelity, representation quality, uncertainty calibration, and decision usefulness measure different properties, but can interact through the same model and data. None alone guarantees the others. A claim that an objective produces a useful world model therefore needs direct evidence on the intended decisions as well as predictive diagnostics.

The next chapter, Actions, Inverse Dynamics, and Latent Actions, takes the objective one level down: not just what the model predicts, but how actions enter the prediction. Does the model see the action as an input, infer it as a latent variable, or discover it through an inverse dynamics model? That question changes both the loss and the state, and it sets up the data-curation and pretraining questions that follow.

Ready to test your understanding of predictive and generative objectives? This quiz reinforces the key ideas about reconstruction, contrastive and masked prediction, multi-step losses, and loss balancing.