Skip to content

Home

xwm
World Models

Action-conditioned latent world models in JAX. One dynamics interface, three training signals, four planners, and no decoder anywhere.

Overview

xwm is a JAX library for action-conditioned latent world models. A model is three parts: an encoder \(h: O \rightarrow Z\), a dynamics model \(d: Z \times A \rightarrow Z\), and whatever prediction heads the training signal requires. Every objective and every planner operates in \(Z\).

The library contains no decoder and no pixel-reconstruction loss. Encoders accept an arbitrary subset of the token grid, so masked positions cost nothing to compute, and planners reach the dynamics through a single \((z, a) \rightarrow z\) closure. That one closure is what lets a single set of planners serve every model in the library.

An action-conditioned latent world model An action-conditioned latent world model

The observation is a real frame from the Franka FR3 environment in xwm.envs, at the resolution the encoder is given. Everything after the encoder happens inside Z, which is why one dynamics closure serves every family and every planner in the library.

Install

pip install xwm                 # core
pip install "xwm[dev]"          # core plus tests
pip install "xwm[newton]"       # core plus the Franka robot environment

Python 3.11 or newer, with jax, equinox, optax and einops. The optional extras are listed in Installation.

Quick start

import xwm

xwm.set_seed(0)

# Self-supervised: learns from observation alone, no reward.
model = xwm.families.jepa.lejepa(size="small", img_size=224, patch_size=16)

# Reward-driven, continuous actions: the natural fit for a robot arm.
agent = xwm.families.tdmpc2.tdmpc2(action_dim=7, observation="state", state_dim=20)

# Reward-driven, discrete actions, plans with tree search.
agent = xwm.families.muzero.muzero(n_actions=15, observation="state", state_dim=20)

trainer = xwm.training.Trainer(model, xwm.training.adamw(1e-4))
state, history = trainer.fit(batches, steps=10_000)

The Quickstart takes this from an untrained model to a working planner in about twenty lines.

Models

Three families share the same encoders, latent dynamics and planners. What separates them is the signal that trains the latent space.

family learning signal reward? actions planner
jepa its own future embeddings no continuous CEM / MPPI
tdmpc2 reward and TD value yes continuous MPPI
muzero search-improved targets yes discrete MCTS

They are complementary rather than competing. JEPA needs no reward, so it can pretrain on passive video, which is abundant and unlabelled. TD-MPC2 and MuZero need interaction, but they learn a value function, so their planner can see past its own horizon. A JEPA encoder is a reasonable initialisation for either, and it is one argument: tdmpc2(encoder=pretrained).

Documentation

Library layout

module contents
xwm.core types, base modules, EMA targets, rollouts, the default key
xwm.nn attention, transformers, RoPE, patch and tubelet embeddings, SimNorm
xwm.encoders observation to latent: image, video, state
xwm.dynamics (z, a) → z', transformer or MLP
xwm.heads reward, value, policy, Q-ensemble, categorical scalars
xwm.masking what a JEPA predicts: blocks, tubes, temporal splits
xwm.objectives latent prediction, SIGReg, VICReg, InfoNCE
xwm.families jepa, tdmpc2, muzero, and a registry
xwm.planning CEM, MPPI, gradient planning, MPC, MCTS, latent costs
xwm.training Trainer, schedules, TrainState, ReplayBuffer
xwm.envs a Franka FR3 arm in Newton
xwm.data batch streams and a synthetic controllable world
xwm.metrics probes and collapse diagnostics
xwm.plots figures, GIFs, JSON and LaTeX tables
xwm.tools checkpointing, model summaries

xwm.dynamics is the centre of the library rather than an add-on. Every family consumes a \((z, a) \rightarrow z\) model from it, and every planner consumes nothing else. Changing family changes how that model is trained, never how it is used.

Project