yam-world-model

YAM World Model — a Dreamer-4 style action-conditioned video world model

An action-conditioned latent video world model for the YAM robot arm (single-arm 7-DoF and bimanual 14-DoF), trained from scratch on a single RTX 4090. Given a few frames of context and a stream of joint commands, it imagines the future.

Built on the visionary Dreamer-4 implementation, retargeted from SO-101 to YAM.

<p align="center"> <img src="assets/yam_action_control.webp" alt="The same clip imagined under true, shuffled, and zero actions" width="800"> </p>

The model actually obeys its actions

This is the property that matters and the one that is easy to fake. A world model trained on a single repetitive scene learns to predict the future from pixels alone and quietly ignores the action input. The control is to re-run the same clip, same seed, and change only the actions:

• true beats shuffled on both metrics.
• Wrong actions produce wrong motion — the arm thrashes at 2.3x ground truth.
• Zero actions nearly freeze the arm, which is correct for "no commanded movement".

⚠️ zero scores highest on PSNR, and that is a metric artifact. On a mostly static scene a near-frozen prediction beats sharp-but-slightly-misaligned motion, because pixel metrics reward blur. PSNR alone would have scored this model as ignoring its actions. Pair the action control with a motion statistic.

Reconstruction

<p align="center"> <img src="assets/yam_reconstruction.webp" alt="Tokenizer reconstruction against ground truth" width="800"> </p>

Rollout quality saturates at the tokenizer's own reconstruction ceiling — past that, the dynamics model is limited by what the tokenizer can decode.

Architecture

Two stages, both JAX/Flax, following Dreamer 4.

Video tokenizer — masked-autoencoder spatiotemporal transformer. 144×192 frames, patch 16, → 96 latent tokens × 16 channels per frame, tanh bottleneck. Video only; never sees actions. 8 encoder + 8 decoder layers, dim 512. MSE + LPIPS, Muon optimizer, 12,000 steps.

Dynamics — shortcut/flow-matching transformer over the frozen latents, conditioned on a continuous action vector through a per-embodiment MLP projection plus an embodiment embedding. 8 layers, dim 768, 24-frame context (4.8 s at 5 Hz). 6,000 steps, EMA 0.999 (exports carry EMA weights).

Two embodiments share one video backbone: yam_absolute_joint_7d_v1 (7-D) and bi_yam_absolute_joint_14d_v1 (14-D), max_action_dim: 14.

Usage

from predict import WorldModel   # scripts/robot/predict.py

wm = WorldModel(
    tokenizer_dir="yam_tokenizer",
    dynamics_dir="yam_dynamics",
    latent_stats="latent_stats.json",
    dynamics_step=6000,
)
# context_frames: (T0, H, W, 3) uint8   actions: (T0+horizon, 7) normalized to [-1,1]
frames = wm.imagine(context_frames, actions, embodiment_id=0)

5.4 fps warm on an RTX 4090 (first call ~28 s of JIT compile; 4.4 s per 24-frame rollout thereafter). Trained at 5 Hz, so it imagines at roughly the rate the data was recorded — enough for single-trajectory lookahead in a 5 Hz control loop, not for dense multi-rollout MPC.

Training data

266 packed records / 118,982 frames, 0 invalid. All are LeRobot v3.0. Only verified-fixed cameras are kept; wrist cameras are dropped. Episodes are trimmed to their non-idle span, actions normalized q01–q99 → [-1,1], video packed at 240×320 and decoded to 144×192 at 5 Hz.

Limitations

• Not a policy. This is a world model — it imagines, it does not act. No planning or RL is included.
• Never drives a real arm. Model output has not been executed on hardware. Check your robot's calibration convention before attempting it.
• Single fixed viewpoint per episode, and the scene diversity is modest (duster picking, towel folding, two eval sets). Expect degradation on unseen scenes, lighting, or camera placements.
• Bimanual is the weaker embodiment — 16 episodes against 250 for the single-arm side.
• Rollouts drift over long horizons, as flow-matching video models do. Results here are for a 32-frame (6.4 s) horizon from 4 context frames.

A negative result worth publishing

Adding 39 episodes from rerun/abc-130k (a Rerun conversion of the gated XDOF/ABC-130k) more than doubled the corpus to 260,731 frames and turned the starved bimanual embodiment into the larger one — and made the model worse. The single-arm side lost its action-conditioning (shuffled began outscoring true) and the bimanual side never gained any.

<p align="center"> <img src="assets/abc_domain_gap.webp" alt="Combined tokenizer: YAM 25.6 dB versus ABC-130k 19.0 dB" width="800"> </p>

Two unseparated causes: the two bimanual robots were merged into one embodiment (same action convention, different kinematics — identical joint angles produce different pixel motion on a different robot), and the combined tokenizer reconstructs ABC at only 19.0 dB against YAM's 25.6, a gap that widened with training as ABC's dense two-arm scenes saturated the 96-latent bottleneck.

The weights published here are the single-corpus YAM model, which passes the action control.

Also in this repo: abc130k/ — an ABC-130k bimanual model

A second pair of checkpoints trained from scratch on ABC-130k (240 episodes across 48 tasks, streamed from a 33.75 TB Rerun bucket without bulk downloading). Larger tokenizer: 192 latents x 24 channels.

Doubling bottleneck width lifted reconstruction 22.90 -> 28.31 dB on identical eval clips — a decisive answer to whether ABC's dense two-arm scenes were capacity-limited. They were.

It did not produce action-conditioning. Feeding another episode's actions yields motion closer to ground truth than the true actions (2.34x vs 3.12x of ground-truth motion), so the model is not following its commands. Rollouts sit at 17.02 dB against a 28.31 dB tokenizer ceiling — the dynamics model, not the tokenizer, is the binding constraint. Full analysis in docs/abc130k_world_model.md.

Use abc130k/ for bimanual reconstruction; use the YAM weights above if you need a world model that responds to actions.

Reproducing

Full recipe, per-stage commands, and the failure modes encountered are documented in docs/yam_world_model.md.

Acknowledgements

• Dreamer 4 — the architecture.
• visionary — the implementation this builds on.
• Dimios45 — the YAM datasets.
• Rerun and the ABC-130k authors — the bimanual corpus.