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

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

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

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

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.

Acknowledgements

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