Predictive and Generative Objectives › Balancing Representation, Dynamics, and Reward Losses

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

$$\mathcal{L} = \alpha \mathcal{L}_{\mathrm{repr}} + \beta \mathcal{L}_{\mathrm{dyn}} + \gamma \mathcal{L}_{\mathrm{rew}} + \delta \mathcal{L}_{\mathrm{aux}},$$

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

$$\mathcal{L} = \sum_i \left( \frac{1}{2 \sigma_i^2} \mathcal{L}_i + \log \sigma_i \right).$$

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.