nicklashansen/mmbench2

TL;DR: Hansen and Wang (UCSD) release the full stack behind a paper arguing that world-model hallucination — the drift into physically implausible futures that makes learned simulators untrustworthy for planning — can be flagged at inference time by cheap runtime signals, and largely trained away by using those same signals to steer data collection. The release is unusually complete: 350M-parameter Dreamer-4-style checkpoints, a new 210-task / 427-hour multi-domain dataset (MMBench2), the predictor and mitigation code, and a browser demo you can poke at directly.

What it is and why it matters

The central practical objection to model-based control with generative world models is that rollouts hallucinate: objects teleport, dynamics dissolve, and a planner optimizing against the model exploits exactly those errors. This work makes two claims. First, hallucination is predictable — three inference-time uncertainty signals correlate with it well enough to overlay live in an interactive session. Second, it is preventable — using those signals as a curiosity reward for closed-loop data collection, plus coverage-aware resampling during training, produces finetuned models that hallucinate less. If both hold up, it changes the calculus on whether world-model planning can be trusted, since it gives you both a runtime abort signal and a recipe for closing coverage gaps.

How it works

The world model itself is a faithful PyTorch reimplementation of Dreamer 4, adapted to multi-task continuous control: a ~100M causal video tokenizer (masked auto-encoding of 224×224 RGB into 64-dim tanh latents) plus a 250M block-causal Transformer dynamics model trained on frozen tokens with a shortcut flow-matching objective, action-conditioned, with reward and BC heads bolted on post-pretraining.

The novel machinery sits on top:

• Three hallucination predictors (src/uncertainty.py), each motion-normalized so fast scenes aren't automatically flagged: tokenizer round-trip residual u_r (does encoding/decoding the generated frame lose information — i.e., is the model off the tokenizer's manifold?), flow instability u_f (does the shortcut flow-matching sampler disagree with itself across step sizes?), and inter-seed denoising variance u_s (do independent denoising seeds diverge?).
• Coverage-aware training: resampling that upweights under-represented regions of state-action space during both tokenizer and dynamics training.
• Targeted data collection: the predictors become a curiosity reward driving closed-loop online collection in the live simulators (collect_data.py, curiosity.py), so new data lands exactly where the model is unreliable.

The dataset is a contribution on its own

MMBench2: 210 tasks across 10 domains (DMControl, Meta-World, ManiSkill3, MuJoCo, RoboDesk, OGBench, continuous Atari, and more), 65,600 mixed-quality trajectories, ~23M frames at 224×224/15fps, with actions, rewards, language instructions, and — crucially — live simulators for every task, enabling the closed-loop collection story. 200 tasks for pretraining, 10 held out for transfer. Actions are 1–16-dim, zero-padded with validity masks.

What's actually released, and what's missing

Everything is public under MIT: training/eval code, three checkpoint variants (base, coverage_aware, combined), the full dataset on HF, and a hosted live demo. The README, however, contains no quantitative results — no hallucination-detection AUCs, no before/after mitigation numbers, no closed-loop control performance. All evidence lives in the paper/interactive page, so you'll need to read those before trusting the headline. Two other caveats: 350M parameters is small by current world-model standards, so whether the predictors and the mitigation scale is open (the round-trip and inter-seed signals should transfer; coverage-aware resampling may matter less at scale); and the README doesn't say whether "preventable" is demonstrated in closed-loop control or only open-loop rollout quality — that's the key thing to check.

Practicalities: inference needs any CUDA GPU (≥4 GB for the interactive UI); training needs ~8×H100 and ~8 TB of preprocessed shards. A minor gymnasium/ale-py version conflict affects continuous Atari only. 94 stars and active updates through August 2026.

Try it

The interactive interface seeds from live simulators, so no dataset download needed:

conda env create -f environment.yaml && conda activate mmbench2
cd src
./run_interactive.sh combined   # serves at http://localhost:7860

Or skip local setup entirely: the hosted demo with live hallucination-signal overlays is at https://www.nicklashansen.com/mmbench2/#live-demo — driving the model into an under-covered state and watching the predictors spike is the fastest way to judge whether the claim is real.