TinyWAM

TinyWAM is a tiny DreamZero-style World Action Model reproduction for a 2D tabletop manipulation task. The installable Python package is tinydreamzero. It trains on scripted pixel demonstrations, predicts future visual latents and action chunks jointly with flow matching, and evaluates closed-loop on held-out layouts.

Setup

uv venv --python python3.14
uv sync --extra dev

Smoke Workflow

uv run --extra dev python -m tinydreamzero smoke --run-dir runs/smoke

Expected artifacts:

• runs/smoke/data/metadata.json
• runs/smoke/checkpoints/autoencoder.pt
• runs/smoke/checkpoints/wam.pt
• runs/smoke/metrics_eval_baselines.json
• runs/smoke/report.md

Default Workflow

uv run --extra dev python -m tinydreamzero generate --run-dir runs/default
uv run --extra dev python -m tinydreamzero train-ae --run-dir runs/default --data-root runs/default/data --device auto
uv run --extra dev python -m tinydreamzero train-wam --run-dir runs/default --data-root runs/default/data --device auto
uv run --extra dev python -m tinydreamzero eval --run-dir runs/default --policy baselines
uv run --extra dev python -m tinydreamzero eval --run-dir runs/default --policy wam --device auto
uv run --extra dev python -m tinydreamzero demo --run-dir runs/default

On an 8GB M1, start with the smoke workflow, then reduce --batch-size before increasing episode counts.

Dataset Playback

After generating a dataset, play episodes in real time with the Tkinter/Pillow viewer:

uv run --extra dev python scripts/play_dataset.py runs/dev/data --split train --fps 12 --scale 8

Use --split all to continue through train, validation, and test episodes in sorted order. Press space to pause, right arrow to skip to the next episode, and q or escape to quit.

WAM Rollout Playback

After training the autoencoder and WAM checkpoints, inspect closed-loop rollouts and the model's predicted future frames:

uv run --extra dev python scripts/play_wam_rollouts.py --run-dir runs/default --split test --episodes 32 --fps 8 --scale 5

The viewer shows the current rendered frame, autoencoder reconstruction, WAM-predicted future frames, and the actual frames produced by executing the predicted action chunk.

License

MIT License. See LICENSE.