Skip to content

API reference

Fifteen modules, in the order they compose. Each page renders the module's own docstring, including the References block naming the paper the code follows, and then every public symbol.

Top level

The package re-exports the handful of names most code needs:

xwm -- building blocks for predictive world models, in JAX.

A world model answers "what happens if I do this?". An action-conditioned one answers it in a latent space of its own choosing, which is what makes it usable for robotics: you can search over imagined action sequences without rendering a single pixel.

Layout

The library separates generic machinery from families from use:
xwm.core types, base modules, EMA targets, rollouts, keys
xwm.nn layers -- attention, transformers, SimNorm, patches
xwm.encoders observation -> latent (image, video, state)
xwm.dynamics (z, a) -> z' -- the heart of a robotics model
xwm.heads reward, value, policy, Q-ensemble
xwm.masking what a JEPA predicts: blocks, tubes, time splits
xwm.objectives losses: latent prediction, SIGReg, TD, categorical
xwm.families jepa, tdmpc2, muzero
xwm.planning CEM, MPPI, gradient, MPC, MCTS
xwm.training one trainer, schedules, replay buffer
xwm.envs simulated robots (a Franka arm in Newton)
xwm.data batch streams and a synthetic controllable world
xwm.metrics probes and collapse diagnostics
xwm.plots figures, GIFs, JSON/LaTeX tables
xwm.tools checkpointing and model summaries

The three families differ only in what trains the latent space -- its own future embeddings (JEPA), reward and TD value (TD-MPC2), or search-improved targets (MuZero). They share encoders, dynamics and planners.

Conventions

Modules follow Equinox: immutable PyTrees, written for a single unbatched sample and vmaped by the caller. Batch-level entry points are the methods named loss. Images are (C, H, W), clips (T, C, H, W), token sequences (N, D), flat latents (D,), actions (A,).

key= is optional wherever a model is built (see xwm.core.random); it stays required wherever a key is consumed inside jit.

Quick start

>>> import xwm
>>> xwm.set_seed(0)
>>> model = xwm.families.jepa.lejepa(size="tiny", img_size=64, patch_size=8)
>>> trainer = xwm.training.Trainer(model, xwm.training.adamw(1e-4))

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.

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 
32
33
34
35
36
37
38
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 
40
41
42
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.

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 
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
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 
74
75
76
77
78
79
80
81
82
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 
84
85
86
87
88
89
90
91
92
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

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 
101
102
103
104
105
106
107
108
109
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

seed 

seed(value: int)

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

Source code in xwm/core/random.py 
112
113
114
115
116
117
118
119
@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)

key_source 

key_source() -> KeySource

The ambient key source, created on first use.

Source code in xwm/core/random.py 
92
93
94
95
96
97
98
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

Citations

Every module carries a References block in its docstring naming the paper the code follows, so the citation sits beside the implementation:

help(xwm.families.tdmpc2.model)
model paper
I-JEPA Assran et al., CVPR 2023 · arXiv:2301.08243
V-JEPA Bardes et al., 2024 · arXiv:2404.08471
V-JEPA 2 / -AC Assran et al., V-JEPA 2, 2025
LeJEPA Balestriero & LeCun, 2025
TD-MPC2 Hansen, Su & Wang, ICLR 2024 · arXiv:2310.16828
TD-MPC Hansen, Wang & Su, ICML 2022 · arXiv:2203.04955
MuZero Schrittwieser et al., Nature 2020 · arXiv:1911.08265
Sampled MuZero Hubert et al., ICML 2021 · arXiv:2104.06303
VICReg Bardes, Ponce & LeCun, ICLR 2022 · arXiv:2105.04906

Component-level citations live in the docstrings of the modules that implement them: SimNorm, two-hot categorical scalars, REDQ, SAC, MPPI, PUCT, Epps–Pulley, RankMe, ViT/ViViT, MAE, RoPE, LayerScale, Mish.