Training¶

torch_world has no trainer. Every model is an nn.Module whose forward returns a dict with a scalar "loss", so a training loop is the usual five lines.

opt = torch.optim.AdamW(model.parameters(), lr=1e-3)
for batch in loader:
    out = model(batch["pixels"], batch["actions"])
    opt.zero_grad()
    out["loss"].backward()
    opt.step()

The other keys of out are the loss components, ready to log.

Data from episodes¶

TrajectoryDataset turns episodes (N, T, ...) into sliding windows of window steps, one dict per window. World models usually train on history + 1 frames.

from torch_world.datasets import TrajectoryDataset

data = PointMass(batch_size=2048).collect(16)
dataset = TrajectoryDataset(data, window=model.history + 1)

EMA targets¶

I-JEPA and V-JEPA predict the features of a target encoder: a frozen copy of the encoder, updated by exponential moving average after every optimizer step. The momentum usually rises linearly to 1.

from torch_world.utils import linear_schedule

for step in range(steps):
    loss = model(images, masks_enc, masks_pred)["loss"]
    opt.zero_grad()
    loss.backward()
    opt.step()
    model.update_target(linear_schedule(0.996, 1.0, step, steps))

Your own JEPA subclass gets the same machinery: build the copy with self.target_encoder = self.make_target_encoder(encoder) and call update_target. For other modules, use ema_update directly.

BatchNorm before evaluation¶

LeWM and LeJEPA projectors use BatchNorm. Their running averages lag behind features that move during training, which makes eval-mode embeddings drift from train-mode ones. Before evaluating, re-estimate the statistics as exact averages:

import itertools

from torch_world.utils import recompute_batchnorm

recompute_batchnorm(model, itertools.islice(loader, 50), lambda b: model(b["pixels"], b["actions"]))
model.eval()

Frozen encoders¶

DINO-WM, RoboJEPA and VL-JEPA keep their pre-trained encoder frozen and in eval mode, even after model.train(). Pass model.parameters() to the optimizer as usual: frozen parameters have no gradient. For VL-JEPA, use model.param_groups(lr) to train the Y-encoder at a lower learning rate.

When the encoder is expensive, encode the dataset once and train on the cached features, as MedicalVQA.encode does.