quickdraw-wm-torus-vl128

isaac-ronald-ward/quickdraw-wm-torus-vl128

A latent world model: it takes 8 steps of context (proprioceptive vector + camera frame) plus a sequence of actions, and rolls forward open-loop — predicting the frames and states that follow with no further observations.

Trained with quickdraw (model=mm_flow, recipe vl128) on isaac-ronald-ward/torus-world. Image head: image at 128px, 32 latent tokens each.

What it looks like

<video controls loop muted playsinline width="100%" src="https://huggingface.co/isaac-ronald-ward/quickdraw-wm-torus-vl128/resolve/main/rollout.mp4"></video>

Open-loop rollout, full horizon at frame rate — predicted beside ground truth. If your viewer does not play it inline: rollout.mp4.

Top row predicted, bottom row ground truth, over the full +1..+128 open-loop horizon on a held-out validation episode. The run's own logged artifact at epoch 17.

Also: a second episode filmstrip_2.png, error against horizon error_vs_step.png.

Headline numbers

All from epoch 17, the published checkpoint — not each metric's own best epoch.

metrics.json carries every logged metric across the whole run, each with the eval index it came from. Lower is better for lpips, higher for psnr. @+128 means 128 prediction steps with no re-grounding.

Checkpoint selection: explicitly requested via +hub.ckpt

Version this was trained with

pip install "git+https://github.com/isaac-ward/quickdraw@e990f47"

quickdraw 0.1.0 · torch 2.10.0+cu128. The package version is not bumped per change, so the commit is the real pin.

Quickstart

from quickdraw import load_pretrained, load_example_context
import torch

model, norm, cfg = load_pretrained("isaac-ronald-ward/quickdraw-wm-torus-vl128", device="cuda")
ex = load_example_context("isaac-ronald-ward/quickdraw-wm-torus-vl128")          # real context windows, shipped with the model

P, H = int(cfg.data.P), 64
head = "image"

ctx = {"proprio": norm.norm_obs(torch.from_numpy(ex["obs"][:, :P])).float().cuda(),
       head:      torch.from_numpy(ex[f"frames__{head}"][:, :P]).float().div(255).cuda()}
acts = norm.norm_act(torch.from_numpy(ex["act"][:, :P + H - 1])).float().cuda()

out = model.imagine_eval(ctx, acts, horizon=H, decode_chunk=16)
frames  = out[head].clamp(0, 1)              # (B, H, 128, 128, 3) imagined frames
proprio = norm.denorm_obs(out["proprio"])    # (B, H, 6) in physical units

Full walkthrough: docs/using_pretrained_models.md and its companion notebook.

Common pitfalls

Vectors are normalised, images are not. proprio in through norm_obs, out through denorm_obs; frames are plain [0, 1] floats. Getting this backwards produces plausible garbage, not an error.

acts needs P + H - 1 steps, not H — the context steps consume actions too.

Pass decode_chunk. The image decoder is ~78% of per-sample memory; without a chunk a long horizon will OOM. 16 is a safe default.

N missing keys on load is expected when the names contain visual._net — the frozen LPIPS loss network, excluded from the published weights and rebuilt on demand.

What ships here