VIScore: Diagnosing Planning-Relevant Quality in Latent World Models
Key results:
- 🔍 Covers all four stages that decide planning success: encoder, predictor, planner, task tolerance
- 🔥 High correlation with success rate: pooled ρ +0.90, against +0.82 for raw empowerment
- 💪 Best calibration across planners and datasets: below the constant-predictor reference on every pool
- 🚀 20× faster than the real evaluation: ~7 s per checkpoint, no environment and no planner rollouts
- 🧩 Model-agnostic: three methods (
encode, action_embed, predict_next) are the whole interface
- ♻️ CPU-reproducible: every table in the paper recomputes in seconds
Available Checkpoints
Each file is a full-object torch.save pickle of jepa.JEPA: ViT-Tiny encoder + action-conditioned AdaLN transformer predictor, ~15M parameters, 72 MB. Filenames say which model the weights are — vis-wm_epoch_<N>.ckpt, lewm_epoch_<N>.ckpt (the SIGReg baseline) or dino-cls_epoch_<N>.ckpt.
| Path | Contents | Files | Size |
|---|
vis-wm/<task>/seed<S>/ | VIS-WM at the epoch the paper reports, + config.yaml | 39 | 0.88 GiB |
baselines-lewm/maze2d/seed729/ | SIGReg (LeWM) baseline on MAZE | 3 | 0.07 GiB |
pools/<pool>/<env>/<run>/ | every checkpoint behind a reported number, plus the calibration pool | 547 | 36.7 GiB |
VIS-WM: PushT / Reacher / Two-Room / Cube × seeds 403, 46, 729, plus MAZE at seed 729. Reported epoch per arm — maximum success on the single-seed d=25 evaluation, ties to the later epoch:
| Task | seed 403 | seed 46 | seed 729 |
|---|
| PushT | 9 | 10 | 10 |
| Reacher | 5 | 10 | 4 |
| Two-Room | 5 | 10 | 9 |
| Cube | 6 | 8 | 4 |
| MAZE | — | — | 8 (SIGReg baseline: 10) |
Pools. The split is by training run, so no checkpoint in a test pool comes from a run used to
fit anything. Every directory says what its checkpoints are for:
| Directory | Ckpt / runs | Which reported number it produces | Evaluation seeds |
|---|
pools/development/ | 137 / 14 | block 1 of tab:viscore; the constants are fitted here | all available |
pools/heldout/ | 103 / 33 | block 2 -- the headline test, runs disjoint from development | 5501, 60601, 90210 |
pools/heldout-method/ | 3 / 3 | block 3, our DINO-CLS arms (the other 27 are other groups', not re-hosted) | its own |
pools/heldout-dataset/ | 20 / 2 | block 4 -- the unseen MAZE task | its own |
pools/cube-reference/ | 50 / 5 | the parenthesised Cube column, and nothing else | 5501, 60601, 90210 |
pools/calibration-fit/ | 232 / 47 | no reported number: these fit the frozen map blocks 3--4 apply | all available |
So: 340 checkpoints produce every number in the paper (313 hosted here + 27 fetched with
reproduce/download_external.py), 572 if you also want to refit the frozen calibration map
from weights, and 0 if you only want to recompute the tables from the shipped metric values.
Per-task counts in the two ranked pools (PushT / Reacher / Two-Room): development 57 / 50 / 30,
held-out 27 / 44 / 32.
heldout absorbed what earlier versions reported as a separate terminal pool: 14 of that pool's
17 runs were also in the held-out fold, and restricted to converged epochs its labels no longer
resolve individual checkpoints.
cube-reference is separate from heldout because Cube entered no fold. Its label spread does not
exceed its own binomial standard error, so nothing there is rankable; the paper prints its
correlation in parentheses and excludes it from every pooled and calibration number. These
checkpoints are published because that exclusion argument is measured on them.
calibration-fit is the part of the 472-cell calibration pool that is not already in
development or heldout (both of which are subsets of it). Its checkpoints enter no correlation
and no reported error -- they only determine the metric-to-success map that the two transfer blocks
apply without refitting. They are published so that a competing metric can be calibrated on the
same 472 cells.
Three SIGReg Two-Room runs landed in neither fold and are not released, which is why Two-Room
carries no SIGReg checkpoint in heldout.
The rest of the held-out-method pool is other groups' released checkpoints (Qantara, RC-aux,
INTACT), which are not re-hosted here; fetch them with reproduce/download_external.py.
pools/pool_manifest.csv gives each checkpoint's pool membership, success labels and seven metric values; pools/pool_assignment.csv gives run → fold.
Training
LeWM architecture with the VISReg regularizer, λ = 4.5. AdamW, linear warmup + cosine decay, 10 epochs, batch 128, lr 1e-4. Two-Room uses history_size=1, the other tasks 3.
Usage
Score a checkpoint
1from viscore import score_checkpoint
2
3f = score_checkpoint("vis-wm/pusht/seed403/vis-wm_epoch_9.ckpt",
4 "probes/probe_pusht.npz", task="pusht")
5print(f) # VIS + veracity / influence / sobriety + their inputs
The pickle names jepa and module as top-level modules; viscore binds the shipped copies automatically, so no LeWorldModel checkout is needed.
Download with huggingface_hub
1from huggingface_hub import hf_hub_download, snapshot_download
2
3# one checkpoint
4path = hf_hub_download(repo_id="BooBooWu/viscore",
5 filename="vis-wm/pusht/seed403/vis-wm_epoch_9.ckpt")
6
7# a whole tree
8snapshot_download(repo_id="BooBooWu/viscore", allow_patterns=["vis-wm/*"])
Build a probe
A probe is one frozen slice of training data (rng(0), 300 episodes, frameskip 5) shared by every checkpoint compared.
1viscore probe --tasks pusht --data-home $STABLEWM_HOME --out-dir probes
2viscore score --task pusht --probe probes/probe_pusht.npz \
3 --run-dir vis-wm/pusht/seed403 --epochs 9 --csv out.csv
Evaluation
Both of the paper's main tables recompute from files shipped in the
GitHub repo — no GPU, dataset or checkpoint download:
1git clone https://github.com/HaiyuWu/viscore && cd viscore && pip install -e .
2
3python reproduce/tables.py # metric vs success, three pools
4python reproduce/planning_tables.py --strict # planning tables, gated against published values
Notes: peak epoch = maximum evaluation success rate, ties to the later epoch. Cube is reported in parentheses and excluded from pooled and calibration columns. All published labels and sobriety probes use CEM. A single 50-episode evaluation has a success sd of ≈6 points.
Citation
1@article{wu2026viscore,
2 title = {VIScore: Diagnosing Planning-Relevant Quality in Latent World Models},
3 author = {Wu, Haiyu and Balestriero, Randall and Levine, Morgan},
4 journal = {arXiv preprint arXiv:2608.11174},
5 year = {2026},
6 eprint = {2608.11174},
7 archivePrefix = {arXiv}
8}
Architecture and training recipe from
LeWorldModel; regularizer from
VISReg.
License
This project (code and checkpoints) is released under the
MIT License.
