Online, Offline, and Continual Learning › Distribution Shift and Replay

Distribution shift is the situation where the data the model sees at deployment differs from the data it was trained on. It is the reason an offline model goes stale and the reason an online model can chase its own tail.

The classical taxonomy distinguishes three kinds of shift, and it is worth mapping each onto world models explicitly.

• Covariate shift: the distribution of inputs $p(x)$ changes, but the conditional $p(y \mid x)$ does not. In a world model this means the agent visits new states and new actions, but the underlying physics is unchanged. This happens constantly, because policies explore.
• Concept shift: the conditional $p(y \mid x)$ changes while inputs stay in the same range. In a world model this is a real change in dynamics: the conveyor belt speeds up, the surface becomes slippery, gravity changes in simulation. This is what we usually mean by the world changing.
• Prior-probability shift: the label prior $p(y)$ changes while the class-conditional input distribution $p(x \mid y)$ remains fixed. In world-model deployments, analogous changes in goal or task frequencies can alter what the planner encounters even when the transition dynamics are unchanged.

The distinction matters because the remedies differ. Covariate shift can often be handled by collecting more data in the new input regions, because the underlying mapping is unchanged. Concept shift cannot be fixed by coverage alone; the model must update its parameters to reflect a new mapping. Prior shift may not require any change to the dynamics model at all if the dynamics are unchanged, but it can still matter for planning if the objective or the distribution of goals has changed.

A supervised model that experiences covariate shift can make worse predictions on new inputs. A world model inside an agent can also induce further shift. The policy queries the model, the model returns predictions, the policy acts on those predictions, and the resulting transitions become the next training data. Model error can therefore alter the state distribution, which then changes what the model learns next. This feedback is what makes distribution shift in world models qualitatively different from distribution shift in ordinary supervised learning.

There are two consequences worth separating.

First, extrapolation is not rare, it is systematic. The policy is chosen to maximize expected return under the model. By construction, it prefers states and actions where the model predicts high return. If the model is optimistic anywhere, the policy will find that region. This is the optimizer's curse, and we will demonstrate it concretely in the code.

Second, the training distribution can narrow over time. If the agent only trains on data from the current policy, it may gradually stop seeing states that earlier policies visited. The model can become confident about a shrinking region and increasingly wrong outside it. This creates the conditions for a self-confirming loop: model errors steer the policy, and the resulting data can entrench those errors.

A useful way to see the difference from supervised learning is to imagine a classifier that is used to decide which emails reach each inbox. If the classifier makes mistakes, recipients see worse emails, but the distribution of incoming emails does not change because of the classifier's decisions. In a world model, the policy is like a classifier that decides which states to visit, and those states become the training data for the next round. The classifier is not a passive observer of its input distribution; it is an active participant in generating that distribution.

The most widely used defense is experience replay. Instead of updating only on the newest transition, the agent stores transitions in a buffer $\mathcal{B}$ and samples mini-batches from it:

$$\theta \leftarrow \theta - \eta\, \nabla_\theta \frac{1}{|\mathcal{B}_{\text{batch}}|} \sum_{(x,y) \in \mathcal{B}_{\text{batch}}} \ell\big(f_\theta(x),\, y\big).$$

where:

• $\theta$: the model parameters
• $\eta$: the learning rate
• $\nabla_\theta$: the gradient with respect to the parameters $\theta$
• $\mathcal{B}_{\text{batch}}$: the mini-batch sampled from the replay buffer $\mathcal{B}$
• $|\mathcal{B}_{\text{batch}}|$: the number of transitions in the mini-batch
• $\ell(f_\theta(x), y)$: the loss on a buffered transition $(x, y)$
• $f_\theta(x)$: the model's prediction for input $x$
• $x, y$: a transition from the buffer, with $x$ the input and $y$ the target
• $\mathcal{B}$: the replay buffer storing past transitions

A replay buffer is a sampled approximation to the retained history. It converts a streaming problem into something closer to a batch problem without giving up the ability to absorb new data, and is one of the main pieces of infrastructure connecting offline and online learning.

Two design choices matter most.

Buffer size controls the effective memory horizon. A small buffer tracks recent dynamics quickly and forgets older ones. A large buffer retains a broad distribution and adapts slowly. This is the same tradeoff as the forgetting factor in RLS, expressed in units of samples rather than a discount rate.

Sampling strategy controls which parts of history are emphasized.

• Uniform sampling treats every stored transition equally. It is simple, unbiased with respect to the buffer contents, and the default in most implementations.
• Prioritized sampling draws transitions with probability proportional to a priority such as the last temporal-difference error or model loss (Schaul et al., 2016). It focuses capacity on the examples the model currently gets wrong. The cost is that the sampling distribution is no longer the data distribution, so importance weights are needed to keep the gradient unbiased.
• Recency-weighted sampling blends the two, up-weighting recent transitions to track drift while still retaining a background of older data.

A replay buffer stores samples and shapes the training distribution. It determines the distribution of data that the model is trained on, and therefore the errors that the model is optimized to reduce. If the buffer is dominated by recent data, the model is optimized for the recent regime. If it contains old data, the model is optimized for a mixture of regimes. The mixture can be desirable when old data remains valid and harmful when it is obsolete.

Replay introduces a second ratio to track alongside the update-to-data ratio: the replay fraction, the proportion of each batch drawn from the buffer rather than from fresh data. A replay fraction of zero is pure streaming. A fraction near one is nearly offline training on a frozen buffer.

Increasing the replay fraction stabilizes learning and reduces forgetting, but it also slows adaptation to real change, because the gradient is dominated by data that may no longer reflect the world. There is no universally correct value. What matters is that the choice is made deliberately and monitored, and that the buffer is large enough to cover the range of conditions the agent is expected to encounter.

The interaction between replay fraction and buffer size can be subtle. A large buffer with a high replay fraction produces a very stable model that changes slowly. A small buffer with a low replay fraction produces a very adaptive model that may forget quickly. A small buffer with a high replay fraction can be worse than either, because the model repeatedly trains on a narrow slice of recent history and can overfit to it while still calling it replay. The design must be considered as a whole.

When the buffer contains transitions generated by older policies, the data are off-policy with respect to the current one. For pure dynamics modeling this often matters less than it does for value learning: in a stationary Markov environment, $p(s_{t+1} \mid s_t, a_t)$ does not depend on which policy selected $a_t$. Data from any policy can therefore train the same conditional dynamics on the state-action support it covers. Hidden state, environment drift, selection effects, or inadequate support break this simple argument.

Coverage is doing real work here. Off-policy data identifies the conditional dynamics only on the region it covers. Outside that region the model is extrapolating, and that is often where the policy goes next.

This asymmetry between dynamics learning and value learning is worth emphasizing. A value function is defined relative to a policy and reward, so naive off-policy targets can be biased; importance sampling, corrected targets, or explicitly off-policy algorithms address that mismatch. A dynamics model is defined by the environment, so data collected under any policy can estimate the same transition conditional when the Markov and stationarity assumptions hold and the state-action pair lies in the data support. Replay is therefore often a natural fit for model-based learning. Its central danger is unsupported extrapolation rather than the policy label attached to an otherwise valid transition.