xwm.core

Types, base modules, EMA targets, rollouts and the ambient key source. Everything else in the library is built on these.

Core abstractions: types, base modules, EMA targets, latent rollouts.

Modules:

Name	Description
ema	Exponential-moving-average targets.
module	Base classes shared by every xwm model.
random	A default PRNG key, so key= is optional without hiding it.
rollout	Latent rollouts.
types	Shared type aliases and structural protocols for xwm.

Classes:

Name	Description
Module	An equinox.Module with a couple of conveniences.
WorldModel	A trainable world model.
KeySource	A counter-based key generator: fold_in(root, n) for n = 0, 1, 2, ...
Encoder	Maps one sample to a token sequence (N, D).
LatentDynamics	One step of action-conditioned latent dynamics: (z, a) -> z'.
Objective	Batch-level training objective, returning (scalar_loss, metrics).
Predictor	Predicts target-position embeddings from context tokens.

Functions:

Name	Description
ema_init	Create the initial target as a detached copy of model.
ema_update	target <- m * target + (1 - m) * model over inexact arrays.
stop_gradient	Detach every array leaf of tree.
batched_apply	Apply fn over the leading axis of x in chunks, then concatenate.
vmap_apply	vmap fn over a batch, splitting key per sample.
default_key	Draw the next key from the ambient source.
key_source	The ambient key source, created on first use.
resolve_key	Return key, or the next ambient key when it is None.
seed	Use value as the ambient seed for the duration of the block.
set_seed	Replace the ambient source with a fresh one for seed.
split	n keys, split from key or drawn from the ambient source.
rollout_cost	Accumulate cost_fn(z, a, t) along a rollout.
teacher_forced_rollout	One-step predictions from ground-truth latents (teacher forcing).

Module

Bases: Module

An equinox.Module with a couple of conveniences.

Everything in xwm subclasses this, so n_params and eval_mode are available on individual layers and on whole world models alike.

Note that inference is deliberately not the name of the method here: Equinox treats an inference attribute as the flag it toggles, so a method by that name would shadow it and break equinox.nn.inference_mode.

Methods:

Name	Description
eval_mode	A copy with dropout and drop-path disabled.
train_mode	A copy with dropout and drop-path re-enabled.

Attributes:

Name	Type	Description
n_params	int	Number of trainable (inexact-array) scalars in this subtree.

n_params

n_params: int

Number of trainable (inexact-array) scalars in this subtree.

eval_mode

eval_mode()

A copy with dropout and drop-path disabled.

Modules are immutable, so this returns the eval-mode model rather than mutating in place: model = model.eval_mode().

Source code in xwm/core/module.py

def eval_mode(self):
    """A copy with dropout and drop-path disabled.

    Modules are immutable, so this *returns* the eval-mode model rather
    than mutating in place: ``model = model.eval_mode()``.
    """
    return eqx.nn.inference_mode(self, value=True)

train_mode

train_mode()

A copy with dropout and drop-path re-enabled.

Source code in xwm/core/module.py

def train_mode(self):
    """A copy with dropout and drop-path re-enabled."""
    return eqx.nn.inference_mode(self, value=False)

WorldModel

Bases: Module

A trainable world model.

Subclasses implement loss, the single contract the trainer relies on. Models whose objective needs a slowly-moving teacher (I-JEPA, V-JEPA, BYOL-style asymmetry) set uses_target to True; the trainer then maintains an EMA copy of the model and passes it in as target. Models that do not need one (LeJEPA, VICReg) leave it False and receive None.

Methods:

Name	Description
loss	Scalar loss for a batched input, plus scalars to log.
trainable	Boolean tree marking which leaves the optimizer may update.
prepare_batch	Host-side hook, run before the jitted step.

loss

loss(batch: Batch, *, key: PRNGKey, target: WorldModel | None = None) -> tuple[Array, Metrics]

Scalar loss for a batched input, plus scalars to log.

Parameters:

Name	Type	Description	Default
batch	Batch	modality-specific dict (see xwm.data).	required
key	PRNGKey	RNG for masking, dropout and any stochastic objective.	required
target	WorldModel | None	EMA copy of self when uses_target, else None.	None

Source code in xwm/core/module.py

def loss(
    self,
    batch: Batch,
    *,
    key: PRNGKey,
    target: WorldModel | None = None,
) -> tuple[Array, Metrics]:
    """Scalar loss for a *batched* input, plus scalars to log.

    Args:
        batch: modality-specific dict (see :mod:`xwm.data`).
        key: RNG for masking, dropout and any stochastic objective.
        target: EMA copy of ``self`` when :attr:`uses_target`, else ``None``.
    """
    raise NotImplementedError

trainable

trainable() -> PyTree

Boolean tree marking which leaves the optimizer may update.

Defaults to every inexact array. Override to freeze part of the model -- an action-conditioned stage trained on top of a fixed video encoder is the canonical case, and freezing it here means the optimizer never allocates state for those parameters at all.

Source code in xwm/core/module.py

def trainable(self) -> PyTree:
    """Boolean tree marking which leaves the optimizer may update.

    Defaults to every inexact array. Override to freeze part of the model --
    an action-conditioned stage trained on top of a fixed video encoder is
    the canonical case, and freezing it here means the optimizer never
    allocates state for those parameters at all.
    """
    return jax.tree_util.tree_map(eqx.is_inexact_array, self)

prepare_batch

prepare_batch(batch: Batch, key: PRNGKey) -> Batch

Host-side hook, run before the jitted step.

Where mask sampling and any other combinatorial, non-jittable batch preparation belongs. It may read the model's static configuration but must not depend on its parameters -- it runs outside jit and outside the gradient.

Source code in xwm/core/module.py

def prepare_batch(self, batch: Batch, key: PRNGKey) -> Batch:
    """Host-side hook, run before the jitted step.

    Where mask sampling and any other combinatorial, non-jittable batch
    preparation belongs. It may read the model's *static* configuration but
    must not depend on its parameters -- it runs outside ``jit`` and outside
    the gradient.
    """
    return batch

KeySource

KeySource(seed: int = DEFAULT_SEED)

A counter-based key generator: fold_in(root, n) for n = 0, 1, 2, ...

Counter-based rather than split-chained so the n-th draw is a pure function of (seed, n). That makes the sequence inspectable and replayable, and avoids a long chain of splits whose state depends on every prior call.

Methods:

Name	Description
next_key	Return the next key and advance the counter.

Attributes:

Name	Type	Description
counter	int	How many keys have been drawn. Useful in tests and logs.

Source code in xwm/core/random.py

def __init__(self, seed: int = DEFAULT_SEED):
    self.seed = int(seed)
    self._root = jr.PRNGKey(self.seed)
    self._counter = 0
    self._lock = threading.Lock()

counter

counter: int

How many keys have been drawn. Useful in tests and logs.

next_key

next_key() -> Array

Return the next key and advance the counter.

Source code in xwm/core/random.py

def next_key(self) -> Array:
    """Return the next key and advance the counter."""
    with self._lock:
        index = self._counter
        self._counter += 1
    return jr.fold_in(self._root, index)

Encoder

Bases: Protocol

Maps one sample to a token sequence (N, D).

LatentDynamics

Bases: Protocol

One step of action-conditioned latent dynamics: (z, a) -> z'.

Objective

Bases: Protocol

Batch-level training objective, returning (scalar_loss, metrics).

Predictor

Bases: Protocol

Predicts target-position embeddings from context tokens.

context is (K_ctx, D) with flat grid positions context_idx; the return value is (K_tgt, D) aligned with target_idx.

ema_init

ema_init(model: PyTree) -> PyTree

Create the initial target as a detached copy of model.

Source code in xwm/core/ema.py

def ema_init(model: PyTree) -> PyTree:
    """Create the initial target as a detached copy of ``model``."""
    return jax.tree_util.tree_map(
        lambda x: jax.lax.stop_gradient(x) if eqx.is_inexact_array(x) else x, model
    )

ema_update

ema_update(target: PyTree, model: PyTree, momentum: Array | float) -> PyTree

target <- m * target + (1 - m) * model over inexact arrays.

Non-array leaves (ints, bools, static config) are taken from model so the target never drifts out of structural sync with the student.

Source code in xwm/core/ema.py

def ema_update(target: PyTree, model: PyTree, momentum: Array | float) -> PyTree:
    """``target <- m * target + (1 - m) * model`` over inexact arrays.

    Non-array leaves (ints, bools, static config) are taken from ``model`` so
    the target never drifts out of structural sync with the student.
    """
    m = jnp.asarray(momentum)

    def step(t, s):
        if eqx.is_inexact_array(t) and eqx.is_inexact_array(s):
            return jax.lax.stop_gradient(m * t + (1.0 - m) * s)
        return s

    return jax.tree_util.tree_map(step, target, model)

stop_gradient

stop_gradient(tree: PyTree) -> PyTree

Detach every array leaf of tree.

Source code in xwm/core/ema.py

def stop_gradient(tree: PyTree) -> PyTree:
    """Detach every array leaf of ``tree``."""
    return jax.tree_util.tree_map(
        lambda x: jax.lax.stop_gradient(x) if eqx.is_inexact_array(x) else x, tree
    )

batched_apply

batched_apply(fn, x, *, batch_size: int = 64) -> PyTree

Apply fn over the leading axis of x in chunks, then concatenate.

fn maps a batch to a batch (already vmapped or jitted). Use this instead of one enormous call whenever the leading axis is a dataset rather than a minibatch: encoding 16k frames at once asks the allocator for tens of gigabytes, and the failure mode is an out-of-memory abort at the end of a long run rather than anything diagnosable.

The trailing chunk may be smaller than batch_size, which costs one extra compilation under jit; padding instead would silently change the result.

Source code in xwm/core/module.py

def batched_apply(fn, x, *, batch_size: int = 64) -> PyTree:
    """Apply ``fn`` over the leading axis of ``x`` in chunks, then concatenate.

    ``fn`` maps a batch to a batch (already vmapped or jitted). Use this instead
    of one enormous call whenever the leading axis is a dataset rather than a
    minibatch: encoding 16k frames at once asks the allocator for tens of
    gigabytes, and the failure mode is an out-of-memory abort at the end of a
    long run rather than anything diagnosable.

    The trailing chunk may be smaller than ``batch_size``, which costs one extra
    compilation under ``jit``; padding instead would silently change the result.
    """
    import numpy as np

    leaves = jax.tree_util.tree_leaves(x)
    if not leaves:
        raise ValueError("nothing to apply over")
    n = leaves[0].shape[0]
    if batch_size < 1:
        raise ValueError(f"batch_size must be positive, got {batch_size}")
    outputs = [
        fn(jax.tree_util.tree_map(lambda a, s=start: a[s : s + batch_size], x))
        for start in range(0, n, batch_size)
    ]
    if len(outputs) == 1:
        return outputs[0]
    return jax.tree_util.tree_map(
        lambda *parts: jnp.concatenate(parts, axis=0) if parts[0].ndim else np.stack(parts),
        *outputs,
    )

vmap_apply

vmap_apply(fn, *args, key: PRNGKey | None = None, n: int | None = None) -> PyTree

vmap fn over a batch, splitting key per sample.

fn is a sample-level callable fn(*args, key=...). Leading axes of args are mapped; when key is None the key argument is omitted entirely (so deterministic modules need no RNG plumbing).

Source code in xwm/core/module.py

def vmap_apply(fn, *args, key: PRNGKey | None = None, n: int | None = None) -> PyTree:
    """vmap ``fn`` over a batch, splitting ``key`` per sample.

    ``fn`` is a sample-level callable ``fn(*args, key=...)``. Leading axes of
    ``args`` are mapped; when ``key`` is ``None`` the ``key`` argument is
    omitted entirely (so deterministic modules need no RNG plumbing).
    """
    if key is None:
        return jax.vmap(lambda *a: fn(*a))(*args)
    if n is None:
        n = jax.tree_util.tree_leaves(args)[0].shape[0]
    keys = jr.split(key, n)
    return jax.vmap(lambda *a: fn(*a[:-1], key=a[-1]))(*args, keys)

default_key

default_key() -> Array

Draw the next key from the ambient source.

Source code in xwm/core/random.py

def default_key() -> Array:
    """Draw the next key from the ambient source."""
    return key_source().next_key()

key_source

key_source() -> KeySource

The ambient key source, created on first use.

Source code in xwm/core/random.py

def key_source() -> KeySource:
    """The ambient key source, created on first use."""
    source = _source.get()
    if source is None:
        source = KeySource(DEFAULT_SEED)
        _source.set(source)
    return source

resolve_key

resolve_key(key: PRNGKey | None) -> Array

Return key, or the next ambient key when it is None.

The one-line helper every constructor calls, so the fallback lives in one place instead of being reimplemented per module.

Source code in xwm/core/random.py

def resolve_key(key: PRNGKey | None) -> Array:
    """Return ``key``, or the next ambient key when it is ``None``.

    The one-line helper every constructor calls, so the fallback lives in one
    place instead of being reimplemented per module.
    """
    return default_key() if key is None else key

seed

seed(value: int)

Use value as the ambient seed for the duration of the block.

Source code in xwm/core/random.py

@contextmanager
def seed(value: int):
    """Use ``value`` as the ambient seed for the duration of the block."""
    token = _source.set(KeySource(value))
    try:
        yield _source.get()
    finally:
        _source.reset(token)

set_seed

set_seed(seed: int) -> KeySource

Replace the ambient source with a fresh one for seed.

Process-wide (within the current context). For a scoped change that restores the previous source, use seed.

Source code in xwm/core/random.py

def set_seed(seed: int) -> KeySource:
    """Replace the ambient source with a fresh one for ``seed``.

    Process-wide (within the current context). For a scoped change that restores
    the previous source, use :func:`seed`.
    """
    source = KeySource(seed)
    _source.set(source)
    return source

split

split(n: int, key: PRNGKey | None = None) -> list[Array]

n keys, split from key or drawn from the ambient source.

Source code in xwm/core/random.py

def split(n: int, key: PRNGKey | None = None) -> list[Array]:
    """``n`` keys, split from ``key`` or drawn from the ambient source."""
    if n < 1:
        raise ValueError(f"need at least one key, got {n}")
    return list(jr.split(resolve_key(key), n))

rollout_cost

rollout_cost(dynamics: LatentDynamics, z0: Array, actions: Array, cost_fn: Callable[[Array, Array, int], Array], *, key: PRNGKey | None = None, cost_on: str = 'next') -> Array

Accumulate cost_fn(z, a, t) along a rollout.

Fused with the rollout so planners never materialise the whole trajectory, which matters when sampling thousands of candidate action sequences.

Parameters:

Name	Type	Description	Default
cost_on	str	which latent the cost sees. "next" -- the state the action led to. Right for a goal cost: "how close did this action get me?" "current" -- the state the action was taken from. Right for a learned reward head, which is trained as r(z_t, a_t); scoring it at z_{t+1} would evaluate the model off-distribution by one step and quietly bias every plan.	'next'

Source code in xwm/core/rollout.py

def rollout_cost(
    dynamics: LatentDynamics,
    z0: Array,
    actions: Array,
    cost_fn: Callable[[Array, Array, int], Array],
    *,
    key: PRNGKey | None = None,
    cost_on: str = "next",
) -> Array:
    """Accumulate ``cost_fn(z, a, t)`` along a rollout.

    Fused with the rollout so planners never materialise the whole trajectory,
    which matters when sampling thousands of candidate action sequences.

    Args:
        cost_on: which latent the cost sees.

            * ``"next"`` -- the state the action led to. Right for a goal cost:
              "how close did this action get me?"
            * ``"current"`` -- the state the action was taken *from*. Right for a
              learned reward head, which is trained as ``r(z_t, a_t)``; scoring
              it at ``z_{t+1}`` would evaluate the model off-distribution by one
              step and quietly bias every plan.
    """
    if cost_on not in ("next", "current"):
        raise ValueError(f"cost_on must be 'next' or 'current', got {cost_on!r}")
    horizon = actions.shape[0]
    keys = jr.split(key, horizon) if key is not None else jnp.zeros((horizon, 2), jnp.uint32)
    on_current = cost_on == "current"

    def step(carry, inputs):
        z, total, t = carry
        a, k = inputs
        cost = cost_fn(z, a, t) if on_current else jnp.zeros(())
        z_next = dynamics(z, a) if key is None else dynamics(z, a, key=k)
        if not on_current:
            cost = cost_fn(z_next, a, t)
        return (z_next, total + cost, t + 1), None

    (_, total, _), _ = jax.lax.scan(step, (z0, jnp.zeros(()), 0), (actions, keys))
    return total

teacher_forced_rollout

teacher_forced_rollout(dynamics: LatentDynamics, latents: Array, actions: Array, *, key: PRNGKey | None = None) -> Array

One-step predictions from ground-truth latents (teacher forcing).

Parameters:

Name	Type	Description	Default
latents	Array	(T, ...) encoded observations.	required
actions	Array	(T - 1, A) actions, actions[t] joining t to t+1.	required

Returns:

Type	Description
Array	(T - 1, ...) predictions of latents[1:].

Source code in xwm/core/rollout.py

def teacher_forced_rollout(
    dynamics: LatentDynamics,
    latents: Array,
    actions: Array,
    *,
    key: PRNGKey | None = None,
) -> Array:
    """One-step predictions from *ground-truth* latents (teacher forcing).

    Args:
        latents: ``(T, ...)`` encoded observations.
        actions: ``(T - 1, A)`` actions, ``actions[t]`` joining ``t`` to ``t+1``.

    Returns:
        ``(T - 1, ...)`` predictions of ``latents[1:]``.
    """
    n = actions.shape[0]
    keys = jr.split(key, n) if key is not None else jnp.zeros((n, 2), jnp.uint32)

    def one(z, a, k):
        return dynamics(z, a) if key is None else dynamics(z, a, key=k)

    return jax.vmap(one)(latents[:n], actions, keys)

rollout

Note

xwm.core.rollout is both a submodule and the function it exports. The function is documented here under its canonical path; xwm.core.rollout(...) is the way to call it.

Roll dynamics forward over an action sequence.

Parameters:

Name	Type	Description	Default
dynamics	LatentDynamics	one-step latent model (z, a) -> z'.	required
z0	Array	initial latent, any shape (typically (N, D) tokens).	required
actions	Array	(H, A) action sequence.	required
key	PRNGKey | None	optional RNG for stochastic dynamics.	None

Returns:

Type	Description
Array	(H, *z0.shape) -- the latent after each action. z0 itself is
Array	not included, so out[t] is the state reached by actions[:t + 1].

Source code in xwm/core/rollout.py

def rollout(
    dynamics: LatentDynamics,
    z0: Array,
    actions: Array,
    *,
    key: PRNGKey | None = None,
) -> Array:
    """Roll ``dynamics`` forward over an action sequence.

    Args:
        dynamics: one-step latent model ``(z, a) -> z'``.
        z0: initial latent, any shape (typically ``(N, D)`` tokens).
        actions: ``(H, A)`` action sequence.
        key: optional RNG for stochastic dynamics.

    Returns:
        ``(H, *z0.shape)`` -- the latent *after* each action. ``z0`` itself is
        not included, so ``out[t]`` is the state reached by ``actions[:t + 1]``.
    """
    horizon = actions.shape[0]
    keys = jr.split(key, horizon) if key is not None else jnp.zeros((horizon, 2), jnp.uint32)

    def step(z, inputs):
        a, k = inputs
        z_next = dynamics(z, a) if key is None else dynamics(z, a, key=k)
        return z_next, z_next

    _, traj = jax.lax.scan(step, z0, (actions, keys))
    return traj