Action-conditioned world models for robotics.

---

xwm is a JAX-based library for action-conditioned latent world models: 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, which is what lets one set of planners serve every model.

The components are independently useful: ViT encoders over 2D patches or 3D tubelets and MLP encoders over state vectors; transformer or residual-MLP dynamics; categorical reward and value heads, pessimistic Q-ensembles, squashed-Gaussian policies; latent-prediction, SIGReg, VICReg, InfoNCE and TD objectives; CEM, MPPI, gradient and PUCT-MCTS planners; a family-agnostic trainer with EMA targets, parameter freezing and a trajectory replay buffer.

Install

Add it to your own project:

uv add xwm                     # core
uv add "xwm[plots]"            # figures, GIFs, tables
uv add "xwm[newton]"           # the Franka robot environment
uv add "xwm[data]"             # recorded datasets: DROID, LIBERO, OGBench, OXE

Or work on it from a clone, where uv.lock pins the whole environment:

git clone https://github.com/kamara-lab/xwm && cd xwm
uv sync --extra dev                       # core + tests
uv sync --extra dev --extra newton        # + the Franka robot environment
uv run pytest                             # run in that environment

--all-extras is the one combination to avoid: it pulls in render, whose ovrtx ships as an sdist and wants a graphics-capable NVIDIA GPU, so it will try to build on machines that can never use it. Add --extra render deliberately, on a host that has one.

Python ≥ 3.11, jax, equinox, optax, einops.

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)

Models

All three share the same encoders, latent dynamics and planners. What separates them is what signal trains the latent space.

They are complementary rather than competing. JEPA needs no reward, so it can pretrain on passive video, 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: tdmpc2(encoder=pretrained) is one argument.

xwm.families.available() lists every registered model; xwm.families.create(name, **kwargs) builds one by name.

Layout

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.

Concepts

What a JEPA predicts

A mask sampler splits the token grid into a visible context and target blocks. The context encoder computes only the visible tokens, which is where the speedup over reconstruction comes from.

Masks are batch-shared and statically shaped, so a training step compiles once. Sampling is combinatorial host-side work and happens in model.prepare_batch(), outside jit; the Trainer calls it for you.

Why it doesn't collapse

Predicting a representation from a representation has a trivial solution: emit a constant. collapse= selects the countermeasure.

SIGReg replaces EMA teachers, stop-gradients, centering and sharpening with one statement: the embedding distribution should be an isotropic Gaussian. It is enforced by a sketch — for z ~ N(0, I_D) and any unit vector v, the projection ⟨z, v⟩ is exactly N(0, 1) regardless of D — so it draws random directions, projects the batch onto each, and penalises deviation from a standard normal. Isotropy and unit scale both fall out, the cost is linear in batch size, and there is one coefficient instead of a schedule.

xwm.objectives.sigreg(z, key, n_proj=256, statistic="epps_pulley")

Planning

model.dynamics_fn() hands a planner a plain (z, a) -> z' closure. Everything in xwm.planning is jittable — candidates are vmaped and refinement is a lax.fori_loop — so a plan is one device call.

planner = xwm.planning.CEM(horizon=8, action_dim=7, n_samples=512, n_elites=64)
cost = xwm.planning.goal_cost(model.encode(goal_image), kind="l2")
plan = planner.plan(key, model.dynamics_fn(), model.encode(observation), cost)

run_mpc closes the loop with replanning and warm starts. For value-based agents, return_cost scores candidates by predicted reward plus a terminal value bootstrap — the term that lets a horizon-3 planner act as though it saw further.

Diagnostics

The loss is not the metric. A collapsing encoder drives its prediction loss down — it is predicting its own degenerate output.

xwm.metrics.collapse_report(z)
# {'rankme': ..., 'rank_ratio': ..., 'feature_std': ..., 'mean_cosine': ...}

feature_std → 0 and mean_cosine → 1 both mean collapse; rankme is the effective rank of the spectrum. All are reported because each misses a case the others catch — rankme is computed after centring, so a constant offset is invisible to it. A linear probe is not a collapse detector: ridge_probe standardises features, so it amplifies a nearly-dead signal back to full scale.

Robotics

xwm.envs wraps a Franka Emika FR3 in Newton (NVIDIA Warp), observed either as pixels or as a 20-D proprioceptive state vector. A dense reach task supplies the reward the value-based families need.

env = xwm.envs.FrankaEnv(xwm.envs.FrankaConfig(image_size=64))

data = xwm.envs.franka_sequences(env, 320, 8, seed=0)   # for JEPA
env.state_observation(), env.reward(action), env.goal_distance()   # for RL

franka_sequences returns exactly what xwm.data.sprite_sequences does, so it drops straight into any family.

Two synthetic worlds cover the same ground on CPU, without a simulator. xwm.data.PushWorld is planar pushing: the puck moves only when the pusher touches it, which puts a hinge in the dynamics that SpriteWorld's linear ones cannot have, and it carries a dense reward. xwm.data.MazeWorld is sparse-reward navigation through corridors, shaped after OGBench's antmaze navigate tasks. Both were tuned by measuring the task rather than by eye: on PushWorld, a single sequence of random actions solves 3% of starts and the best of 80 solves 94%, which is the gap a planner has to exploit.

Recorded datasets

xwm.datasets reads the corpora the field benchmarks on, in the four formats they come in, and returns exactly the field layout sprite_sequences does.

xwm.datasets.describe("lerobot/droid-100")               # size, episodes, licence, citation
data = xwm.datasets.create("libero/10", length=16, limit=20, resize=64)
buffer = xwm.datasets.to_replay_buffer(xwm.datasets.create("robomimic/lift", length=32))

Mixtures are what OXE is actually for, so xwm.datasets.mixture(xwm.datasets.OXE_MIXTURE, action_dim=7, resize=64) interleaves several members by weight, zero-padding narrower action spaces to a common width; xwm.datasets.stream(name) is the same registry access without materialising anything.

Every registered source was resolved against the live host, and the episode counts come from the datasets' own metadata rather than from their papers. The one thing to know before using it: recorded formats store one action per frame, including a last one whose result was never recorded, and xwm needs T - 1 actions for T frames. The readers drop it and the contract check refuses an episode that did not — a silent one-step shift looks exactly like underfitting.

Rendering

Training and figures want opposite things from a renderer, so there are two paths. Use xwm.envs.which_backends() to see what is installed.

env.observe()                        # warp, at config.image_size -- for training
env.render(384, samples=3)           # warp, supersampled -- for a clean figure

with env.high_quality_renderer(backend="rtx", size=(768, 768)) as r:
    r.add(env.state)                 # path traced, one frame per state
frames = r.frames                    # (T, 3, H, W) -- feeds save_gif directly

with env.high_quality_renderer(backend="usd", output_path="ep.usd") as r:
    r.add(env.state)                 # a stage to render in Omniverse or Blender

env.render casts one ray per pixel, so samples renders at samples× and averages down — the only anti-aliasing the Warp raytracer has. It is the right tool for observations and for tidy figures, but it will not produce a photorealistic image: for that use rtx, or export a USD stage and render it offline. Every backend shares one camera definition (env.camera_framing), so the path-traced figure and the observations the model trains on show the same view from the same place.

Because the physics is deterministic given a seed and an action sequence, a path-traced figure is produced by replaying an episode rather than by storing its pixels — the render is of the same episode the numbers came from. Example 06 writes both: episode_frames.png is what the encoder sees, episode_rtx.gif and planning_episode_rtx.gif are what the robot is doing. Set XWM_RTX=0 to skip them, or XWM_RTX_SIZE to change the resolution. deploy/app_render.py runs every available backend on a GPU and writes the results side by side; app_render.py::vulkan_probe reports in about a minute whether OVRTX can get a device at all.

rtx needs more than an NVIDIA GPU: it needs graphics access. Many GPU cloud containers — Modal's among them — expose a compute-only device set (no /dev/nvidia-modeset), which satisfies CUDA but not NVIDIA's Vulkan driver, so OVRTX cannot create an instance there however complete the library stack is. Example 06 therefore picks its renderer from which_backends() at run time, and where OVRTX is unavailable it writes supersampled Warp figures plus episode.usd to path trace offline. See docs/findings.md for the diagnosis.

The two-stage V-JEPA 2-AC recipe — learn a representation from passive video, freeze it, learn action-conditioned dynamics in its latent space — is one call. Freezing is not only a compute saving: with the encoder fixed the prediction targets are fixed functions of the observations, so the dynamics model has nothing to gain from degrading the representation.

model = xwm.families.jepa.action_world_model(
    action_dim=7, encoder=pretrained_encoder, freeze_encoder=True,
)
trainer.n_trainable == model.dynamics.n_params   # the encoder gets no optimizer state

Training mixes teacher forcing (one step from ground-truth latents — a dense signal) with rollout (the full horizon from a single latent, the model consuming its own predictions — the only term that penalises compounding error).

Training

Trainer is family-agnostic: it needs only loss, prepare_batch and trainable. It owns the jit boundary, the EMA teacher, and the parameter filter, so frozen submodules never reach the optimizer.

trainer = xwm.training.Trainer(model, xwm.training.adamw(xwm.training.cosine_warmup(1e-3, 1000)))
state, history = trainer.fit(batches, steps=1000)

Batches come from xwm.data.iter_batches for a fixed dataset, or from xwm.training.ReplayBuffer for the reward-driven families, whose losses need contiguous slices of a single episode. The buffer rejects slices that straddle an episode boundary — training a dynamics model to predict through a reset is the one transition it can never get right. xwm.datasets.to_replay_buffer fills it from a recorded corpus, one clip per episode.

Keys

key= is optional wherever a model is built. Omit it and the key comes from an ambient source; pass one and nothing ambient is touched.

xwm.set_seed(0)
model = xwm.families.jepa.ijepa(img_size=64)                    # ambient
other = xwm.families.jepa.ijepa(img_size=64, key=jr.PRNGKey(7)) # explicit

with xwm.seed(123):                                             # scoped
    model = xwm.families.jepa.ijepa(img_size=64)

The source advances on every draw — it has to, or every transformer block would be initialised identically — so a fixed sequence of calls under a fixed seed is reproducible, but inserting a construction shifts everything built after it. Pass explicit keys for anything that must survive refactors.

Only construction defaults. loss, sigreg and the planners still require a key, because those are consumed inside jit, where a key drawn at trace time would be baked in as a constant and reused for every step.

Examples

Examples 01–05 run on CPU against the synthetic worlds in xwm.data, so there is no dataset to download. 06–08 need the newton extra and download the Franka asset on first run. 09 needs the data extra and downloads ~30 MB of recorded robot data.

Measured results, including the negative ones, are collected in docs/findings.md.

Running on GPU

Each experiment is its own Modal app, so the nine can run concurrently on separate GPUs and be started, watched and stopped independently.

./deploy/run_all.sh                            # all nine, gpu preset
modal run deploy/app_tdmpc2.py --preset xl     # one, at higher fidelity

Presets (cpu-parity, gpu, xl) raise resolution, episode count, model size and render quality through XWM_* environment variables, so there is one copy of each pipeline rather than a laptop version and a cluster version. cpu-parity exists to isolate hardware from settings when comparing runs.

Conventions

• Modules are unbatched. Written for a single sample and vmaped by the caller, the Equinox idiom. Batch-level entry points are the methods named loss.
• Shapes. Images (C, H, W), clips (T, C, H, W), token sequences (N, D), flat latents (D,), actions (A,). Masks are int32 index arrays.
• Immutability. model.eval_mode() returns a dropout-free copy.

Tests

uv run pytest

The tests are written to fail on broken behaviour, not just broken shapes: mask samplers must never leak a target token into the context, dynamics must respond to their action input, planners must reach a reachable goal, MCTS must find a payoff one step away, frozen parameters must not move, SIGReg must actually pull a skewed distribution toward isotropy, and a dataset reader must drop the final recorded action rather than shift a whole corpus by one step.

Two groups are gated. The Franka tests need the newton extra and skip without it. The dataset tests that download run only under XWM_DATASET_TESTS=1; the rest of them exercise every reader against fixtures written into tmp_path in the real on-disk formats — a real parquet shard, a real mp4, a real HDF5 — because a reader tested against a mock of a format is a reader tested against a belief about the format.

References

Every module carries a References block in its docstring naming the paper the code follows, so the citation sits beside the implementation — try help(xwm.families.tdmpc2.model).

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

Simulation: Newton with a Franka Emika FR3; MuJoCo via mujoco_warp where a CUDA GPU is available, Featherstone otherwise.

Contributors

Supported by

Get in touch kleyton.vsc@gmail.com

License

Apache-2.0