wms — World Models on Minecraft (VPT + DreamerV3)

Training a DreamerV3 world model (action-conditioned video prediction) on OpenAI's VPT Minecraft contractor data, with a full data pipeline, a sweep harness, and kernel experiments. World-model only: VPT data is offline (image, action) with no rewards/env, so this learns to predict future frames given past frames + actions — not an agent policy.

Headline result: best config dyn_deter=512, lr=6e-4; an 80k-step run reaches image reconstruction loss 16.15 at 64×64. Full analysis in FINDINGS.md; the build journey in DEVLOG.md; architecture in SPEC.md.

Heads-up: the data is large. The full VPT contractor set is ~11.3 TB (measured). You do not need all of it — a few hundred GB subset trains a solid world model. The downloader estimates size first and refuses to fill your disk.

What's here

src/minecraft_wm/
  vpt_actions.py   # VPT keyboard/mouse -> 25-D action vector (minerl-free)
  download.py      # selective, resumable, disk-safe VPT downloader (HEAD-estimates first)
  preprocess.py    # mp4 + jsonl -> uncompressed mmappable .npy episodes (parallel)
  dataset.py       # random length-L windows over episodes (true mmap)
  config_util.py   # build a DreamerV3 config from vendored configs.yaml + --set overrides
  spaces.py        # gym-free obs/act space shims for the WorldModel
  wm_train.py      # offline world-model trainer (TB logs, ckpts, video preds)
  sweep.py         # sequential hyperparameter sweep + ranked summary
  kernels/gru_fused.py  # Triton fused GRU-cell epilogue (prototype, forward-only)
scripts/
  setup.sh         # clone vendored baselines + create the venv
  validate_slice.py    # vertical-slice sanity: data -> WorldModel -> train step
  campaign.py      # multi-stage overnight sweep campaign (self-budgeting)
  profile_wm.py    # per-component + launch-bound profiler
  bench_gru.py / bench_batch.py  # kernel + batch-size benchmarks
configs/           # sweep configs
data/indexes/      # the 9 official VPT index manifests (relpath lists; the download inputs)

Vendored baselines are not committed (clone via setup.sh): NM512/dreamerv3-torch (the WorldModel) and openai/Video-Pre-Training (data format).

Requirements

• Linux + NVIDIA GPU (developed on an RTX PRO 6000 96 GB; B16/64×64 needs only ~5 GB).
• Python 3.12, a CUDA build of PyTorch 2.x, plus triton, opencv-python, and ffmpeg.
• uv for env management.

Reproduce

# 1. clone + set up (clones vendored repos, builds .venv, installs deps)
git clone https://github.com/infatoshi/wms && cd wms
bash scripts/setup.sh && . .venv/bin/activate

# 2. estimate, then download a subset (NOT the full 11 TB). Refuses below --min-free-gb.
python -m minecraft_wm.download --index data/indexes/all_10xx_Jun_29.json \
    --count 2000 --estimate-only
python -m minecraft_wm.download --index data/indexes/all_10xx_Jun_29.json \
    --count 2000 --min-free-gb 50 --workers 12 --out data/raw

# 3. preprocess mp4+jsonl -> uncompressed .npy episodes (parallel)
python -m minecraft_wm.preprocess --jobs 14 --raw-dir data/raw --out-dir data/episodes

# 4. sanity-check the full path on real data (one WorldModel train step)
python scripts/validate_slice.py --steps 5

# 5. train the best config
python -m minecraft_wm.wm_train --name run0 --steps 50000 --precision 16 \
    --set dyn_deter=512 --set model_lr=6e-4 --set batch_size=16 --set batch_length=64

# 6. (optional) sweep, or the overnight campaign
python -m minecraft_wm.sweep --config configs/sweep.yaml --name lr_deter
python scripts/campaign.py --hours 9

Outputs land in runs/<name>/: TensorBoard logs, latest.pt, and pred_*.png ([truth | model | error] video-prediction strips).

Key gotchas (learned the hard way — see DEVLOG.md)

• Store episodes uncompressed. numpy ignores mmap_mode on compressed .npz, so a compressed dataset decompresses a whole episode per window — an 11× slowdown. The pipeline uses uncompressed .npy + true mmap.
• Frame/action alignment: VPT has N+1 frames vs N actions; action[t] led into frame[t] ((s_{t-1}, a_t) -> s_t), with is_first[0]=True per window.
• Train with AMP (--precision 16); keep dataloader workers low (data is tiny mmap).
• The RSSM scan is launch-bound at small batch; large batch (≥256) makes it compute-bound.

License

MIT. VPT data and the vendored repos retain their own licenses.