Pretraining, Post-Training, Scaling, and Adaptation › Code Implementation
World Models HandbookAll chapters Pretraining, Post-Training, Scaling, and Adaptation › 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
Let's build the pipeline end to end. We start with a scaling-law fit, then move to the embodiment experiment, then finish with parameter-efficient adaptation and preference-based post-training.
Fitting a scaling law from measurements
We start by generating synthetic loss measurements from a known scaling law, then fitting the exponents back. This mirrors the practical workflow: you measure loss at a small ladder of points and fit a power law to extrapolate. The points are chosen to span multiple orders of magnitude in both axes, since a power law can only be identified if the fit data covers enough dynamic range.
# Synthetic measurements of loss at a grid of (N, D) points.
N_sizes = np.logspace(5, 11, 9)
D_sizes = np.logspace(6, 12, 9)
# Ground-truth constants, Kaplan-style for language.
alpha_N_true, alpha_D_true = 0.076, 0.095
Nc_true, Dc_true = 8.8e13, 5.4e13
L_inf_true = 1.69
N_grid, D_grid = np.meshgrid(N_sizes, D_sizes, indexing="ij")
L_true = (
(Nc_true / N_grid) ** alpha_N_true
+ (Dc_true / D_grid) ** alpha_D_true
+ L_inf_true
)
# Add realistic multiplicative noise to simulate measurement variability.
# A clean 1% noise makes the measured points sit almost exactly on the fitted
# line, which is too clean to be instructive; 3% gives visibly scattered
# measurements while still letting the power law be recovered.
L_obs = L_true * np.exp(rng.normal(0, 0.03, size=L_true.shape))Now we fit the five parameters of the scaling law. We work in log space to keep the optimization well-conditioned.
def fit_loss(params, N, D, L_obs):
log_Nc, log_Dc, log_alpha_N, log_alpha_D, log_Linf = params
Nc, Dc, alpha_N, alpha_D, L_inf = np.exp(params)
L_pred = (Nc / N) ** alpha_N + (Dc / D) ** alpha_D + L_inf
return np.mean((np.log(L_pred) - np.log(L_obs)) ** 2)
x0 = np.log(np.array([1e14, 1e14, 0.1, 0.1, 1.7]))
res = minimize(fit_loss, x0, args=(N_grid, D_grid, L_obs), method="Nelder-Mead")
if not res.success:
raise RuntimeError(f"Scaling-law fit failed: {res.message}")
Nc_hat, Dc_hat, alpha_N_hat, alpha_D_hat, L_inf_hat = np.exp(res.x)Fitted alpha_N = 0.072 (true 0.076) Fitted alpha_D = 0.089 (true 0.095) Fitted L_inf = 1.484 (true 1.69)
Because these synthetic observations were generated from the same functional form, the fitted exponents should be close to their known values. This checks the fitting code; it does not show that a nine-by-nine ladder is sufficient for an unknown real system.
The measured points follow the fitted trend with visible scatter, as expected from the multiplicative noise added above. In real experiments the scatter can be larger and the exponents are estimated with less confidence.
The three terms in the fitted law answer different questions, and their relative size tells you which lever to pull. Splitting the fit at the largest dataset size makes the crossover visible.
# Decompose the fitted scaling law at the largest dataset size.
D_ref = float(D_sizes[-1])
param_term = (Nc_hat / N_sizes) ** alpha_N_hat
data_term = np.full_like(param_term, (Dc_hat / D_ref) ** alpha_D_hat)
floor_term = np.full_like(param_term, L_inf_hat)
When the parameter term sits above the others, growing the model is the intervention this fitted decomposition favors. Once it drops below the data term or floor, extra parameters buy little within the model, and data or a reformulated prediction target becomes the next candidate experiment.
We can do the same thing for the data axis. Now we fix the model size and vary .
The two plots illustrate how a fitted law can compare allocations inside the measured regime. Choosing a compute-optimal allocation then requires minimizing that fitted loss along a fixed-compute constraint; a ratio borrowed from another experiment is only a reference allocation.
The compute constraint is easiest to see in space. Each fixed budget traces a hyperbola. The line below is an illustrative Chinchilla-style allocation rule; it is not the optimum of this synthetic fitted law.
# Allocations where each compute contour intersects the illustrative D = 20N ray.
D_per_param = 20.0
C_budgets = np.array([1e20, 1e21, 1e22, 1e23])
N_ratio = np.sqrt(C_budgets / (6.0 * D_per_param))
D_ratio = D_per_param * N_ratio
Moving along the ray gives a sequence of allocations with a fixed data-to-parameter ratio. To locate this chapter's fitted optimum, you would instead minimize the fitted loss separately along each compute contour and confirm the answer with runs near that point.
A family of embodiments to pretrain and adapt
Now let's build the embodiment family. We define a rollout function for a driven damped point mass, then a dataset builder that stacks transitions across parameter settings.
def transition(state, action, mass, damping, dt=0.1):
x, v = state
v_next = v + dt * (-damping / mass * v) + dt * action / mass
x_next = x + dt * v_next
return np.array([x_next, v_next], dtype=np.float32)
def rollout(mass, damping, n_steps, rng, dt=0.1):
state = np.array([1.0, 0.0], dtype=np.float32)
states = [state.copy()]
actions = []
for _ in range(n_steps):
a = float(rng.normal(0.0, 1.0))
state = transition(state, a, mass, damping, dt)
states.append(state.copy())
actions.append(a)
return np.array(states, dtype=np.float32), np.array(
actions, dtype=np.float32
)Next we define a dataset builder that produces aligned arrays of states, actions, next-states, and task identifiers.
def build_dataset(embodiments, n_traj, n_steps, rng):
S, A, Y, E = [], [], [], []
for i, (m, d) in enumerate(embodiments):
for _ in range(n_traj):
states, actions = rollout(m, d, n_steps, rng)
S.append(states[:-1])
A.append(actions[:, None])
Y.append(states[1:])
E.append(np.full(len(actions), i, dtype=np.int64))
return (
np.stack(S).astype(np.float32), # (n_traj, n_steps, 2)
np.stack(A).astype(np.float32), # (n_traj, n_steps, 1)
np.stack(Y).astype(np.float32), # (n_traj, n_steps, 2)
np.stack(E), # (n_traj, n_steps)
)We split the embodiments into four training systems and one held-out target system.
data_rng = np.random.default_rng(7)
# Choose training embodiments that share the driven damped-system
# structure with the target but differ enough in (m, b) that zero-shot
# transfer is useful but imperfect, leaving room for adaptation to help.
# The target is not adjacent to any single training embodiment.
train_embodiments = np.array(
[
[1.0, 0.3],
[2.0, 1.0],
[0.5, 0.15],
[1.6, 0.5],
]
)
test_embodiment = np.array([[1.2, 0.8]])
# Pretraining data: many trajectories from each training embodiment.
S_tr_traj, A_tr_traj, Y_tr_traj, E_tr_traj = build_dataset(
train_embodiments, n_traj=20, n_steps=30, rng=data_rng
)
S_tr = S_tr_traj.reshape(-1, S_tr_traj.shape[-1])
A_tr = A_tr_traj.reshape(-1, A_tr_traj.shape[-1])
Y_tr = Y_tr_traj.reshape(-1, Y_tr_traj.shape[-1])
E_tr = E_tr_traj.reshape(-1)
# Held-out evaluation data: separate trajectories from the target embodiment.
S_te_traj, A_te_traj, Y_te_traj, _ = build_dataset(
test_embodiment, n_traj=5, n_steps=30, rng=data_rng
)
S_te = S_te_traj.reshape(-1, S_te_traj.shape[-1])
A_te = A_te_traj.reshape(-1, A_te_traj.shape[-1])
Y_te = Y_te_traj.reshape(-1, Y_te_traj.shape[-1])
# Small adaptation dataset: just two short trajectories from the target embodiment.
S_ft_traj, A_ft_traj, Y_ft_traj, _ = build_dataset(
test_embodiment, n_traj=2, n_steps=20, rng=data_rng
)
S_ft = S_ft_traj.reshape(-1, S_ft_traj.shape[-1])
A_ft = A_ft_traj.reshape(-1, A_ft_traj.shape[-1])
Y_ft = Y_ft_traj.reshape(-1, Y_ft_traj.shape[-1])
E_ft = np.full(len(S_ft), 4, dtype=np.int64)The key contrast is between the pretraining data (four embodiments, 80 trajectories) and the adaptation data (one embodiment, two trajectories). The pretrained model has seen a lot of dynamics; the adaptation data is tiny.
Pretraining a task-conditioned dynamics model
We now define a task-conditioned dynamics model. The network takes a state, an action, and a task embedding, and predicts the next state.
class TaskEmbedding(nn.Module):
"""A list containing one trainable embedding vector per task."""
def __init__(self, n_tasks, emb_dim):
super().__init__()
self.emb_dim = emb_dim
self.embs = nn.ParameterList(
[nn.Parameter(torch.randn(emb_dim) * 0.1) for _ in range(n_tasks)]
)
def forward(self, task_id):
return torch.stack([self.embs[int(i)] for i in task_id], dim=0)
class DynamicsModel(nn.Module):
def __init__(self, emb_dim=8, hidden=64):
super().__init__()
self.hidden = hidden
self.net = nn.Sequential(
nn.Linear(3 + emb_dim, hidden),
nn.ReLU(),
nn.Linear(hidden, hidden),
nn.ReLU(),
nn.Linear(hidden, 2),
)
def forward(self, s, a, e):
return self.net(torch.cat([s, a, e], dim=-1))We use a ParameterList rather than a single nn.Embedding tensor so that we can freeze individual task embeddings during adaptation. The training function accepts an explicit list of parameters to update, so we can choose between embedding-only and full fine-tuning.
def train(model, embs, S, A, Y, E, epochs, lr, params=None):
S_t = torch.from_numpy(S)
A_t = torch.from_numpy(A)
Y_t = torch.from_numpy(Y)
E_t = torch.from_numpy(E)
if params is None:
params = list(model.parameters()) + list(embs.parameters())
opt = torch.optim.Adam([p for p in params if p.requires_grad], lr=lr)
loss = None
for _ in range(epochs):
opt.zero_grad()
e = embs(E_t)
pred = model(S_t, A_t, e)
loss = ((pred - Y_t) ** 2).mean()
loss.backward()
opt.step()
return float(loss.item())Now we pretrain on four embodiments.
embs = TaskEmbedding(n_tasks=4, emb_dim=8)
model = DynamicsModel(emb_dim=8, hidden=64)
pretrain_loss = train(model, embs, S_tr, A_tr, Y_tr, E_tr, epochs=500, lr=2e-3)Final pretraining MSE: 0.00008
The pretrained model has learned a shared dynamics representation plus four task embeddings. Now we adapt.
Adapting to a new embodiment
We evaluate five strategies on the held-out embodiment:
- Persistence baseline: predict . This is the trivial baseline that any model must beat.
- Zero-shot: use the pretrained network with the mean of the training embeddings. This represents the case where you have no target data at all.
- Embedding-only adaptation: freeze the network and the four training embeddings, and fit a single new embedding to the two target trajectories.
- Full fine-tuning: unfreeze everything and continue training on the target data.
- From scratch: train a fresh model with a single embedding on the target data alone.
Let's compute each.
def evaluate_with_emb(model, S, A, Y, e):
S_t = torch.from_numpy(S)
A_t = torch.from_numpy(A)
Y_t = torch.from_numpy(Y)
e_batch = e.unsqueeze(0).expand(len(S), -1)
with torch.no_grad():
pred = model(S_t, A_t, e_batch)
return float(((pred - Y_t) ** 2).mean())
# Persistence baseline.
persistence_mse = float(
((torch.from_numpy(S_te) - torch.from_numpy(Y_te)) ** 2).mean()
)
# Zero-shot: pretrained model with the mean task embedding.
mean_emb = torch.stack([embs.embs[i] for i in range(4)]).mean(dim=0)
zeroshot_mse = evaluate_with_emb(model, S_te, A_te, Y_te, mean_emb)Now embedding-only adaptation. We build a new embedding module with five entries, copy over the pretrained ones, and freeze everything except the new embedding.
embs_adapt = TaskEmbedding(n_tasks=5, emb_dim=8)
for p in model.parameters():
p.requires_grad_(False)
with torch.no_grad():
for i in range(4):
embs_adapt.embs[i].copy_(embs.embs[i])
embs_adapt.embs[4].copy_(mean_emb)
for i in range(4):
embs_adapt.embs[i].requires_grad_(False)
_ = train(
model,
embs_adapt,
S_ft,
A_ft,
Y_ft,
E_ft,
epochs=600,
lr=5e-3,
params=[embs_adapt.embs[4]],
)
embed_mse = evaluate_with_emb(model, S_te, A_te, Y_te, embs_adapt.embs[4])Now full fine-tuning. We copy the pretrained network, initialize a fifth embedding, and update the network plus that target embedding. The four stored source embeddings are not used by the target-only batch, so we freeze them rather than counting parameters that receive no gradient.
embs_full = TaskEmbedding(n_tasks=5, emb_dim=8)
with torch.no_grad():
for i in range(4):
embs_full.embs[i].copy_(embs.embs[i])
embs_full.embs[i].requires_grad_(False)
embs_full.embs[4].copy_(mean_emb)
model_full = DynamicsModel(emb_dim=8, hidden=64)
model_full.load_state_dict(model.state_dict())
_ = train(model_full, embs_full, S_ft, A_ft, Y_ft, E_ft, epochs=400, lr=1e-3)
full_mse = evaluate_with_emb(model_full, S_te, A_te, Y_te, embs_full.embs[4])Finally, from-scratch training with a single task embedding and no pretraining.
embs_scratch = TaskEmbedding(n_tasks=1, emb_dim=8)
model_scratch = DynamicsModel(emb_dim=8, hidden=64)
E_scratch = np.zeros(len(S_ft), dtype=np.int64)
_ = train(
model_scratch,
embs_scratch,
S_ft,
A_ft,
Y_ft,
E_scratch,
epochs=400,
lr=1e-2,
)
scratch_mse = evaluate_with_emb(
model_scratch, S_te, A_te, Y_te, embs_scratch.embs[0]
)Let's collect the results.
Persistence baseline : 0.00368 Zero-shot (no target) : 0.00034 Embedding-only adapt : 0.00006 Full fine-tune : 0.00004 From scratch : 0.00103
In this fixed-seed toy run, persistence has the highest one-step MSE. Zero-shot transfer improves on it, and fitting eight embedding values improves further. Full fine-tuning achieves a similar error scale, while the target-only model has less data from which to learn the shared dynamics. These are outcomes of one controlled example, not estimates of average method performance.
Let's visualize the comparison.
The bar chart shows the tradeoff for this seed. Embedding-only adaptation uses eight trainable values and reaches the same order of one-step error as full fine-tuning. The comparison should be repeated across seeds and target systems before drawing a general performance conclusion.
Parameter-efficient adaptation with LoRA
Embedding-only adaptation is a special case of PEFT: it updates a small number of parameters. A more general version is LoRA, which we can apply to the network's linear layers. Let's implement a minimal LoRA wrapper.
class LoRALinear(nn.Module):
def __init__(self, base, rank=4, alpha=8.0):
super().__init__()
self.base = base
for p in self.base.parameters():
p.requires_grad = False
self.A = nn.Parameter(torch.zeros(rank, base.in_features))
self.B = nn.Parameter(torch.zeros(base.out_features, rank))
nn.init.normal_(self.A, std=0.01)
self.scale = alpha / rank
def forward(self, x):
return self.base(x) + self.scale * (x @ self.A.T @ self.B.T)We can wrap the first linear layer of the pretrained network. Only the LoRA matrices are trainable, so the number of updated parameters is tiny.
model_lora = DynamicsModel(emb_dim=8, hidden=64)
model_lora.load_state_dict(model.state_dict())
for p in model_lora.parameters():
p.requires_grad_(False)
# Replace the first linear layer with a LoRA-wrapped version.
model_lora.net[0] = LoRALinear(model_lora.net[0], rank=4, alpha=8.0)
lora_params = list(model_lora.net[0].parameters())
n_lora_trainable = sum(p.numel() for p in lora_params if p.requires_grad)We now adapt the LoRA parameters and a new task embedding on the target data. Everything else stays frozen.
embs_lora = TaskEmbedding(n_tasks=5, emb_dim=8)
with torch.no_grad():
for i in range(4):
embs_lora.embs[i].copy_(embs.embs[i])
embs_lora.embs[i].requires_grad_(False)
embs_lora.embs[4].copy_(mean_emb)
_ = train(
model_lora,
embs_lora,
S_ft,
A_ft,
Y_ft,
E_ft,
epochs=400,
lr=1e-2,
params=lora_params + [embs_lora.embs[4]],
)
lora_mse = evaluate_with_emb(model_lora, S_te, A_te, Y_te, embs_lora.embs[4])Trainable parameters in LoRA setup: 308 LoRA adaptation MSE: 0.00003
The accuracy numbers do not show how differently the three adaptation strategies spend parameters. Counting trainable parameters makes that contrast explicit.
# Count trainable parameters for each adaptation strategy.
full_ft_params = sum(
p.numel()
for p in list(model_full.parameters()) + list(embs_full.parameters())
if p.requires_grad
)
embed_only_params = embs_lora.emb_dim
lora_trainable_params = n_lora_trainable + embs_lora.emb_dim
iso_strategies = ["Embedding-only", "LoRA", "Full fine-tune"]
iso_params = [embed_only_params, lora_trainable_params, full_ft_params]
For this seed, the three strategies land within a small range of one-step error but use very different trainable parameter counts. A frozen-backbone design can store many small adapted variants without duplicating the backbone.
In this example, LoRA reaches the same order of error as full fine-tuning while updating only a few hundred parameters. The adapter is smaller than a full model copy. Whether accuracy remains comparable is task- and rank-dependent.
Preference-based post-training with a DPO-style objective
Finally, let's look at the post-training half. We want to shape the model so it prefers useful rollouts over harmful ones, using a preference signal. We will use the DPO loss as a conceptual demonstration.
The DPO loss takes log-probabilities of preferred and rejected outputs under the current model and a reference model, and increases the margin between them.
import torch.nn.functional as F
def dpo_loss(logp_w, logp_l, ref_logp_w, ref_logp_l, beta=0.1):
"""DPO loss for a batch of (preferred, rejected) pairs.
logp_w, logp_l: log-probabilities under the current model
ref_logp_w, ref_logp_l: log-probabilities under the reference model
"""
logits = beta * ((logp_w - ref_logp_w) - (logp_l - ref_logp_l))
return -F.logsigmoid(logits).mean()To use this loss with a deterministic world model, we need a likelihood. The simplest approach is to interpret the model's output as the mean of a Gaussian with fixed variance , so that the log-likelihood of a target given a prediction is
where:
- : the target next state or transition
- : the model's predicted next state or transition
- : the fixed variance of the Gaussian likelihood
- : a constant that does not depend on
This likelihood supplies a sequence score that a preference objective can compare. The preference label still comes from a human, expert, or explicit proxy; it does not follow from prediction accuracy. Let's simulate the loss using synthetic values shaped like log-probabilities.
# Synthetic log-probs: the current model prefers y_w but is close to the reference.
torch.manual_seed(0)
n_pairs = 64
ref_logp_w = torch.randn(n_pairs) * 0.5
ref_logp_l = torch.randn(n_pairs) * 0.5
logp_w = ref_logp_w + 0.3 + torch.randn(n_pairs) * 0.05
logp_l = ref_logp_l - 0.2 + torch.randn(n_pairs) * 0.05
dpo_value = dpo_loss(logp_w, logp_l, ref_logp_w, ref_logp_l, beta=0.5)DPO loss on synthetic batch: 0.5775
A low batch DPO loss means the current model's preferred-versus-rejected log-ratio margin is large relative to the reference. In a real loop, the compared trajectories must be sampled and labeled under a documented protocol. The reference ratios arise from the KL-regularized derivation, but they do not guarantee safety or robustness: biased preference data, poor coverage, or an inappropriate can still move the model in an undesirable direction.
The shape of the DPO objective is worth seeing directly. Sweeping the log-ratio margin for several values of traces a family of sigmoid curves. Larger values make the loss change more sharply with the margin and saturate sooner.
# Sweep the DPO margin for several betas so the effect of beta on how
# aggressively the loss saturates is visible, as the prose claims.
margin_grid = np.linspace(-3.0, 5.0, 200)
margin_t = torch.tensor(margin_grid, dtype=torch.float32)
betas_demo = [0.05, 0.1, 0.5, 1.0]
dpo_curves = {b: -F.logsigmoid(b * margin_t).numpy() for b in betas_demo}
At a zero margin every curve passes through . As the margin grows, the loss approaches zero. Larger makes the loss and its gradient transition more sharply as a function of this fixed margin; the resulting parameter update also depends on how the margin itself changes with the model parameters.