Online, Offline, and Continual Learning › Code Implementation
World Models HandbookAll chapters Online, Offline, and Continual Learning › Code ImplementationOne section · Sections in this chapter7 / 11 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
We will now build three small experiments that isolate the mechanisms discussed above. The first trains a small neural dynamics model on two environments in sequence, comparing naive fine-tuning against replay and against Fisher-weighted regularization. The second tracks a drifting coefficient with and without a replay buffer. The third demonstrates model exploitation in a planner. Each experiment precomputes its results from variables and is followed by a discussion of what the numbers mean.
The experiments are deliberately small. The goal is not to achieve state-of-the-art performance but to make the mechanisms visible. The first experiment uses a two-environment sequence so that forgetting and retention can be measured directly. The second uses a scalar drifting coefficient so that the adaptation and noise tradeoff can be seen in a single trace. The third uses a planner with a widening search budget so that the gap between predicted and true return can be measured. Each experiment isolates one part of the problem.
We start with imports and a seeded random generator. Everything downstream uses this generator so results are reproducible.
from collections import deque
import numpy as np
import torch
import torch.nn as nn
torch.manual_seed(0)
np.random.seed(0)
rng = np.random.default_rng(0)Next we define two environments that share a state-action space but have different dynamics. Both are linear in state and action, differing in the state coefficient and the action gain. The difference is large enough that a single shared model cannot serve both without being told which one it is in.
ENV_PARAMS = {0: (0.85, 0.40), 1: (0.55, 0.70)}
def simulate(env_id, n_steps, rng):
a, b = ENV_PARAMS[env_id]
s = rng.normal()
states, actions, next_states = [], [], []
for _ in range(n_steps):
u = rng.uniform(-1.0, 1.0)
s_next = a * s + b * u + 0.05 * rng.normal()
states.append(s)
actions.append(u)
next_states.append(s_next)
s = s_next
return np.array(states), np.array(actions), np.array(next_states)Each transition is packaged as a model input where is a one-hot indicator of the environment. Giving the model the environment identity isolates forgetting from the harder problem of inferring which environment it is in.
def make_dataset(env_id, n_steps, rng):
s, u, s_next = simulate(env_id, n_steps, rng)
onehot = np.zeros((n_steps, 2), dtype=np.float32)
onehot[:, env_id] = 1.0
X = np.concatenate([s[:, None], u[:, None], onehot], axis=1).astype(
np.float32
)
y = s_next[:, None].astype(np.float32)
return X, y
X_A, y_A = make_dataset(0, 3000, rng)
X_B, y_B = make_dataset(1, 3000, rng)
X_A_test, y_A_test = make_dataset(0, 500, rng)
X_B_test, y_B_test = make_dataset(1, 500, rng)The model is a small multilayer perceptron. It is deliberately modest so that forgetting appears quickly and the comparison between methods is visible within a short training run.
class DynamicsMLP(nn.Module):
def __init__(self, in_dim=4, hidden=32):
super().__init__()
self.net = nn.Sequential(
nn.Linear(in_dim, hidden),
nn.ReLU(),
nn.Linear(hidden, hidden),
nn.ReLU(),
nn.Linear(hidden, 1),
)
def forward(self, x):
return self.net(x)We need a held-out evaluation helper. This computes mean squared error on a given dataset without altering the model's parameters.
def mse(model, X, y):
model.eval()
with torch.no_grad():
pred = model(torch.from_numpy(X))
value = float(((pred - torch.from_numpy(y)) ** 2).mean())
model.train()
return value
def rollout_mse(model, X, y, horizon=8):
"""Mean squared error over recursive one-step-ahead rollouts."""
model.eval()
with torch.no_grad():
s = float(X[0, 0])
total = 0.0
for t in range(min(horizon, len(X))):
inp = torch.tensor([[s, X[t, 1], *X[t, 2:]]], dtype=torch.float32)
pred = model(inp).item()
total += (pred - y[t, 0]) ** 2
s = torch.tensor(pred)
value = total / min(horizon, len(X))
model.train()
return valueFor the regularization baseline we need a diagonal Fisher estimate. We approximate the Fisher by summing squared gradients of the squared-error loss over a sample of task-A transitions, then normalize so that the regularization strength is comparable across parameters. We compute it once, at the moment of the task switch, and hold it fixed while training on environment B.
def compute_fisher(model, X, y, n_samples=400):
fisher = {n: torch.zeros_like(p) for n, p in model.named_parameters()}
idx = np.random.choice(len(X), size=min(n_samples, len(X)), replace=False)
for i in idx:
model.zero_grad()
out = model(torch.from_numpy(X[i : i + 1]))
target = torch.from_numpy(y[i : i + 1])
loss = ((out - target) ** 2).sum()
loss.backward()
for n, p in model.named_parameters():
fisher[n] += p.grad.detach() ** 2
for n in fisher:
fisher[n] = fisher[n] / len(idx)
scale = sum(f.mean() for f in fisher.values()) / len(fisher)
for n in fisher:
fisher[n] = fisher[n] / (scale + 1e-8)
return fisher
def ewc_penalty(model, fisher, anchor, lam):
penalty = torch.zeros(())
for name, param in model.named_parameters():
penalty = penalty + (fisher[name] * (param - anchor[name]) ** 2).sum()
return 0.5 * lam * penaltyThe training loop below is the core of the experiment. It runs a fixed number of steps on environment A, then the same number on environment B, logging held-out error on both environments throughout. Three methods share the loop: naive uses only the current task's data, replay adds half as many task-A samples as current-task samples (so old data is one third of the combined batch), and ewc adds the Fisher-weighted penalty computed at the moment of the switch.
STEPS_PER_TASK = 800
def train_continual(
method,
X_A,
y_A,
X_B,
y_B,
X_A_test,
y_A_test,
X_B_test,
y_B_test,
steps_per_task=STEPS_PER_TASK,
batch=64,
lr=1e-3,
lam=5.0,
log_every=25,
):
model = DynamicsMLP()
opt = torch.optim.Adam(model.parameters(), lr=lr)
XA, yA = torch.from_numpy(X_A), torch.from_numpy(y_A)
XB, yB = torch.from_numpy(X_B), torch.from_numpy(y_B)
nA, nB = len(XA), len(XB)
fisher, anchor = None, None
log = {"step": [], "taskA": [], "taskB": []}
for step in range(2 * steps_per_task):
on_task_b = step >= steps_per_task
if not on_task_b:
idx = torch.randint(0, nA, (batch,))
X, y = XA[idx], yA[idx]
else:
idx = torch.randint(0, nB, (batch,))
X, y = XB[idx], yB[idx]
if method == "replay":
k = batch // 2
ia = torch.randint(0, nA, (k,))
X = torch.cat([X, XA[ia]])
y = torch.cat([y, yA[ia]])
opt.zero_grad()
loss = ((model(X) - y) ** 2).mean()
if method == "ewc" and on_task_b:
loss = loss + ewc_penalty(model, fisher, anchor, lam)
loss.backward()
opt.step()
if method == "ewc" and step == steps_per_task - 1:
fisher = compute_fisher(model, X_A, y_A)
anchor = {
n: p.detach().clone() for n, p in model.named_parameters()
}
if (
step % log_every == 0
or step == steps_per_task - 1
or step == 2 * steps_per_task - 1
):
log["step"].append(step)
log["taskA"].append(mse(model, X_A_test, y_A_test))
log["taskB"].append(mse(model, X_B_test, y_B_test))
return model, logWe now run all three methods with the same seed so that the comparison is as paired as possible. This cell only computes the results; the next cell displays them and the figure cell below plots them from the precomputed logs.
results = {}
models = {}
for method in ["naive", "replay", "ewc"]:
torch.manual_seed(0)
np.random.seed(0)
model, log = train_continual(
method, X_A, y_A, X_B, y_B, X_A_test, y_A_test, X_B_test, y_B_test
)
results[method] = log
models[method] = modelrollout_stats = {}
for method, model in models.items():
rollout_stats[method] = {
"taskA": rollout_mse(model, X_A_test, y_A_test),
"taskB": rollout_mse(model, X_B_test, y_B_test),
}The printed summary reports held-out error on each environment at the end of training. The interesting comparison is task A: a method that has forgotten environment A will show a much higher error there than a method that has retained it.
naive | final task-A MSE: 0.0189 | final task-B MSE: 0.0029 | task-A rollout MSE: 0.0630 | task-B rollout MSE: 0.0027 replay | final task-A MSE: 0.0025 | final task-B MSE: 0.0028 | task-A rollout MSE: 0.0018 | task-B rollout MSE: 0.0037 ewc | final task-A MSE: 0.0034 | final task-B MSE: 0.0029 | task-A rollout MSE: 0.0114 | task-B rollout MSE: 0.0023
The final one-step and rollout MSE values show the stability-plasticity tradeoff. Naive fine-tuning achieves the lowest task-B error but the largest task-A error, while replay and EWC retain task-A performance at some cost on task B. Rollout error is higher for all methods because prediction errors compound over the horizon, but the ranking between methods is preserved, which indicates that the models learned dynamics rather than merely memorizing one-step transitions.
Read the two panels together. In the right panel, naive training reaches the lowest error on environment B. It pays for that with the largest regression in the left panel. Replay and EWC sit higher on task B and much lower on task A. This is the stability-plasticity tradeoff made visible: the method that is most plastic is the one that forgets most.
The comparison also shows that replay and EWC are not identical. Replay directly reintroduces old data, so it retains task A by continuing to train on it. EWC does not reintroduce data; it constrains parameters that were important for task A. Replay can be seen as changing the data distribution, while EWC can be seen as changing the loss landscape. Both reduce forgetting, but they do so through different mechanisms, and their relative advantage depends on the buffer size, the regularization strength, and the overlap between tasks.
The next experiment isolates streaming estimation. A single scalar coefficient in the true dynamics changes abruptly halfway through a long stream. We estimate it with plain SGD and with SGD over a replay buffer, and we record the state coefficient over time.
def stream_parameters(
n_steps=4000, switch_at=2000, eta=0.02, buffer_size=400, batch=16
):
a_true = np.where(np.arange(n_steps) < switch_at, 0.85, 0.55)
b_true = np.where(np.arange(n_steps) < switch_at, 0.40, 0.70)
s = 0.0
theta_sgd = np.zeros(2)
theta_replay = np.zeros(2)
buf_s = deque(maxlen=buffer_size)
buf_u = deque(maxlen=buffer_size)
buf_y = deque(maxlen=buffer_size)
trace = {"t": [], "a_true": [], "a_sgd": [], "a_replay": []}
for t in range(n_steps):
u = rng.uniform(-1.0, 1.0)
y = a_true[t] * s + b_true[t] * u + 0.05 * rng.normal()
# Plain streaming SGD on the newest transition only.
pred_sgd = theta_sgd[0] * s + theta_sgd[1] * u
theta_sgd = theta_sgd - eta * (pred_sgd - y) * np.array([s, u])
# Streaming SGD over a mini-batch drawn from the replay buffer.
buf_s.append(s)
buf_u.append(u)
buf_y.append(y)
idx = rng.integers(0, len(buf_s), size=min(batch, len(buf_s)))
bs = np.array(buf_s)[idx]
bu = np.array(buf_u)[idx]
by = np.array(buf_y)[idx]
residual = theta_replay[0] * bs + theta_replay[1] * bu - by
grad = (residual[:, None] * np.stack([bs, bu], axis=1)).mean(axis=0)
theta_replay = theta_replay - eta * grad
s = y
if t % 20 == 0:
trace["t"].append(t)
trace["a_true"].append(a_true[t])
trace["a_sgd"].append(theta_sgd[0])
trace["a_replay"].append(theta_replay[0])
return tracetrace = stream_parameters()
The plot shows both effects at once. The plain SGD curve reacts almost immediately at the switch, but it never settles: it wobbles around the new value with the noise ball predicted by the analysis. The replay curve is visibly smoother, because each update averages over sixteen transitions, but it takes longer to arrive at the new value because the buffer still contains old-regime data. Neither curve is strictly better. The right choice depends on whether you care more about responsiveness or variance.
The experiment also illustrates the role of the buffer as a memory horizon. When the true coefficient changes, the replay estimate continues to be pulled toward the old value until the old transitions have been overwritten in the buffer. The plain SGD estimate has no such inertia, but it pays for its responsiveness with variance. If the buffer were larger, the replay curve would lag even more. If it were smaller, the replay curve would track more quickly but would retain less protection against forgetting.
The final experiment demonstrates model exploitation. We train a dynamics model on transitions whose actions are restricted to , then plan over a much wider action range . The true dynamics saturate because of a hyperbolic tangent; a model trained on the narrow range has no way to know that.
def true_step(s, u):
return 0.9 * s + 0.4 * np.tanh(u)
def true_rollout(s0, actions):
s = s0
total = 0.0
for u in actions:
s = true_step(s, u)
total += -((s - 0.8) ** 2)
return total
n_train = 3000
u_train = rng.uniform(-1.0, 1.0, size=n_train)
s_train = np.zeros(n_train)
y_train = np.zeros(n_train)
s_cur = rng.normal()
for i in range(n_train):
y_train[i] = true_step(s_cur, u_train[i]) + 0.02 * rng.normal()
s_train[i] = s_cur
s_cur = y_train[i]
X_ex = np.stack([s_train, u_train], axis=1).astype(np.float32)
y_ex = y_train[:, None].astype(np.float32)torch.manual_seed(1)
exploit_model = DynamicsMLP(in_dim=2)
opt = torch.optim.Adam(exploit_model.parameters(), lr=1e-2)
Xe, ye = torch.from_numpy(X_ex), torch.from_numpy(y_ex)
for _ in range(1500):
idx = torch.randint(0, len(Xe), (128,))
opt.zero_grad()
loss = ((exploit_model(Xe[idx]) - ye[idx]) ** 2).mean()
loss.backward()
opt.step()Now we plan by random shooting. For each search budget we sample that many random action sequences, score them with the learned model, keep the best, and evaluate its true return. As the budget grows, the optimizer has more chances to find a sequence the model is optimistic about.
def model_rollout_return(model, s0, U):
n = U.shape[0]
s = torch.full((n, 1), s0)
total = torch.zeros(n)
for t in range(U.shape[1]):
inp = torch.cat([s, U[:, t : t + 1]], dim=1)
s = model(inp)
total = total - (s.squeeze(1) - 0.8) ** 2
return total
horizon = 6
s0 = 0.0
budgets = [1, 2, 5, 10, 25, 50, 100, 250, 500, 1000, 2500, 5000, 10000]
pred_returns, true_returns, chosen_actions = [], [], []
exploit_model.eval()
with torch.no_grad():
for n in budgets:
U = torch.empty(n, horizon).uniform_(-3.0, 3.0)
predicted = model_rollout_return(exploit_model, s0, U)
best = int(predicted.argmax())
pred_returns.append(float(predicted[best]))
true_returns.append(true_rollout(s0, U[best].numpy()))
chosen_actions.append(float(U[best].abs().max()))candidates= 1 | predicted= -3.998 | true= -4.002 | max|u|=0.84 candidates= 2 | predicted= -2.537 | true= -1.936 | max|u|=2.61 candidates= 5 | predicted= -0.880 | true= -0.835 | max|u|=2.31 candidates= 10 | predicted= -0.711 | true= -1.084 | max|u|=2.60 candidates= 25 | predicted= -0.477 | true= -0.406 | max|u|=2.14 candidates= 50 | predicted= -0.585 | true= -0.854 | max|u|=2.21 candidates= 100 | predicted= -0.561 | true= -0.469 | max|u|=2.88 candidates= 250 | predicted= -0.371 | true= -0.488 | max|u|=2.81 candidates= 500 | predicted= -0.261 | true= -0.354 | max|u|=2.92 candidates= 1000 | predicted= -0.191 | true= -0.269 | max|u|=2.76 candidates= 2500 | predicted= -0.187 | true= -0.435 | max|u|=2.95 candidates= 5000 | predicted= -0.180 | true= -0.358 | max|u|=2.70 candidates= 10000 | predicted= -0.165 | true= -0.378 | max|u|=2.89
The table illustrates the optimizer's curse. Across these independently drawn candidate sets, predicted and true returns fluctuate rather than changing monotonically, but larger searches tend to locate sequences with optimistic model estimates. Selected actions can also approach the edge of the search range, where the model was never trained and extrapolates most aggressively.
The gap is the point, not monotonicity. Because each budget draws a fresh candidate set, either curve can reverse locally. In this seeded run the broader searches more often find optimistic extrapolations, and their selected model return tends to exceed the measured true return. This single experiment illustrates the selection mechanism; it does not establish that every larger budget is worse or that a small budget is safe.
This is the mechanism behind unsafe drift. It is not that the model is bad on the data it was trained on. It is that the planner is an optimizer, and optimizers search until they find something, including the model's mistakes.
The experiment also shows why uncertainty penalties are helpful. If the planner were penalized for actions where the model is uncertain, the edge of the search range would be less attractive. The model was never trained on actions outside , so a well-calibrated uncertainty estimate would be large there. Without such a penalty, the planner has no reason to avoid the edge. With it, the planner can still explore but is discouraged from exploiting the model's most extreme extrapolations.