README.md

---
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.

<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:

| `--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

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

| 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.

<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.**

## 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.