Predictive and Generative Objectives › Multi-step and Overshooting Losses

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

$$\begin{aligned} e_H &\le \epsilon \sum_{k=0}^{H-1}\rho^k \\ &= \begin{cases} \epsilon(\rho^H-1)/(\rho-1), & \rho\ne 1, \\ H\epsilon, & \rho=1. \end{cases} \end{aligned}$$

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

$$\mathcal{L}_K = \mathbb{E}_{(o,a)\sim \mathcal{D},\, t} \left[ \sum_{k=1}^{K} w_k \, \lVert \hat{o}_{t+k} - o_{t+k} \rVert_2^2 \right],$$

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

$$\mathcal{L}_{\mathrm{over}}^{(d)} = \mathbb{E}_{q_\phi(z_t\mid o_{\le t},a_{<t})}\!\left[D_{\mathrm{KL}}\!\left(q_\phi(z_{t+d}\mid o_{\le t+d},a_{<t+d})\,\|\,p_\theta^{(d)}(z_{t+d}\mid z_t,a_{t:t+d-1})\right)\right].$$

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

$$\mathcal{L}_K^{\mathrm{latent}} = \mathbb{E}_{(o,a)\sim \mathcal{D},\, t} \left[ \sum_{k=1}^{K} w_k \, \lVert \hat{z}_{t+k} - \mathrm{sg}(z_{t+k}) \rVert_2^2 \right],$$

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.