--- license: apache-2.0 library_name: jax pipeline_tag: video-to-video tags: - robotics - world-model - dreamer - dreamer-v4 - yam - lerobot - jax - flax datasets: - Dimios45/yam-pick-duster-200 - Dimios45/yam-pick-duster - Dimios45/yam_towel_fold - Dimios45/molmo_eval_90cmtopcam - Dimios45/molmo_eval_hitl --- # 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](https://github.com/james0248/visionary) Dreamer-4 implementation, retargeted from SO-101 to YAM.

The same clip imagined under true, shuffled, and zero actions

## 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: | `--action_source` | PSNR | SSIM | rollout motion | vs ground truth | | --- | --- | --- | --- | --- | | `true` | **15.03** | **0.6675** | 3.684 | 1.6x | | `shuffled` (another episode's actions) | 14.55 | 0.6484 | 5.311 | **2.3x** | | `zero` | 17.85 | 0.7459 | 1.264 | 0.55x | | ground truth | — | — | 2.295 | 1.0x | * `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

Tokenizer reconstruction against ground truth

| stage | metric | value | | --- | --- | --- | | tokenizer | held-out PSNR | **25.6 dB** | | tokenizer | mse / lpips | 0.0009 / 0.0166 | | tokenizer | dead latent channels | **0 / 16** | | dynamics | rollout PSNR / SSIM | **24.82 / 0.913** | | dynamics | `flow_mse` (baseline 1.0) | 0.0048 | 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 ```python 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 | Corpus | Robot | Action | Episodes | Frames | | --- | --- | --- | --- | --- | | `yam-pick-duster` + `-200` | YAM | 7-D joint, **radians** | 250 | ~104k | | `yam_towel_fold` | bi-YAM | 14-D joint, radians | 11 | 9,468 | | `molmo_eval_90cmtopcam` + `_hitl` | molmoact dual-arm | 14-D joint, radians | 5 | 13,003 | 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`](https://huggingface.co/buckets/rerun/abc-130k) (a Rerun conversion of the gated [`XDOF/ABC-130k`](https://huggingface.co/datasets/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.

Combined tokenizer: YAM 25.6 dB versus ABC-130k 19.0 dB

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: `abc12/` — 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**. | | | | --- | --- | | `abc12/tokenizer` | 24k steps, **28.31 dB** held-out reconstruction | | `abc12/dynamics` | 16k steps, EMA rollout 19.10 dB | **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`](https://github.com/Dimios45/visionary/blob/main/docs/abc130k_world_model.md). **Use `abc12/` 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`](https://github.com/Dimios45/visionary/blob/main/docs/yam_world_model.md). ## Acknowledgements * [Dreamer 4](https://danijar.com/project/dreamer4) — the architecture. * [visionary](https://github.com/james0248/visionary) — the implementation this builds on. * [Dimios45](https://huggingface.co/Dimios45) — the YAM datasets. * [Rerun](https://huggingface.co/rerun) and the ABC-130k authors — the bimanual corpus.