xwm.training

One family-agnostic Trainer, optimiser schedules, TrainState, and a trajectory ReplayBuffer that refuses to cross an episode boundary.

Training: one loop for every world model in xwm.

Trainer is family-agnostic -- it only needs loss, prepare_batch and trainable. What differs is where batches come from: xwm.data.iter_batches for a fixed dataset (JEPA), and ReplayBuffer for the reward-driven families, whose losses need contiguous slices of a single episode.

Modules:

Name	Description
replay	A trajectory replay buffer for the reward-driven families.
schedules	Learning-rate, weight-decay and EMA-momentum schedules.
state	Training state.
trainer	The training loop.

Classes:

Name	Description
ReplayBuffer	Fixed-capacity ring buffer over trajectory steps.
TrainState	Everything needed to resume training.
Trainer	Trains any WorldModel.

Functions:

Name	Description
adamw	AdamW with the defaults xwm uses for JEPA training.
cosine_warmup	Linear warmup then cosine decay -- the standard ViT recipe.
ema_momentum	Teacher momentum, increasing toward final over training.
weight_decay_schedule	Weight decay increasing over training, as in the DINO/I-JEPA recipes.
print_metrics	A callback that prints selected metrics.

ReplayBuffer

ReplayBuffer(capacity: int, observation_shape: tuple[int, ...], action_shape: tuple[int, ...] = (), *, extra: dict[str, tuple[int, ...]] | None = None, seed: int = 0)

Fixed-capacity ring buffer over trajectory steps.

Parameters:

Name	Type	Description	Default
capacity	int	number of steps to retain.	required
observation_shape	tuple[int, ...]	shape of a single observation.	required
action_shape	tuple[int, ...]	shape of a single action (() for a discrete index).	()
extra	dict[str, tuple[int, ...]] | None	additional per-step fields to store, as name -> shape. Use it for MuZero's search targets (value_target, policy_target).	None
seed	int	RNG seed for sampling.	0

NumPy rather than JAX: this is host-side mutable storage with random writes, which is exactly what JAX arrays are bad at. Batches are handed over as NumPy and converted at the jit boundary.

Methods:

Name	Description
add_episode	Append one episode.
sample	Sample batch_size slices of horizon steps.

Source code in xwm/training/replay.py

def __init__(
    self,
    capacity: int,
    observation_shape: tuple[int, ...],
    action_shape: tuple[int, ...] = (),
    *,
    extra: dict[str, tuple[int, ...]] | None = None,
    seed: int = 0,
):
    if capacity < 2:
        raise ValueError(f"capacity must be at least 2, got {capacity}")
    self.capacity = capacity
    self._observation = np.zeros((capacity, *observation_shape), np.float32)
    self._action = np.zeros((capacity, *action_shape), np.float32)
    self._reward = np.zeros((capacity,), np.float32)
    self._episode = np.full((capacity,), -1, np.int64)
    self._extra = {
        name: np.zeros((capacity, *shape), np.float32)
        for name, shape in (extra or {}).items()
    }
    self._cursor = 0
    self._size = 0
    self._episode_counter = 0
    self._rng = np.random.default_rng(seed)

add_episode

add_episode(observations: ndarray, actions: ndarray, rewards: ndarray, **extra: ndarray) -> None

Append one episode.

Parameters:

Name	Type	Description	Default
observations	ndarray	(T + 1, ...) -- one more than the actions, since the final observation is the state the last action led to.	required
actions	ndarray	(T, ...).	required
rewards	ndarray	(T,).	required
extra	ndarray	any fields declared in extra at construction, (T + 1, ...) or (T, ...).	{}

Source code in xwm/training/replay.py

def add_episode(
    self,
    observations: np.ndarray,
    actions: np.ndarray,
    rewards: np.ndarray,
    **extra: np.ndarray,
) -> None:
    """Append one episode.

    Args:
        observations: ``(T + 1, ...)`` -- one more than the actions, since the
            final observation is the state the last action led to.
        actions: ``(T, ...)``.
        rewards: ``(T,)``.
        extra: any fields declared in ``extra`` at construction, ``(T + 1, ...)``
            or ``(T, ...)``.
    """
    observations = np.asarray(observations, np.float32)
    actions = np.asarray(actions, np.float32)
    rewards = np.asarray(rewards, np.float32)
    steps = actions.shape[0]
    if observations.shape[0] != steps + 1:
        raise ValueError(
            f"expected {steps + 1} observations for {steps} actions, "
            f"got {observations.shape[0]}"
        )
    if rewards.shape[0] != steps:
        raise ValueError(f"expected {steps} rewards, got {rewards.shape[0]}")

    episode = self._episode_counter
    self._episode_counter += 1
    for t in range(steps):
        i = self._cursor
        self._observation[i] = observations[t]
        self._action[i] = actions[t]
        self._reward[i] = rewards[t]
        self._episode[i] = episode
        for name, buffer in self._extra.items():
            if name not in extra:
                raise KeyError(f"missing extra field {name!r}")
            buffer[i] = np.asarray(extra[name], np.float32)[t]
        self._cursor = (self._cursor + 1) % self.capacity
        self._size = min(self._size + 1, self.capacity)

sample

sample(batch_size: int, horizon: int) -> Batch

Sample batch_size slices of horizon steps.

Returns a dict with observation (B, horizon + 1, ...), action and reward (B, horizon, ...), plus any extra fields at (B, horizon + 1, ...).

Raises:

Type	Description
ValueError	if no slice of that length fits inside a single episode.

Source code in xwm/training/replay.py

def sample(self, batch_size: int, horizon: int) -> Batch:
    """Sample ``batch_size`` slices of ``horizon`` steps.

    Returns a dict with ``observation`` ``(B, horizon + 1, ...)``, ``action``
    and ``reward`` ``(B, horizon, ...)``, plus any extra fields at
    ``(B, horizon + 1, ...)``.

    Raises:
        ValueError: if no slice of that length fits inside a single episode.
    """
    if horizon < 1:
        raise ValueError(f"horizon must be positive, got {horizon}")
    starts = self._valid_starts(horizon)
    if starts.size == 0:
        raise ValueError(
            f"no episode in the buffer holds {horizon + 1} consecutive steps; "
            "collect longer episodes or lower the horizon"
        )
    chosen = self._rng.choice(starts, size=batch_size, replace=True)
    offsets = np.arange(horizon + 1)
    index = (chosen[:, None] + offsets[None, :]) % self.capacity

    batch: Batch = {
        "observation": self._observation[index],
        "action": self._action[index[:, :horizon]],
        "reward": self._reward[index[:, :horizon]],
    }
    for name, buffer in self._extra.items():
        batch[name] = buffer[index]
    return batch

TrainState

TrainState(model: WorldModel, target: WorldModel | None, opt_state: PyTree, step: Array | int = 0)

Bases: Module

Everything needed to resume training.

Attributes:

Name	Type	Description
model	WorldModel	the world model.
target	WorldModel | None	EMA teacher, or None for models that don't use one.
opt_state	PyTree	optimizer state, covering only trainable parameters.
step	Array	steps completed.

Source code in xwm/training/state.py

def __init__(
    self,
    model: WorldModel,
    target: WorldModel | None,
    opt_state: PyTree,
    step: Array | int = 0,
):
    self.model = model
    self.target = target
    self.opt_state = opt_state
    self.step = jnp.asarray(step, jnp.int32)

Trainer

Trainer(model: WorldModel, optimizer: GradientTransformation, *, ema_momentum: float | Schedule = 0.996)

Trains any WorldModel.

Parameters:

Name	Type	Description	Default
model	WorldModel	the model to train.	required
optimizer	GradientTransformation	an optax transformation (see xwm.training.adamw for sensible JEPA defaults).	required
ema_momentum	float | Schedule	teacher momentum. A float is constant; a optax.Schedule is evaluated per step (see xwm.training.ema_momentum). Ignored unless the model sets uses_target.	0.996

Example

>>> trainer = Trainer(model, adamw(cosine_warmup(1e-3, 1000)))
>>> state, history = trainer.fit(batches, steps=1000, key=key)

Methods:

Name	Description
init	Fresh training state, with the teacher and optimizer state allocated.
step	Prepare the batch, take one optimizer step, update the teacher.
fit	Train over an iterable of batches.
evaluate	Average the loss over batches with the model in eval mode.

Source code in xwm/training/trainer.py

def __init__(
    self,
    model: WorldModel,
    optimizer: optax.GradientTransformation,
    *,
    ema_momentum: float | optax.Schedule = 0.996,
):
    self.model = model
    self.optimizer = optimizer
    self.trainable_spec = model.trainable()
    self.ema_momentum = ema_momentum
    self._step_fn = self._compile()

init

init(model: WorldModel | None = None) -> TrainState

Fresh training state, with the teacher and optimizer state allocated.

Source code in xwm/training/trainer.py

def init(self, model: WorldModel | None = None) -> TrainState:
    """Fresh training state, with the teacher and optimizer state allocated."""
    model = self.model if model is None else model
    params, _ = eqx.partition(model, self.trainable_spec)
    return TrainState(
        model=model,
        target=ema_init(model) if model.uses_target else None,
        opt_state=self.optimizer.init(params),
        step=0,
    )

step

step(state: TrainState, batch: Batch, key: PRNGKey) -> tuple[TrainState, Metrics]

Prepare the batch, take one optimizer step, update the teacher.

Source code in xwm/training/trainer.py

def step(self, state: TrainState, batch: Batch, key: PRNGKey) -> tuple[TrainState, Metrics]:
    """Prepare the batch, take one optimizer step, update the teacher."""
    k_prep, k_loss = jr.split(key)
    batch = state.model.prepare_batch(batch, k_prep)
    momentum = jnp.asarray(self.momentum_at(int(state.step)), jnp.float32)
    return self._step_fn(state, batch, k_loss, momentum)

fit

fit(batches: Iterable[Batch], *, key: PRNGKey | None = None, steps: int | None = None, state: TrainState | None = None, log_every: int = 10, callbacks: Sequence[Callback] = ()) -> tuple[TrainState, list[dict[str, float]]]

Train over an iterable of batches.

Parameters:

Name	Type	Description	Default
batches	Iterable[Batch]	any iterable of batch dicts (see xwm.data).	required
steps	int | None	stop after this many steps; None exhausts the iterable.	None
state	TrainState | None	resume from this state instead of a fresh one.	None
log_every	int	how often to pull metrics back to the host. Metrics stay on device otherwise, so the loop does not block on the step it just dispatched.	10
callbacks	Sequence[Callback]	called as cb(state, metrics) on logged steps.	()

Returns:

Type	Description
tuple[TrainState, list[dict[str, float]]]	(final_state, history).

Source code in xwm/training/trainer.py

def fit(
    self,
    batches: Iterable[Batch],
    *,
    key: PRNGKey | None = None,
    steps: int | None = None,
    state: TrainState | None = None,
    log_every: int = 10,
    callbacks: Sequence[Callback] = (),
) -> tuple[TrainState, list[dict[str, float]]]:
    """Train over an iterable of batches.

    Args:
        batches: any iterable of batch dicts (see :mod:`xwm.data`).
        steps: stop after this many steps; ``None`` exhausts the iterable.
        state: resume from this state instead of a fresh one.
        log_every: how often to pull metrics back to the host. Metrics stay
            on device otherwise, so the loop does not block on the step it
            just dispatched.
        callbacks: called as ``cb(state, metrics)`` on logged steps.

    Returns:
        ``(final_state, history)``.
    """
    # Safe to default: this runs once per run, outside jit. The per-step
    # keys below are derived from it, so the whole run stays reproducible.
    key = resolve_key(key)
    state = self.init() if state is None else state
    history: list[dict[str, float]] = []
    for i, batch in enumerate(batches):
        if steps is not None and i >= steps:
            break
        state, metrics = self.step(state, batch, jr.fold_in(key, i))
        if log_every and (i % log_every == 0 or (steps is not None and i == steps - 1)):
            row = {"step": int(state.step), **{k: float(v) for k, v in metrics.items()}}
            history.append(row)
            for cb in callbacks:
                cb(state, metrics)
    return state, history

evaluate

evaluate(state: TrainState, batches: Iterable[Batch], *, key: PRNGKey | None = None) -> dict[str, float]

Average the loss over batches with the model in eval mode.

Source code in xwm/training/trainer.py

def evaluate(
    self,
    state: TrainState,
    batches: Iterable[Batch],
    *,
    key: PRNGKey | None = None,
) -> dict[str, float]:
    """Average the loss over ``batches`` with the model in eval mode."""
    key = resolve_key(key)
    model = state.model.eval_mode()
    target = None if state.target is None else state.target.eval_mode()

    @eqx.filter_jit
    def one(batch, key):
        return model.loss(batch, key=key, target=target)[1]

    totals: dict[str, float] = {}
    count = 0
    for i, batch in enumerate(batches):
        batch = model.prepare_batch(batch, jr.fold_in(key, i))
        metrics = one(batch, jr.fold_in(key, i + 1_000_000))
        for k, v in metrics.items():
            totals[k] = totals.get(k, 0.0) + float(v)
        count += 1
    return {k: v / max(count, 1) for k, v in totals.items()}

adamw

adamw(learning_rate: float | Schedule, *, weight_decay: float | Schedule = 0.05, b1: float = 0.9, b2: float = 0.95, grad_clip: float | None = 1.0) -> GradientTransformation

AdamW with the defaults xwm uses for JEPA training.

b2 = 0.95 rather than 0.999: the loss is a moving target (the teacher moves, or the regularizer's random projections change every step), so a shorter second-moment window tracks it better.

Source code in xwm/training/schedules.py

def adamw(
    learning_rate: float | optax.Schedule,
    *,
    weight_decay: float | optax.Schedule = 0.05,
    b1: float = 0.9,
    b2: float = 0.95,
    grad_clip: float | None = 1.0,
) -> optax.GradientTransformation:
    """AdamW with the defaults xwm uses for JEPA training.

    ``b2 = 0.95`` rather than 0.999: the loss is a moving target (the teacher
    moves, or the regularizer's random projections change every step), so a
    shorter second-moment window tracks it better.
    """
    tx = optax.adamw(learning_rate, b1=b1, b2=b2, weight_decay=weight_decay)
    if grad_clip is not None:
        tx = optax.chain(optax.clip_by_global_norm(grad_clip), tx)
    return tx

cosine_warmup

cosine_warmup(peak: float, total_steps: int, *, warmup_steps: int = 0, init: float = 0.0, final: float = 0.0) -> Schedule

Linear warmup then cosine decay -- the standard ViT recipe.

Source code in xwm/training/schedules.py

def cosine_warmup(
    peak: float,
    total_steps: int,
    *,
    warmup_steps: int = 0,
    init: float = 0.0,
    final: float = 0.0,
) -> optax.Schedule:
    """Linear warmup then cosine decay -- the standard ViT recipe."""
    if warmup_steps <= 0:
        return optax.cosine_decay_schedule(peak, total_steps, alpha=final / max(peak, 1e-12))
    return optax.warmup_cosine_decay_schedule(
        init_value=init,
        peak_value=peak,
        warmup_steps=warmup_steps,
        decay_steps=total_steps,
        end_value=final,
    )

ema_momentum

ema_momentum(total_steps: int, *, base: float = 0.996, final: float = 1.0) -> Schedule

Teacher momentum, increasing toward final over training.

Starting lower lets the teacher track the student while the representation is still changing fast; ending near 1.0 freezes it into a stable target once it is worth imitating. I-JEPA and V-JEPA both ramp it this way.

Source code in xwm/training/schedules.py

def ema_momentum(
    total_steps: int,
    *,
    base: float = 0.996,
    final: float = 1.0,
) -> optax.Schedule:
    """Teacher momentum, increasing toward ``final`` over training.

    Starting lower lets the teacher track the student while the representation
    is still changing fast; ending near ``1.0`` freezes it into a stable target
    once it is worth imitating. I-JEPA and V-JEPA both ramp it this way.
    """
    return optax.linear_schedule(base, final, total_steps)

weight_decay_schedule

weight_decay_schedule(total_steps: int, *, init: float = 0.04, final: float = 0.4) -> Schedule

Weight decay increasing over training, as in the DINO/I-JEPA recipes.

Source code in xwm/training/schedules.py

def weight_decay_schedule(
    total_steps: int,
    *,
    init: float = 0.04,
    final: float = 0.4,
) -> optax.Schedule:
    """Weight decay increasing over training, as in the DINO/I-JEPA recipes."""
    return optax.linear_schedule(init, final, total_steps)

print_metrics

print_metrics(every: int = 1, keys: Sequence[str] | None = None) -> Callback

A callback that prints selected metrics.

Source code in xwm/training/trainer.py

def print_metrics(every: int = 1, keys: Sequence[str] | None = None) -> Callback:
    """A callback that prints selected metrics."""
    seen = {"n": 0}

    def cb(state: TrainState, metrics: Metrics) -> None:
        seen["n"] += 1
        if (seen["n"] - 1) % every:
            return
        items = metrics if keys is None else {k: metrics[k] for k in keys if k in metrics}
        body = "  ".join(f"{k}={float(v):.4f}" for k, v in items.items())
        print(f"step {int(state.step):>6}  {body}")

    return cb