xwm.families

The three model families and a registry. Each family is a different answer to what trains the latent space; they share encoders, dynamics and planners. See Models for the narrative version.

Model families.

Three ways to learn an action-conditioned world model, all sharing the same encoders (xwm.encoders), latent dynamics (xwm.dynamics) and planners (xwm.planning). What separates them is what signal trains the latent space:

family	learning signal	planner
jepa	its own future embeddings	CEM / MPPI
tdmpc2	reward + TD value	MPPI
muzero	search-improved targets	MCTS

JEPA needs no reward, so it can pretrain on passive video -- abundant, unlabelled robot footage. TD-MPC2 and MuZero need reward and therefore interaction, but they learn a value function, so their planner can look beyond its horizon. They are complementary: a JEPA encoder is a reasonable initialisation for either.

Functions:

Name	Description
available	Registered model names.
create	Build a registered model by name.
families	The family each name belongs to.

Attributes:

Name	Type	Description
REGISTRY	dict[str, Callable[..., WorldModel]]

REGISTRY

REGISTRY: dict[str, Callable[..., WorldModel]] = {'jepa/image': _jepa_image, 'jepa/image-lejepa': _jepa_image_lejepa, 'jepa/video': _jepa_video, 'jepa/video-lejepa': _jepa_video_lejepa, 'jepa/action': _jepa_action, 'tdmpc2': _tdmpc2, 'muzero': _muzero}

available

available() -> list[str]

Registered model names.

Source code in xwm/families/registry.py

def available() -> list[str]:
    """Registered model names."""
    return sorted(REGISTRY)

create

create(name: str, **kwargs: Any) -> WorldModel

Build a registered model by name.

Source code in xwm/families/registry.py

def create(name: str, **kwargs: Any) -> WorldModel:
    """Build a registered model by name."""
    if name not in REGISTRY:
        raise KeyError(f"unknown model {name!r}; available: {available()}")
    return REGISTRY[name](**kwargs)

families

families() -> list[str]

The family each name belongs to.

Source code in xwm/families/registry.py

def families() -> list[str]:
    """The family each name belongs to."""
    return sorted({name.split("/")[0] for name in REGISTRY})

xwm.families.jepa

The JEPA family: predict masked latents, don't collapse.

Self-supervised world models that learn a representation by predicting their own embeddings at positions they were not shown. No reward, no reconstruction.

• JEPA -- the model; ijepa, vjepa, lejepa are recipes over it.
• ActionWorldModel -- the V-JEPA 2-AC stage: freeze a JEPA encoder and learn action-conditioned dynamics in its latent space. This is the family's bridge to robotics.

Modules:

Name	Description
action	Action-conditioned world models (the V-JEPA 2-AC recipe).
model	The JEPA world model: predict masked latents, and don't collapse.
predictor	The JEPA predictor: predict embeddings at positions it has not seen.
recipes	Recipes: the published JEPA configurations, assembled from the parts.

Classes:

Name	Description
ActionWorldModel	A frozen (or fine-tuned) observation encoder plus learned latent dynamics.
JEPA	Joint-embedding predictive architecture over a token grid.
JEPAPredictor	Predict embeddings at positions it has not seen, given context tokens.

Functions:

Name	Description
action_world_model	Assemble an ActionWorldModel.
ijepa	I-JEPA: multi-block latent prediction with an EMA teacher.
lejepa	LeJEPA for images: the same prediction problem, SIGReg instead of a teacher.
video_lejepa	LeJEPA for video: tube-masked prediction with SIGReg instead of a teacher.
vjepa	V-JEPA: tube-masked latent prediction with an EMA teacher.

ActionWorldModel

ActionWorldModel(encoder: VisionEncoder, dynamics: ActionConditionedPredictor, *, freeze_encoder: bool = True, teacher_forcing_weight: float = 1.0, rollout_weight: float = 1.0, loss_kind: LossKind = 'l1', normalize_target: bool = False, collapse: str = 'none', reg_weight: float = 1.0, n_proj: int = 256, video_key: str = 'video', action_key: str = 'action')

Bases: WorldModel

A frozen (or fine-tuned) observation encoder plus learned latent dynamics.

Parameters:

Name	Type	Description	Default
encoder	VisionEncoder	maps one observation to (N, D) tokens. Applied per frame, so an ImageEncoder is the usual choice; a VideoEncoder works if you feed it short clips.	required
dynamics	ActionConditionedPredictor	the one-step action-conditioned model.	required
freeze_encoder	bool	exclude the encoder from trainable, so the optimizer never touches it and allocates no state for it.	True
teacher_forcing_weight, rollout_weight		mixture of the two losses.	required
loss_kind	LossKind	how latents are compared (V-JEPA 2-AC uses L1).	'l1'
collapse	str	"none" is correct with a frozen encoder. Use "sigreg" if you unfreeze it, so the encoder cannot trivialise its own targets.	'none'
reg_weight, n_proj		SIGReg settings, used only when collapse="sigreg".	required

Methods:

Name	Description
trainable	Everything inexact, minus the encoder when it is frozen.
dynamics_fn	A plain (z, a) -> z' closure in eval mode, ready for a planner.
encode	Encode one observation to (N, D) latent tokens.
embed	Mean-pooled (D,) representation of one observation, in eval mode.
encode_sequence	Encode (T, ...) observations independently to (T, N, D).
imagine	Roll the dynamics forward from z0: (H, A) -> (H, N, D).
loss	Args:

Source code in xwm/families/jepa/action.py

def __init__(
    self,
    encoder: VisionEncoder,
    dynamics: ActionConditionedPredictor,
    *,
    freeze_encoder: bool = True,
    teacher_forcing_weight: float = 1.0,
    rollout_weight: float = 1.0,
    loss_kind: LossKind = "l1",
    normalize_target: bool = False,
    collapse: str = "none",
    reg_weight: float = 1.0,
    n_proj: int = 256,
    video_key: str = "video",
    action_key: str = "action",
):
    if encoder.embed_dim != dynamics.embed_dim:
        raise ValueError(
            f"encoder width {encoder.embed_dim} != dynamics width {dynamics.embed_dim}"
        )
    self.encoder = encoder
    self.dynamics = dynamics
    self.freeze_encoder = freeze_encoder
    self.teacher_forcing_weight = teacher_forcing_weight
    self.rollout_weight = rollout_weight
    self.loss_kind = loss_kind
    self.normalize_target = normalize_target
    self.collapse = collapse
    self.reg_weight = reg_weight
    self.n_proj = n_proj
    self.video_key = video_key
    self.action_key = action_key
    self.uses_target = False

trainable

trainable() -> PyTree

Everything inexact, minus the encoder when it is frozen.

Source code in xwm/families/jepa/action.py

def trainable(self) -> PyTree:
    """Everything inexact, minus the encoder when it is frozen."""
    spec = super().trainable()
    if self.freeze_encoder:
        spec = eqx.tree_at(
            lambda m: m.encoder,
            spec,
            jax.tree_util.tree_map(lambda _: False, spec.encoder),
        )
    return spec

dynamics_fn

dynamics_fn(*, key: PRNGKey | None = None) -> Callable[[Array, Array], Array]

A plain (z, a) -> z' closure in eval mode, ready for a planner.

Source code in xwm/families/jepa/action.py

def dynamics_fn(self, *, key: PRNGKey | None = None) -> Callable[[Array, Array], Array]:
    """A plain ``(z, a) -> z'`` closure in eval mode, ready for a planner."""
    model = self.dynamics.eval_mode()
    return lambda z, a: model(z, a, key=key)

encode

encode(observation: Array, *, key: PRNGKey | None = None) -> Array

Encode one observation to (N, D) latent tokens.

Source code in xwm/families/jepa/action.py

def encode(self, observation: Array, *, key: PRNGKey | None = None) -> Array:
    """Encode one observation to ``(N, D)`` latent tokens."""
    return self.encoder(observation, key=key)

embed

embed(observation: Array) -> Array

Mean-pooled (D,) representation of one observation, in eval mode.

Mirrors xwm.JEPA.embed, so a probe or a nearest-neighbour lookup works the same way whichever model family produced the encoder.

Source code in xwm/families/jepa/action.py

def embed(self, observation: Array) -> Array:
    """Mean-pooled ``(D,)`` representation of one observation, in eval mode.

    Mirrors :meth:`xwm.JEPA.embed`, so a probe or a nearest-neighbour lookup
    works the same way whichever model family produced the encoder.
    """
    return jnp.mean(self.encoder.eval_mode()(observation), axis=0)

encode_sequence

encode_sequence(frames: Array, *, key: PRNGKey | None = None) -> Array

Encode (T, ...) observations independently to (T, N, D).

Source code in xwm/families/jepa/action.py

def encode_sequence(self, frames: Array, *, key: PRNGKey | None = None) -> Array:
    """Encode ``(T, ...)`` observations independently to ``(T, N, D)``."""
    encoder = self.encoder if key is not None else self.encoder.eval_mode()
    if key is None:
        return jax.vmap(encoder)(frames)
    keys = jr.split(key, frames.shape[0])
    return jax.vmap(lambda f, k: encoder(f, key=k))(frames, keys)

imagine

imagine(z0: Array, actions: Array) -> Array

Roll the dynamics forward from z0: (H, A) -> (H, N, D).

Wrap with equinox.filter_jit rather than jax.jit -- the bound method carries the model's parameters, which plain jit would try to treat as static.

Source code in xwm/families/jepa/action.py

def imagine(self, z0: Array, actions: Array) -> Array:
    """Roll the dynamics forward from ``z0``: ``(H, A) -> (H, N, D)``.

    Wrap with :func:`equinox.filter_jit` rather than :func:`jax.jit` -- the
    bound method carries the model's parameters, which plain ``jit`` would
    try to treat as static.
    """
    return rollout(self.dynamics_fn(), z0, actions)

loss

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

batch: {"video": (B, T, ...), "action": (B, T - 1, A)}. action[b, t] is the action taken between frames t and t + 1.

Source code in xwm/families/jepa/action.py

def loss(
    self,
    batch: Batch,
    *,
    key: PRNGKey,
    target: ActionWorldModel | None = None,
) -> tuple[Array, Metrics]:
    """Args:
        batch: ``{"video": (B, T, ...), "action": (B, T - 1, A)}``.
            ``action[b, t]`` is the action taken between frames ``t`` and
            ``t + 1``.
    """
    frames, actions = batch[self.video_key], batch[self.action_key]
    if actions.shape[1] != frames.shape[1] - 1:
        raise ValueError(
            f"expected {frames.shape[1] - 1} actions for {frames.shape[1]} frames, "
            f"got {actions.shape[1]}"
        )
    k_enc, k_tf, k_roll, k_reg = jr.split(key, 4)

    keys = jr.split(k_enc, frames.shape[0])
    z = jax.vmap(lambda f, k: self.encode_sequence(f, key=None if self.freeze_encoder else k))(
        frames, keys
    )  # (B, T, N, D)
    if self.freeze_encoder:
        z = stop_gradient(z)
    # Targets are always detached: the dynamics model must chase the
    # representation, never move it toward something easier to predict.
    targets = stop_gradient(z[:, 1:])

    metrics: Metrics = {}
    total = jnp.zeros(())

    if self.teacher_forcing_weight:
        pred = jax.vmap(
            lambda zb, ab, k: teacher_forced_rollout(
                lambda s, a, *, key=None: self.dynamics(s, a, key=key), zb, ab, key=k
            )
        )(z, actions, jr.split(k_tf, z.shape[0]))
        loss_tf = prediction_loss(
            pred, targets, kind=self.loss_kind, normalize_target=self.normalize_target
        )
        metrics["loss_teacher_forcing"] = loss_tf
        total = total + self.teacher_forcing_weight * loss_tf

    if self.rollout_weight:
        pred = jax.vmap(
            lambda zb, ab, k: rollout(
                lambda s, a, *, key=None: self.dynamics(s, a, key=key), zb[0], ab, key=k
            )
        )(z, actions, jr.split(k_roll, z.shape[0]))
        loss_roll = prediction_loss(
            pred, targets, kind=self.loss_kind, normalize_target=self.normalize_target
        )
        metrics["loss_rollout"] = loss_roll
        total = total + self.rollout_weight * loss_roll

    if self.collapse == "sigreg":
        reg = sigreg(z, k_reg, n_proj=self.n_proj)
        metrics["loss_reg"] = reg
        total = total + self.reg_weight * reg

    metrics["latent_std"] = jnp.mean(jnp.std(z.reshape(-1, z.shape[-1]), axis=0))
    metrics["loss"] = total
    return total, metrics

JEPA

JEPA(encoder: VisionEncoder, predictor: JEPAPredictor, mask_sampler: Module, *, input_key: str = 'image', collapse: Collapse = 'ema', loss_kind: LossKind = 'smooth_l1', normalize_target: bool = False, reg_weight: float = 1.0, n_proj: int = 256, statistic: Statistic = 'epps_pulley', n_nodes: int = 32, sigma: float = 1.0)

Bases: WorldModel

Joint-embedding predictive architecture over a token grid.

Parameters:

Name	Type	Description	Default
encoder	VisionEncoder	context encoder; also the teacher unless collapse="ema".	required
predictor	JEPAPredictor	maps context tokens to target-position embeddings.	required
mask_sampler	Module	callable key -> MaskBatch (see xwm.masking).	required
input_key	str	which field of the batch holds the input tensor.	'image'
collapse	Collapse	anti-collapse strategy; see the module docstring.	'ema'
loss_kind	LossKind	how predicted and target embeddings are compared.	'smooth_l1'
normalize_target	bool	LayerNorm targets before comparing (V-JEPA does).	False
reg_weight	float	coefficient on the sigreg / vicreg penalty. This is LeJEPA's single hyperparameter.	1.0
n_proj, statistic, n_nodes, sigma		SIGReg settings.	required

Methods:

Name	Description
encode	Encode one sample to (N, D) tokens. This is the transferable part.
embed	Mean-pooled (D,) representation of one sample, in eval mode.
prepare_batch	Attach a freshly sampled mask. Runs on the host, outside jit.
predict	Predict every target block from one sample's context tokens.

Source code in xwm/families/jepa/model.py

def __init__(
    self,
    encoder: VisionEncoder,
    predictor: JEPAPredictor,
    mask_sampler: eqx.Module,
    *,
    input_key: str = "image",
    collapse: Collapse = "ema",
    loss_kind: LossKind = "smooth_l1",
    normalize_target: bool = False,
    reg_weight: float = 1.0,
    n_proj: int = 256,
    statistic: Statistic = "epps_pulley",
    n_nodes: int = 32,
    sigma: float = 1.0,
):
    if encoder.embed_dim != predictor.embed_dim:
        raise ValueError(
            f"encoder width {encoder.embed_dim} != predictor width {predictor.embed_dim}"
        )
    if tuple(encoder.grid) != tuple(predictor.grid):
        raise ValueError(f"grid mismatch: {encoder.grid} vs {predictor.grid}")
    self.encoder = encoder
    self.predictor = predictor
    self.mask_sampler = mask_sampler
    self.input_key = input_key
    self.collapse = collapse
    self.loss_kind = loss_kind
    self.normalize_target = normalize_target
    self.reg_weight = reg_weight
    self.n_proj = n_proj
    self.statistic = statistic
    self.n_nodes = n_nodes
    self.sigma = sigma
    # Only the EMA strategy needs the trainer to maintain a teacher.
    self.uses_target = collapse == "ema"

encode

encode(x: Array, *, keep: Array | None = None, key: PRNGKey | None = None) -> Array

Encode one sample to (N, D) tokens. This is the transferable part.

Source code in xwm/families/jepa/model.py

def encode(self, x: Array, *, keep: Array | None = None, key: PRNGKey | None = None) -> Array:
    """Encode one sample to ``(N, D)`` tokens. This is the transferable part."""
    return self.encoder(x, keep=keep, key=key)

embed

embed(x: Array) -> Array

Mean-pooled (D,) representation of one sample, in eval mode.

Source code in xwm/families/jepa/model.py

def embed(self, x: Array) -> Array:
    """Mean-pooled ``(D,)`` representation of one sample, in eval mode."""
    return jnp.mean(self.encoder.eval_mode()(x), axis=0)

prepare_batch

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

Attach a freshly sampled mask. Runs on the host, outside jit.

Source code in xwm/families/jepa/model.py

def prepare_batch(self, batch: Batch, key: PRNGKey) -> Batch:
    """Attach a freshly sampled mask. Runs on the host, outside ``jit``."""
    if "masks" in batch:
        return batch
    return {**batch, "masks": self.mask_sampler(key)}

predict

predict(context: Array, masks: MaskBatch, *, key: PRNGKey | None = None) -> Array

Predict every target block from one sample's context tokens.

Returns (M, K_tgt, D). Each block gets its own mask token (cycling if there are more blocks than tokens), so the predictor can distinguish the concurrent prediction problems it is being asked to solve.

Source code in xwm/families/jepa/model.py

def predict(
    self,
    context: Array,
    masks: MaskBatch,
    *,
    key: PRNGKey | None = None,
) -> Array:
    """Predict every target block from one sample's context tokens.

    Returns ``(M, K_tgt, D)``. Each block gets its own mask token (cycling
    if there are more blocks than tokens), so the predictor can distinguish
    the concurrent prediction problems it is being asked to solve.
    """
    n_mask_tokens = self.predictor.mask_tokens.shape[0]
    ids = jnp.arange(masks.n_targets) % n_mask_tokens
    return jax.vmap(
        lambda tgt, i: self.predictor(context, masks.context, tgt, mask_index=i, key=key)
    )(masks.targets, ids)

JEPAPredictor

JEPAPredictor(grid: tuple[int, ...], embed_dim: int, *, pred_dim: int, depth: int, num_heads: int, key: PRNGKey | None = None, n_mask_tokens: int = 1, pos: PosKind = 'sincos', mlp_ratio: float = 4.0, dropout: float = 0.0, drop_path: float = 0.0, layer_scale: float | None = None, remat: bool = False)

Bases: Module

Predict embeddings at positions it has not seen, given context tokens.

This is the module that makes a JEPA predictive rather than merely contrastive. Context tokens are projected into a narrower predictor width, concatenated with one learned mask token per target position -- each carrying the target's positional embedding, and nothing else about its content -- and the stack attends over the whole sequence. Reading out the mask-token positions gives the prediction.

Keeping the predictor narrower than the encoder (pred_dim < embed_dim) is deliberate: a predictor strong enough to model detail removes the pressure on the encoder to make its representation predictable.

Parameters:

Name	Type	Description	Default
grid	tuple[int, ...]	token grid the positions index into.	required
embed_dim	int	encoder width (both the input and the output width).	required
pred_dim	int	internal predictor width.	required
n_mask_tokens	int	distinct learned mask tokens. Using one per target block lets the predictor tell concurrent prediction problems apart.	1

Source code in xwm/families/jepa/predictor.py

def __init__(
    self,
    grid: tuple[int, ...],
    embed_dim: int,
    *,
    pred_dim: int,
    depth: int,
    num_heads: int,
    key: PRNGKey | None = None,
    n_mask_tokens: int = 1,
    pos: PosKind = "sincos",
    mlp_ratio: float = 4.0,
    dropout: float = 0.0,
    drop_path: float = 0.0,
    layer_scale: float | None = None,
    remat: bool = False,
):
    key = resolve_key(key)
    k_in, k_out, k_mask, k_pos, k_blocks = jr.split(key, 5)
    self.embed_in = eqx.nn.Linear(embed_dim, pred_dim, key=k_in)
    self.embed_out = eqx.nn.Linear(pred_dim, embed_dim, key=k_out)
    self.mask_tokens = 0.02 * jr.normal(k_mask, (n_mask_tokens, pred_dim))
    self.pos_embed, self.rope = make_pos(pos, tuple(grid), pred_dim, num_heads, key=k_pos)
    self.blocks = Transformer(
        pred_dim,
        depth,
        num_heads,
        key=k_blocks,
        mlp_ratio=mlp_ratio,
        dropout=dropout,
        drop_path=drop_path,
        layer_scale=layer_scale,
        rope=self.rope,
        remat=remat,
    )
    self.embed_dim = embed_dim
    self.pred_dim = pred_dim
    self.grid = tuple(grid)

action_world_model

action_world_model(*, key: PRNGKey | None = None, action_dim: int, encoder: VisionEncoder | None = None, size: str = 'small', img_size: int | tuple[int, int] = 224, patch_size: int = 16, in_channels: int = 3, action_scale: float = 1.0, conditioning: Conditioning = 'both', freeze_encoder: bool = True, dynamics_kwargs: dict | None = None, **model_kwargs) -> ActionWorldModel

Assemble an ActionWorldModel.

Pass encoder to build dynamics on top of an existing (typically pretrained and frozen) encoder -- the V-JEPA 2-AC setup. Omit it and a fresh ImageEncoder of the given size is created, which is what you want for a from-scratch experiment.

Source code in xwm/families/jepa/action.py

def action_world_model(
    *,
    key: PRNGKey | None = None,
    action_dim: int,
    encoder: VisionEncoder | None = None,
    size: str = "small",
    img_size: int | tuple[int, int] = 224,
    patch_size: int = 16,
    in_channels: int = 3,
    action_scale: float = 1.0,
    conditioning: Conditioning = "both",
    freeze_encoder: bool = True,
    dynamics_kwargs: dict | None = None,
    **model_kwargs,
) -> ActionWorldModel:
    """Assemble an :class:`ActionWorldModel`.

    Pass ``encoder`` to build dynamics on top of an existing (typically
    pretrained and frozen) encoder -- the V-JEPA 2-AC setup. Omit it and a fresh
    :class:`~xwm.image.ImageEncoder` of the given size is created, which is what
    you want for a from-scratch experiment.
    """
    k_enc, k_act, k_dyn = jr.split(resolve_key(key), 3)
    if encoder is None:
        encoder = ImageEncoder(
            key=k_enc,
            img_size=img_size,
            patch_size=patch_size,
            in_channels=in_channels,
            **preset(size, "encoder"),
        )
    action_embed = ContinuousActionEmbed(
        action_dim, encoder.embed_dim, key=k_act, scale=action_scale
    )
    dyn_cfg = {**preset(size, "predictor"), **(dynamics_kwargs or {})}
    dynamics = ActionConditionedPredictor(
        encoder.grid,
        encoder.embed_dim,
        action_embed,
        key=k_dyn,
        conditioning=conditioning,
        **dyn_cfg,
    )
    return ActionWorldModel(
        encoder, dynamics, freeze_encoder=freeze_encoder, **model_kwargs
    )

ijepa

ijepa(*, key: PRNGKey | None = None, size: str = 'small', img_size: int | tuple[int, int] = 224, patch_size: int = 16, in_channels: int = 3, n_targets: int = 4, target_scale: float = 0.15, context_scale: float | None = None, loss_kind: str = 'smooth_l1', encoder_kwargs: dict | None = None, predictor_kwargs: dict | None = None) -> JEPA

I-JEPA: multi-block latent prediction with an EMA teacher.

Train it with xwm.training.Trainer, which maintains the teacher.

Source code in xwm/families/jepa/recipes.py

def ijepa(
    *,
    key: PRNGKey | None = None,
    size: str = "small",
    img_size: int | tuple[int, int] = 224,
    patch_size: int = 16,
    in_channels: int = 3,
    n_targets: int = 4,
    target_scale: float = 0.15,
    context_scale: float | None = None,
    loss_kind: str = "smooth_l1",
    encoder_kwargs: dict | None = None,
    predictor_kwargs: dict | None = None,
) -> JEPA:
    """I-JEPA: multi-block latent prediction with an EMA teacher.

    Train it with :class:`xwm.training.Trainer`, which maintains the teacher.
    """
    return _build_image(
        key=key,
        size=size,
        img_size=img_size,
        patch_size=patch_size,
        in_channels=in_channels,
        n_targets=n_targets,
        target_scale=target_scale,
        context_scale=context_scale,
        encoder_kwargs=encoder_kwargs,
        predictor_kwargs=predictor_kwargs,
        collapse="ema",
        loss_kind=loss_kind,
    )

lejepa

lejepa(*, key: PRNGKey | None = None, size: str = 'small', img_size: int | tuple[int, int] = 224, patch_size: int = 16, in_channels: int = 3, n_targets: int = 4, target_scale: float = 0.15, context_scale: float | None = None, reg_weight: float = 1.0, n_proj: int = 256, loss_kind: str = 'smooth_l1', encoder_kwargs: dict | None = None, predictor_kwargs: dict | None = None) -> JEPA

LeJEPA for images: the same prediction problem, SIGReg instead of a teacher.

reg_weight is the only knob that governs the collapse/expressivity trade-off, and there is no teacher for the trainer to maintain.

Source code in xwm/families/jepa/recipes.py

def lejepa(
    *,
    key: PRNGKey | None = None,
    size: str = "small",
    img_size: int | tuple[int, int] = 224,
    patch_size: int = 16,
    in_channels: int = 3,
    n_targets: int = 4,
    target_scale: float = 0.15,
    context_scale: float | None = None,
    reg_weight: float = 1.0,
    n_proj: int = 256,
    loss_kind: str = "smooth_l1",
    encoder_kwargs: dict | None = None,
    predictor_kwargs: dict | None = None,
) -> JEPA:
    """LeJEPA for images: the same prediction problem, SIGReg instead of a teacher.

    ``reg_weight`` is the only knob that governs the collapse/expressivity
    trade-off, and there is no teacher for the trainer to maintain.
    """
    return _build_image(
        key=key,
        size=size,
        img_size=img_size,
        patch_size=patch_size,
        in_channels=in_channels,
        n_targets=n_targets,
        target_scale=target_scale,
        context_scale=context_scale,
        encoder_kwargs=encoder_kwargs,
        predictor_kwargs=predictor_kwargs,
        collapse="sigreg",
        loss_kind=loss_kind,
        reg_weight=reg_weight,
        n_proj=n_proj,
    )

video_lejepa

video_lejepa(*, key: PRNGKey | None = None, size: str = 'small', img_size: int | tuple[int, int] = 224, patch_size: int = 16, num_frames: int = 16, tubelet_size: int = 2, in_channels: int = 3, n_targets: int = 8, spatial_scale: float = 0.15, temporal_extent: int | None = None, context_scale: float | None = None, reg_weight: float = 1.0, n_proj: int = 256, loss_kind: str = 'l1', encoder_kwargs: dict | None = None, predictor_kwargs: dict | None = None) -> JEPA

LeJEPA for video: tube-masked prediction with SIGReg instead of a teacher.

Source code in xwm/families/jepa/recipes.py

def video_lejepa(
    *,
    key: PRNGKey | None = None,
    size: str = "small",
    img_size: int | tuple[int, int] = 224,
    patch_size: int = 16,
    num_frames: int = 16,
    tubelet_size: int = 2,
    in_channels: int = 3,
    n_targets: int = 8,
    spatial_scale: float = 0.15,
    temporal_extent: int | None = None,
    context_scale: float | None = None,
    reg_weight: float = 1.0,
    n_proj: int = 256,
    loss_kind: str = "l1",
    encoder_kwargs: dict | None = None,
    predictor_kwargs: dict | None = None,
) -> JEPA:
    """LeJEPA for video: tube-masked prediction with SIGReg instead of a teacher."""
    return _build_video(
        key=key,
        size=size,
        img_size=img_size,
        patch_size=patch_size,
        num_frames=num_frames,
        tubelet_size=tubelet_size,
        in_channels=in_channels,
        n_targets=n_targets,
        spatial_scale=spatial_scale,
        temporal_extent=temporal_extent,
        context_scale=context_scale,
        encoder_kwargs=encoder_kwargs,
        predictor_kwargs=predictor_kwargs,
        collapse="sigreg",
        loss_kind=loss_kind,
        reg_weight=reg_weight,
        n_proj=n_proj,
    )

vjepa

vjepa(*, key: PRNGKey | None = None, size: str = 'small', img_size: int | tuple[int, int] = 224, patch_size: int = 16, num_frames: int = 16, tubelet_size: int = 2, in_channels: int = 3, n_targets: int = 8, spatial_scale: float = 0.15, temporal_extent: int | None = None, context_scale: float | None = None, loss_kind: str = 'l1', normalize_target: bool = True, encoder_kwargs: dict | None = None, predictor_kwargs: dict | None = None) -> JEPA

V-JEPA: tube-masked latent prediction with an EMA teacher.

Defaults follow the paper: L1 loss on LayerNorm-ed targets, eight short-range tubes spanning the whole clip.

Source code in xwm/families/jepa/recipes.py

def vjepa(
    *,
    key: PRNGKey | None = None,
    size: str = "small",
    img_size: int | tuple[int, int] = 224,
    patch_size: int = 16,
    num_frames: int = 16,
    tubelet_size: int = 2,
    in_channels: int = 3,
    n_targets: int = 8,
    spatial_scale: float = 0.15,
    temporal_extent: int | None = None,
    context_scale: float | None = None,
    loss_kind: str = "l1",
    normalize_target: bool = True,
    encoder_kwargs: dict | None = None,
    predictor_kwargs: dict | None = None,
) -> JEPA:
    """V-JEPA: tube-masked latent prediction with an EMA teacher.

    Defaults follow the paper: L1 loss on LayerNorm-ed targets, eight
    short-range tubes spanning the whole clip.
    """
    return _build_video(
        key=key,
        size=size,
        img_size=img_size,
        patch_size=patch_size,
        num_frames=num_frames,
        tubelet_size=tubelet_size,
        in_channels=in_channels,
        n_targets=n_targets,
        spatial_scale=spatial_scale,
        temporal_extent=temporal_extent,
        context_scale=context_scale,
        encoder_kwargs=encoder_kwargs,
        predictor_kwargs=predictor_kwargs,
        collapse="ema",
        loss_kind=loss_kind,
        normalize_target=normalize_target,
    )

xwm.families.tdmpc2

TD-MPC2: latent dynamics trained by reward and temporal-difference value.

Modules:

Name	Description
model	TD-MPC2: a latent world model trained by reward and TD value.

Classes:

Name	Description
TDMPC2	A TD-MPC2 agent.

Functions:

Name	Description
planner	An MPPI planner wired to model's own reward and value heads.
tdmpc2	Assemble a TD-MPC2 agent.

TDMPC2

TDMPC2(encoder: Encoder, dynamics: MLPDynamics, reward: ScalarHead, critic: QEnsemble, policy: GaussianPolicy, *, horizon: int = 3, discount: float = 0.99, consistency_coef: float = 20.0, reward_coef: float = 0.1, value_coef: float = 0.1, rho: float = 0.5, entropy_coef: float = 0.0001, obs_key: str = 'observation', action_key: str = 'action', reward_key: str = 'reward')

Bases: WorldModel

A TD-MPC2 agent.

Parameters:

Name	Type	Description	Default
encoder	Encoder	observation encoder. Token output is mean-pooled to a flat latent, since the dynamics and heads here are MLPs.	required
dynamics	MLPDynamics	latent dynamics; SimNorm-normalised by default.	required
reward	ScalarHead	reward head, action-conditioned.	required
critic	QEnsemble	Q-ensemble.	required
policy	GaussianPolicy	policy prior.	required
horizon	int	unroll length for the consistency/reward/value losses.	3
discount	float	RL discount.	0.99
consistency_coef, reward_coef, value_coef		loss weights.	required
rho	float	per-step decay on the unroll losses. Later steps are less reliable because they start from a predicted latent, so they count less.	0.5
entropy_coef	float	SAC-style entropy bonus for the policy.	0.0001

Methods:

Name	Description
encode	Observation -> flat latent (D,).
dynamics_fn	A plain (z, a) -> z' closure in eval mode, ready for a planner.
act	Policy-prior action, no planning. See xwm.planning to plan.
td_target	r + gamma * Q_target(z', pi(z')), fully detached.
loss	Args:

Source code in xwm/families/tdmpc2/model.py

def __init__(
    self,
    encoder: Encoder,
    dynamics: MLPDynamics,
    reward: ScalarHead,
    critic: QEnsemble,
    policy: GaussianPolicy,
    *,
    horizon: int = 3,
    discount: float = 0.99,
    consistency_coef: float = 20.0,
    reward_coef: float = 0.1,
    value_coef: float = 0.1,
    rho: float = 0.5,
    entropy_coef: float = 1e-4,
    obs_key: str = "observation",
    action_key: str = "action",
    reward_key: str = "reward",
):
    self.encoder = encoder
    self.dynamics = dynamics
    self.reward = reward
    self.critic = critic
    self.policy = policy
    self.horizon = horizon
    self.discount = discount
    self.consistency_coef = consistency_coef
    self.reward_coef = reward_coef
    self.value_coef = value_coef
    self.rho = rho
    self.entropy_coef = entropy_coef
    self.obs_key = obs_key
    self.action_key = action_key
    self.reward_key = reward_key
    # The critic target is an EMA of the whole model; only `critic` is read.
    self.uses_target = True

encode

encode(observation: Array, *, key: PRNGKey | None = None) -> Array

Observation -> flat latent (D,).

Encoders emit (N, D) tokens; the dynamics and heads here are MLPs over a single vector, so tokens are mean-pooled.

Source code in xwm/families/tdmpc2/model.py

def encode(self, observation: Array, *, key: PRNGKey | None = None) -> Array:
    """Observation -> flat latent ``(D,)``.

    Encoders emit ``(N, D)`` tokens; the dynamics and heads here are MLPs
    over a single vector, so tokens are mean-pooled.
    """
    tokens = self.encoder(observation, key=key)
    return jnp.mean(tokens, axis=0)

dynamics_fn

dynamics_fn()

A plain (z, a) -> z' closure in eval mode, ready for a planner.

Source code in xwm/families/tdmpc2/model.py

def dynamics_fn(self):
    """A plain ``(z, a) -> z'`` closure in eval mode, ready for a planner."""
    model = self.dynamics.eval_mode()
    return lambda z, a: model(z, a)

act

act(observation: Array, *, key: PRNGKey | None = None) -> Array

Policy-prior action, no planning. See xwm.planning to plan.

Source code in xwm/families/tdmpc2/model.py

def act(self, observation: Array, *, key: PRNGKey | None = None) -> Array:
    """Policy-prior action, no planning. See :mod:`xwm.planning` to plan."""
    return self.policy.eval_mode().act(self.encode(observation), key=key)

td_target

td_target(target: TDMPC2, z_next: Array, reward: Array, key: PRNGKey) -> Array

r + gamma * Q_target(z', pi(z')), fully detached.

The next action comes from the current policy but the Q from the EMA target critic: bootstrapping off the online critic is what makes value learning diverge.

Source code in xwm/families/tdmpc2/model.py

def td_target(
    self,
    target: TDMPC2,
    z_next: Array,
    reward: Array,
    key: PRNGKey,
) -> Array:
    """``r + gamma * Q_target(z', pi(z'))``, fully detached.

    The next action comes from the *current* policy but the Q from the EMA
    target critic: bootstrapping off the online critic is what makes value
    learning diverge.
    """
    k_policy, k_q = jr.split(key)
    next_action = self.policy.act(z_next, key=k_policy)
    q_next = target.critic.pessimistic(z_next, next_action, key=k_q)
    return stop_gradient(reward + self.discount * q_next)

loss

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

batch: {"observation": (B, H + 1, ...), "action": (B, H, A), "reward": (B, H)} -- a contiguous slice of a trajectory.

Source code in xwm/families/tdmpc2/model.py

def loss(
    self,
    batch: Batch,
    *,
    key: PRNGKey,
    target: TDMPC2 | None = None,
) -> tuple[Array, Metrics]:
    """Args:
        batch: ``{"observation": (B, H + 1, ...), "action": (B, H, A),
            "reward": (B, H)}`` -- a contiguous slice of a trajectory.
    """
    if target is None:
        raise ValueError("TD-MPC2 needs an EMA target; the Trainer supplies it")
    observations = batch[self.obs_key]
    actions = batch[self.action_key]
    rewards = batch[self.reward_key]
    horizon = min(self.horizon, actions.shape[1])
    k_enc, k_td, k_q, k_pi = jr.split(key, 4)

    # Encode the whole window once. Targets are detached: the consistency
    # loss must pull the *dynamics* toward the encoder, never the reverse,
    # or the pair collapses to a constant latent.
    encode = jax.vmap(jax.vmap(lambda o: self.encode(o)))
    latents = encode(observations)  # (B, H + 1, D)
    targets_z = stop_gradient(latents[:, 1:])

    consistency = jnp.zeros(())
    reward_loss = jnp.zeros(())
    value_loss = jnp.zeros(())
    z = latents[:, 0]
    for step in range(horizon):
        action = actions[:, step]
        weight = self.rho**step

        z_pred = jax.vmap(self.dynamics)(z, action)
        consistency = consistency + weight * jnp.mean(
            jnp.square(z_pred - targets_z[:, step])
        )
        reward_loss = reward_loss + weight * jnp.mean(
            jax.vmap(self.reward.loss)(z, rewards[:, step], action)
        )

        td = jax.vmap(lambda zn, r, k: self.td_target(target, zn, r, k))(
            targets_z[:, step],
            rewards[:, step],
            jr.split(jr.fold_in(k_td, step), z.shape[0]),
        )
        value_loss = value_loss + weight * jnp.mean(
            jax.vmap(self.critic.loss)(z, action, td)
        )
        z = z_pred

    # Policy: maximise Q with an entropy bonus, on detached latents so the
    # actor cannot reshape the world model to make itself look good.
    flat_z = stop_gradient(latents.reshape(-1, latents.shape[-1]))
    keys = jr.split(k_pi, flat_z.shape[0])
    sampled = jax.vmap(self.policy.sample)(flat_z, keys)
    q_keys = jr.split(k_q, flat_z.shape[0])
    q_values = jax.vmap(
        lambda z_, a_, k_: self.critic.pessimistic(z_, a_, key=k_)
    )(flat_z, sampled.action, q_keys)
    policy_loss = jnp.mean(self.entropy_coef * sampled.log_prob - q_values)

    total = (
        self.consistency_coef * consistency
        + self.reward_coef * reward_loss
        + self.value_coef * value_loss
        + policy_loss
    )
    return total, {
        "loss": total,
        "loss_consistency": consistency,
        "loss_reward": reward_loss,
        "loss_value": value_loss,
        "loss_policy": policy_loss,
        "q_mean": jnp.mean(q_values),
        "entropy": -jnp.mean(sampled.log_prob),
        "latent_std": jnp.mean(jnp.std(flat_z, axis=0)),
    }

planner

planner(model: TDMPC2, *, horizon: int | None = None, n_samples: int = 512, n_iters: int = 6, temperature: float = 0.5, noise_std: float = 0.5)

An MPPI planner wired to model's own reward and value heads.

Returns (planner, cost_fn); call planner.plan(key, model.dynamics_fn(), z, cost_fn). The planner scores candidates by discounted predicted reward plus a terminal value bootstrap, evaluated at the pre-transition latent because that is how the reward head was trained.

Source code in xwm/families/tdmpc2/model.py

def planner(
    model: TDMPC2,
    *,
    horizon: int | None = None,
    n_samples: int = 512,
    n_iters: int = 6,
    temperature: float = 0.5,
    noise_std: float = 0.5,
):
    """An MPPI planner wired to ``model``'s own reward and value heads.

    Returns ``(planner, cost_fn)``; call ``planner.plan(key, model.dynamics_fn(),
    z, cost_fn)``. The planner scores candidates by discounted predicted reward
    plus a terminal value bootstrap, evaluated at the pre-transition latent
    because that is how the reward head was trained.
    """
    from ...planning.cost import return_cost
    from ...planning.sampling import MPPI

    horizon = horizon or model.horizon
    reward = model.reward.eval_mode()
    critic = model.critic.eval_mode()
    policy = model.policy.eval_mode()

    search = MPPI(
        horizon,
        model.reward.action_dim,
        n_samples=n_samples,
        n_iters=n_iters,
        temperature=temperature,
        noise_std=noise_std,
        cost_on="current",
    )
    cost = return_cost(
        lambda z, a: reward.value(z, a),
        lambda z: critic.pessimistic(z, policy.act(z)),
        horizon=horizon,
        discount=model.discount,
    )
    return search, cost

tdmpc2

tdmpc2(*, action_dim: int, encoder: Encoder | None = None, observation: str = 'state', state_dim: int | None = None, img_size: int = 64, patch_size: int = 8, latent_dim: int = 512, hidden_dim: int = 512, horizon: int = 3, n_bins: int = 101, key: PRNGKey | None = None, **model_kwargs) -> TDMPC2

Assemble a TD-MPC2 agent.

Parameters:

Name	Type	Description	Default
observation	str	"state" for a vector observation (needs state_dim) or "image" for pixels.	'state'
encoder	Encoder | None	supply your own -- e.g. a frozen JEPA encoder -- and the rest is built around it.	None
latent_dim	int	width of the pooled latent the dynamics and heads act on.	512

Source code in xwm/families/tdmpc2/model.py

def tdmpc2(
    *,
    action_dim: int,
    encoder: Encoder | None = None,
    observation: str = "state",
    state_dim: int | None = None,
    img_size: int = 64,
    patch_size: int = 8,
    latent_dim: int = 512,
    hidden_dim: int = 512,
    horizon: int = 3,
    n_bins: int = 101,
    key: PRNGKey | None = None,
    **model_kwargs,
) -> TDMPC2:
    """Assemble a TD-MPC2 agent.

    Args:
        observation: ``"state"`` for a vector observation (needs ``state_dim``)
            or ``"image"`` for pixels.
        encoder: supply your own -- e.g. a frozen JEPA encoder -- and the rest
            is built around it.
        latent_dim: width of the pooled latent the dynamics and heads act on.
    """
    from ...encoders.image import ImageEncoder

    key = resolve_key(key)
    k_enc, k_dyn, k_rew, k_q, k_pi = jr.split(key, 5)
    if encoder is None:
        if observation == "state":
            if state_dim is None:
                raise ValueError('observation="state" needs state_dim')
            encoder = StateEncoder(state_dim, latent_dim, key=k_enc)
        elif observation == "image":
            encoder = ImageEncoder(
                key=k_enc, img_size=img_size, patch_size=patch_size,
                embed_dim=latent_dim, depth=4, num_heads=8,
            )
        else:
            raise ValueError(f"unknown observation kind {observation!r}")
    latent_dim = encoder.embed_dim

    scalar = CategoricalScalar(n_bins=n_bins)
    return TDMPC2(
        encoder,
        MLPDynamics(latent_dim, action_dim, key=k_dyn, hidden_dim=hidden_dim),
        ScalarHead(latent_dim, action_dim=action_dim, key=k_rew,
                   hidden_dim=hidden_dim, scalar=scalar),
        QEnsemble(latent_dim, action_dim, key=k_q, hidden_dim=hidden_dim, scalar=scalar),
        GaussianPolicy(latent_dim, action_dim, key=k_pi, hidden_dim=hidden_dim),
        horizon=horizon,
        **model_kwargs,
    )

xwm.families.muzero

MuZero: a latent model trained to agree with its own tree search.

Modules:

Name	Description
model	MuZero: a latent model trained to make search consistent.
targets	Building MuZero's training targets from self-play.

Classes:

Name	Description
MuZero	A MuZero agent over a discrete action space.
PolicyHead	s -> logits over a discrete action set.

Functions:

Name	Description
muzero	Assemble a MuZero agent over a discrete action space.
n_step_value_targets	(T + 1,) value targets for an episode of T transitions.

MuZero

MuZero(encoder: Encoder, dynamics: MLPDynamics, reward: ScalarHead, value: ScalarHead, policy: PolicyHead, *, n_actions: int, horizon: int = 5, reward_coef: float = 1.0, value_coef: float = 0.25, policy_coef: float = 1.0, obs_key: str = 'observation', action_key: str = 'action', reward_key: str = 'reward', value_key: str = 'value_target', policy_key: str = 'policy_target')

Bases: WorldModel

A MuZero agent over a discrete action space.

Parameters:

Name	Type	Description	Default
encoder	Encoder	the representation network.	required
dynamics	MLPDynamics	the recurrent latent dynamics.	required
reward	ScalarHead	reward head, conditioned on the one-hot action.	required
value	ScalarHead	value head.	required
policy	PolicyHead	policy logits head.	required
n_actions	int	size of the action set.	required
horizon	int	unroll length during training.	5
reward_coef, value_coef, policy_coef		loss weights.	required

Methods:

Name	Description
represent	Observation -> latent (D,).
recurrent	(z, action_index) -> (z', reward). The signature MCTS wants.
predict	z -> (policy_logits, value). The other signature MCTS wants.
search_fns	Eval-mode (recurrent, predict) closures for MCTS.
loss	Args:

Source code in xwm/families/muzero/model.py

def __init__(
    self,
    encoder: Encoder,
    dynamics: MLPDynamics,
    reward: ScalarHead,
    value: ScalarHead,
    policy: PolicyHead,
    *,
    n_actions: int,
    horizon: int = 5,
    reward_coef: float = 1.0,
    value_coef: float = 0.25,
    policy_coef: float = 1.0,
    obs_key: str = "observation",
    action_key: str = "action",
    reward_key: str = "reward",
    value_key: str = "value_target",
    policy_key: str = "policy_target",
):
    self.encoder = encoder
    self.dynamics = dynamics
    self.reward = reward
    self.value = value
    self.policy = policy
    self.n_actions = n_actions
    self.horizon = horizon
    self.reward_coef = reward_coef
    self.value_coef = value_coef
    self.policy_coef = policy_coef
    self.obs_key = obs_key
    self.action_key = action_key
    self.reward_key = reward_key
    self.value_key = value_key
    self.policy_key = policy_key
    self.uses_target = False

represent

represent(observation: Array, *, key: PRNGKey | None = None) -> Array

Observation -> latent (D,).

Source code in xwm/families/muzero/model.py

def represent(self, observation: Array, *, key: PRNGKey | None = None) -> Array:
    """Observation -> latent ``(D,)``."""
    return jnp.mean(self.encoder(observation, key=key), axis=0)

recurrent

recurrent(z: Array, action: Array) -> tuple[Array, Array]

(z, action_index) -> (z', reward). The signature MCTS wants.

Source code in xwm/families/muzero/model.py

def recurrent(self, z: Array, action: Array) -> tuple[Array, Array]:
    """``(z, action_index) -> (z', reward)``. The signature MCTS wants."""
    one_hot = jax.nn.one_hot(action, self.n_actions)
    return self.dynamics(z, one_hot), self.reward.value(z, one_hot)

predict

predict(z: Array) -> tuple[Array, Array]

z -> (policy_logits, value). The other signature MCTS wants.

Source code in xwm/families/muzero/model.py

def predict(self, z: Array) -> tuple[Array, Array]:
    """``z -> (policy_logits, value)``. The other signature MCTS wants."""
    return self.policy(z), self.value.value(z)

search_fns

search_fns()

Eval-mode (recurrent, predict) closures for MCTS.

Source code in xwm/families/muzero/model.py

def search_fns(self):
    """Eval-mode ``(recurrent, predict)`` closures for :class:`~xwm.planning.MCTS`."""
    model = self.eval_mode()
    return model.recurrent, model.predict

loss

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

batch: exactly what xwm.training.ReplayBuffer.sample returns -- {"observation": (B, H + 1, ...), "action": (B, H), "reward": (B, H), "value_target": (B, H + 1), "policy_target": (B, H + 1, n_actions)}.

Only ``observation[:, 0]`` is read. Everything after it comes
from the model's own unroll, and that is the whole point: MuZero
never re-encodes an observation mid-unroll, so the latent is free
to be whatever makes the predictions work. The later
observations are accepted (and ignored) so the same batch feeds
both this family and TD-MPC2.

Source code in xwm/families/muzero/model.py

def loss(
    self,
    batch: Batch,
    *,
    key: PRNGKey,
    target: MuZero | None = None,
) -> tuple[Array, Metrics]:
    """Args:
        batch: exactly what :meth:`xwm.training.ReplayBuffer.sample` returns --
            ``{"observation": (B, H + 1, ...), "action": (B, H),
            "reward": (B, H), "value_target": (B, H + 1),
            "policy_target": (B, H + 1, n_actions)}``.

            Only ``observation[:, 0]`` is read. Everything after it comes
            from the model's own unroll, and that is the whole point: MuZero
            never re-encodes an observation mid-unroll, so the latent is free
            to be whatever makes the predictions work. The later
            observations are accepted (and ignored) so the same batch feeds
            both this family and TD-MPC2.
    """
    observations = batch[self.obs_key]
    # Take the root observation. A time axis is present because the batch
    # comes from a trajectory buffer; MuZero simply does not use the rest.
    if observations.ndim > 1 and observations.shape[1] == batch[self.action_key].shape[1] + 1:
        observations = observations[:, 0]
    actions = batch[self.action_key].astype(jnp.int32)
    rewards = batch[self.reward_key]
    value_targets = batch[self.value_key]
    policy_targets = batch[self.policy_key]
    horizon = min(self.horizon, actions.shape[1])

    z = jax.vmap(self.represent)(observations)  # (B, D)

    # Step 0 has no reward to predict: it is the root.
    logits = jax.vmap(self.policy)(z)
    policy_loss = cross_entropy(logits, policy_targets[:, 0])
    value_loss = jnp.mean(jax.vmap(self.value.loss)(z, value_targets[:, 0]))
    reward_loss = jnp.zeros(())

    for step in range(horizon):
        action = actions[:, step]
        one_hot = jax.nn.one_hot(action, self.n_actions)
        reward_loss = reward_loss + jnp.mean(
            jax.vmap(self.reward.loss)(z, rewards[:, step], one_hot)
        )
        z = jax.vmap(self.dynamics)(z, one_hot)
        # Scale the recurrent path so the gradient reaching the encoder does
        # not grow with the unroll length.
        z = 0.5 * z + 0.5 * stop_gradient(z)
        policy_loss = policy_loss + cross_entropy(
            jax.vmap(self.policy)(z), policy_targets[:, step + 1]
        )
        value_loss = value_loss + jnp.mean(
            jax.vmap(self.value.loss)(z, value_targets[:, step + 1])
        )

    scale = 1.0 / (horizon + 1)
    total = scale * (
        self.reward_coef * reward_loss
        + self.value_coef * value_loss
        + self.policy_coef * policy_loss
    )
    return total, {
        "loss": total,
        "loss_reward": scale * reward_loss,
        "loss_value": scale * value_loss,
        "loss_policy": scale * policy_loss,
        "latent_std": jnp.mean(jnp.std(z, axis=0)),
    }

PolicyHead

PolicyHead(latent_dim: int, n_actions: int, *, key: PRNGKey | None = None, hidden_dim: int = 256, depth: int = 2)

Bases: WorldModel

s -> logits over a discrete action set.

Source code in xwm/families/muzero/model.py

def __init__(
    self,
    latent_dim: int,
    n_actions: int,
    *,
    key: PRNGKey | None = None,
    hidden_dim: int = 256,
    depth: int = 2,
):
    key = resolve_key(key)
    keys = jr.split(key, depth + 1)
    dims = [latent_dim] + [hidden_dim] * depth
    self.layers = [eqx.nn.Linear(dims[i], hidden_dim, key=keys[i]) for i in range(depth)]
    self.norms = [LayerNorm(hidden_dim) for _ in range(depth)]
    self.out = eqx.nn.Linear(hidden_dim, n_actions, key=keys[-1])
    self.n_actions = n_actions

muzero

muzero(*, n_actions: int, encoder: Encoder | None = None, observation: str = 'state', state_dim: int | None = None, img_size: int = 64, patch_size: int = 8, latent_dim: int = 256, hidden_dim: int = 256, horizon: int = 5, n_bins: int = 101, key: PRNGKey | None = None, **model_kwargs) -> MuZero

Assemble a MuZero agent over a discrete action space.

Source code in xwm/families/muzero/model.py

def muzero(
    *,
    n_actions: int,
    encoder: Encoder | None = None,
    observation: str = "state",
    state_dim: int | None = None,
    img_size: int = 64,
    patch_size: int = 8,
    latent_dim: int = 256,
    hidden_dim: int = 256,
    horizon: int = 5,
    n_bins: int = 101,
    key: PRNGKey | None = None,
    **model_kwargs,
) -> MuZero:
    """Assemble a MuZero agent over a discrete action space."""
    from ...encoders.image import ImageEncoder

    key = resolve_key(key)
    k_enc, k_dyn, k_rew, k_val, k_pi = jr.split(key, 5)
    if encoder is None:
        if observation == "state":
            if state_dim is None:
                raise ValueError('observation="state" needs state_dim')
            encoder = StateEncoder(state_dim, latent_dim, key=k_enc)
        elif observation == "image":
            encoder = ImageEncoder(
                key=k_enc, img_size=img_size, patch_size=patch_size,
                embed_dim=latent_dim, depth=4, num_heads=8,
            )
        else:
            raise ValueError(f"unknown observation kind {observation!r}")
    latent_dim = encoder.embed_dim
    scalar = CategoricalScalar(n_bins=n_bins)
    return MuZero(
        encoder,
        # LayerNorm, not SimNorm: simplicial normalisation is TD-MPC2's
        # contribution, and applied here it pins every latent near the uniform
        # point of each simplex, leaving almost no variance for the prediction
        # heads to read.
        MLPDynamics(
            latent_dim, n_actions, key=k_dyn, hidden_dim=hidden_dim,
            normalize="layernorm",
        ),
        ScalarHead(latent_dim, action_dim=n_actions, key=k_rew,
                   hidden_dim=hidden_dim, scalar=scalar),
        ScalarHead(latent_dim, key=k_val, hidden_dim=hidden_dim, scalar=scalar),
        PolicyHead(latent_dim, n_actions, key=k_pi, hidden_dim=hidden_dim),
        n_actions=n_actions,
        horizon=horizon,
        **model_kwargs,
    )

n_step_value_targets

n_step_value_targets(rewards: ndarray, root_values: ndarray, *, discount: float = 0.997, n_steps: int = 5) -> ndarray

(T + 1,) value targets for an episode of T transitions.

Parameters:

Name	Type	Description	Default
rewards	ndarray	(T,) rewards received.	required
root_values	ndarray	(T + 1,) MCTS root values, one per visited state.	required
discount	float	RL discount.	0.997
n_steps	int	bootstrap horizon. Longer reduces bias and adds variance; beyond the episode end the sum simply truncates.	5

Returns:

Type	Description
ndarray	target[t] = sum_{k<n} gamma^k r_{t+k} + gamma^n V_root[t+n], with the
ndarray	bootstrap dropped when t + n runs past the episode.

Source code in xwm/families/muzero/targets.py

def n_step_value_targets(
    rewards: np.ndarray,
    root_values: np.ndarray,
    *,
    discount: float = 0.997,
    n_steps: int = 5,
) -> np.ndarray:
    """``(T + 1,)`` value targets for an episode of ``T`` transitions.

    Args:
        rewards: ``(T,)`` rewards received.
        root_values: ``(T + 1,)`` MCTS root values, one per visited state.
        discount: RL discount.
        n_steps: bootstrap horizon. Longer reduces bias and adds variance;
            beyond the episode end the sum simply truncates.

    Returns:
        ``target[t] = sum_{k<n} gamma^k r_{t+k} + gamma^n V_root[t+n]``, with the
        bootstrap dropped when ``t + n`` runs past the episode.
    """
    rewards = np.asarray(rewards, np.float32)
    root_values = np.asarray(root_values, np.float32)
    horizon = rewards.shape[0]
    if root_values.shape[0] != horizon + 1:
        raise ValueError(
            f"expected {horizon + 1} root values for {horizon} rewards, "
            f"got {root_values.shape[0]}"
        )

    targets = np.zeros((horizon + 1,), np.float32)
    for t in range(horizon + 1):
        bootstrap_index = t + n_steps
        if bootstrap_index <= horizon:
            value = (discount**n_steps) * root_values[bootstrap_index]
        else:
            value = 0.0  # past the end of the episode: nothing left to bootstrap
        for k in range(min(n_steps, horizon - t)):
            value += (discount**k) * rewards[t + k]
        targets[t] = value
    return targets