xwm.encoders

Observation → latent. ViT encoders over 2-D patches or 3-D tubelets, and an MLP encoder over state vectors. All of them accept an arbitrary subset of the token grid.

Observation encoders: pixels or state vectors to latent tokens.

Every family in xwm.families takes an encoder from here. Which one you pick is a question about the observation, not about the algorithm.

Modules:

Name	Description
image	Image encoders.
presets	Standard ViT and predictor widths, shared across families.
state	Encoder for low-dimensional state observations.
video	Video encoders.
vision	The vision encoder every family shares: tokenise, add position, transform.

Classes:

Name	Description
ImageEncoder	A ViT over 2-D patches, maskable via keep.
StateEncoder	Embed a state vector to (1, D) tokens.
VideoEncoder	A ViT over space-time tubelets, maskable via keep.
VisionEncoder	Tokenise, add position, run a transformer. Returns (N, D) tokens.

Functions:

Name	Description
image_encoder	Build an ImageEncoder from a size preset.
preset	Look up a size preset by name ("tiny" ... "huge").
video_encoder	Build a VideoEncoder from a size preset.
make_pos	Build the additive table and/or the rotary embedding for kind.

ImageEncoder

ImageEncoder(*, key: PRNGKey | None = None, img_size: int | tuple[int, int] = 224, patch_size: int = 16, in_channels: int = 3, embed_dim: int = 384, depth: int = 12, num_heads: int = 6, pos: PosKind = 'sincos', **kwargs)

Bases: VisionEncoder

A ViT over 2-D patches, maskable via keep.

Parameters:

Name	Type	Description	Default
img_size	int | tuple[int, int]	input resolution, int or (H, W).	224
patch_size	int	side length of a square patch.	16
in_channels	int	input channels.	3
embed_dim	int	token width.	384
depth	int	transformer blocks.	12
num_heads	int	attention heads.	6
pos	PosKind	positional scheme -- "sincos", "learned", "rope".	'sincos'

Source code in xwm/encoders/image.py

def __init__(
    self,
    *,
    key: PRNGKey | None = None,
    img_size: int | tuple[int, int] = 224,
    patch_size: int = 16,
    in_channels: int = 3,
    embed_dim: int = 384,
    depth: int = 12,
    num_heads: int = 6,
    pos: PosKind = "sincos",
    **kwargs,
):
    key = resolve_key(key)
    import jax.random as jr

    k_patch, k_rest = jr.split(key)
    patch_embed = PatchEmbed2d(img_size, patch_size, in_channels, embed_dim, key=k_patch)
    super().__init__(
        patch_embed, depth=depth, num_heads=num_heads, key=k_rest, pos=pos, **kwargs
    )

StateEncoder

StateEncoder(state_dim: int, embed_dim: int = 256, *, key: PRNGKey | None = None, hidden_dim: int | None = None, depth: int = 2, normalize: bool = True)

Bases: Module

Embed a state vector to (1, D) tokens.

Parameters:

Name	Type	Description	Default
state_dim	int	length of the observation vector.	required
embed_dim	int	token width.	256
hidden_dim	int | None	MLP width; defaults to 4 * embed_dim.	None
depth	int	number of MLP blocks.	2
normalize	bool	LayerNorm the input, which matters when the components have wildly different units (radians next to metres next to velocities).	True

Source code in xwm/encoders/state.py

def __init__(
    self,
    state_dim: int,
    embed_dim: int = 256,
    *,
    key: PRNGKey | None = None,
    hidden_dim: int | None = None,
    depth: int = 2,
    normalize: bool = True,
):
    key = resolve_key(key)
    keys = jr.split(key, depth)
    hidden_dim = hidden_dim or 4 * embed_dim
    dims = [state_dim] + [embed_dim] * depth
    self.layers = [
        Mlp(dims[i], hidden_dim, dims[i + 1], key=keys[i]) for i in range(depth)
    ]
    self.input_norm = LayerNorm(state_dim) if normalize else None
    self.out_norm = LayerNorm(embed_dim)
    self.state_dim = state_dim
    self.embed_dim = embed_dim

VideoEncoder

VideoEncoder(*, key: PRNGKey | None = None, img_size: int | tuple[int, int] = 224, patch_size: int = 16, num_frames: int = 16, tubelet_size: int = 2, in_channels: int = 3, embed_dim: int = 384, depth: int = 12, num_heads: int = 6, pos: PosKind = 'sincos', **kwargs)

Bases: VisionEncoder

A ViT over space-time tubelets, maskable via keep.

Tokenising time jointly with space (tubelet_size > 1) is what makes the token count tractable for clips, and it is what makes tube masking a non-trivial prediction problem rather than a frame-copy.

Parameters:

Name	Type	Description	Default
img_size	int | tuple[int, int]	spatial resolution, int or (H, W).	224
num_frames	int	clip length in frames.	16
tubelet_size	int	frames per token; num_frames must be divisible by it.	2

Source code in xwm/encoders/video.py

def __init__(
    self,
    *,
    key: PRNGKey | None = None,
    img_size: int | tuple[int, int] = 224,
    patch_size: int = 16,
    num_frames: int = 16,
    tubelet_size: int = 2,
    in_channels: int = 3,
    embed_dim: int = 384,
    depth: int = 12,
    num_heads: int = 6,
    pos: PosKind = "sincos",
    **kwargs,
):
    key = resolve_key(key)
    import jax.random as jr

    k_patch, k_rest = jr.split(key)
    patch_embed = PatchEmbed3d(
        img_size, patch_size, num_frames, tubelet_size, in_channels, embed_dim, key=k_patch
    )
    super().__init__(
        patch_embed, depth=depth, num_heads=num_heads, key=k_rest, pos=pos, **kwargs
    )

VisionEncoder

VisionEncoder(patch_embed: PatchEmbed, *, depth: int, num_heads: int, key: PRNGKey | None = None, pos: PosKind = 'sincos', mlp_ratio: float = 4.0, qk_norm: bool = False, dropout: float = 0.0, drop_path: float = 0.0, layer_scale: float | None = None, remat: bool = False)

Bases: Module

Tokenise, add position, run a transformer. Returns (N, D) tokens.

The one feature that separates this from a stock ViT is keep: the encoder can be run on an arbitrary subset of the token grid. A JEPA context encoder sees only visible tokens, so masked positions cost nothing to compute -- which is where most of the training speedup over pixel reconstruction comes from.

There is deliberately no [CLS] token. Predictive world models are trained with no global objective to attach one to; pool the tokens instead (xwm.nn.mean_pool or xwm.nn.AttentivePooler).

Source code in xwm/encoders/vision.py

def __init__(
    self,
    patch_embed: PatchEmbed,
    *,
    depth: int,
    num_heads: int,
    key: PRNGKey | None = None,
    pos: PosKind = "sincos",
    mlp_ratio: float = 4.0,
    qk_norm: bool = False,
    dropout: float = 0.0,
    drop_path: float = 0.0,
    layer_scale: float | None = None,
    remat: bool = False,
):
    key = resolve_key(key)
    k_pos, k_blocks = jr.split(key)
    dim = patch_embed.embed_dim
    grid = tuple(patch_embed.grid)
    self.patch_embed = patch_embed
    self.pos_embed, self.rope = make_pos(pos, grid, dim, num_heads, key=k_pos)
    self.blocks = Transformer(
        dim,
        depth,
        num_heads,
        key=k_blocks,
        mlp_ratio=mlp_ratio,
        qk_norm=qk_norm,
        dropout=dropout,
        drop_path=drop_path,
        layer_scale=layer_scale,
        rope=self.rope,
        remat=remat,
    )
    self.embed_dim = dim
    self.grid = grid

image_encoder

image_encoder(size: str = 'small', **kwargs) -> ImageEncoder

Build an ImageEncoder from a size preset.

>>> enc = image_encoder("base", img_size=224, patch_size=16, key=key)

Source code in xwm/encoders/image.py

def image_encoder(size: str = "small", **kwargs) -> ImageEncoder:
    """Build an :class:`ImageEncoder` from a size preset.

    >>> enc = image_encoder("base", img_size=224, patch_size=16, key=key)
    """
    return ImageEncoder(**{**preset(size), **kwargs})

preset

preset(name: str, kind: str = 'encoder') -> dict[str, int]

Look up a size preset by name ("tiny" ... "huge").

Source code in xwm/encoders/presets.py

def preset(name: str, kind: str = "encoder") -> dict[str, int]:
    """Look up a size preset by name (``"tiny"`` ... ``"huge"``)."""
    table = VIT_PRESETS if kind == "encoder" else PREDICTOR_PRESETS
    if name not in table:
        raise KeyError(f"unknown {kind} preset {name!r}; choose from {sorted(table)}")
    return dict(table[name])

video_encoder

video_encoder(size: str = 'small', **kwargs) -> VideoEncoder

Build a VideoEncoder from a size preset.

Source code in xwm/encoders/video.py

def video_encoder(size: str = "small", **kwargs) -> VideoEncoder:
    """Build a :class:`VideoEncoder` from a size preset."""
    return VideoEncoder(**{**preset(size), **kwargs})

make_pos

make_pos(kind: PosKind, grid: tuple[int, ...], dim: int, num_heads: int, *, key: PRNGKey | None = None) -> tuple[SinCosPosEmbed | LearnedPosEmbed | None, AxialRoPE | None]

Build the additive table and/or the rotary embedding for kind.

Source code in xwm/encoders/vision.py

def make_pos(
    kind: PosKind,
    grid: tuple[int, ...],
    dim: int,
    num_heads: int,
    *,
    key: PRNGKey | None = None,
) -> tuple[SinCosPosEmbed | LearnedPosEmbed | None, AxialRoPE | None]:
    """Build the additive table and/or the rotary embedding for ``kind``."""
    n_tokens = int(jnp.prod(jnp.asarray(grid)))
    if kind == "sincos":
        return SinCosPosEmbed(grid, dim), None
    if kind == "learned":
        return LearnedPosEmbed(n_tokens, dim, key=key), None
    if kind == "rope":
        return None, AxialRoPE(grid, dim // num_heads)
    if kind == "none":
        return None, None
    raise ValueError(f"unknown positional kind {kind!r}")