1 Introduction↩︎

Modern generative world models render strikingly realistic, action-controllable futures across diverse environments [1]–[4]. However, the rollouts they produce frequently hallucinate: they remain visually fluent and superficially plausible while drifting away from the ground-truth dynamics [5]. The term is borrowed from the language modeling literature [6], [7], where it typically denotes generation of factually incorrect text; analogous failures have also been studied in image [8] and video generation [9]. In a world model, the failure is arguably more consequential: hallucinated trajectories are fed directly into downstream planners and policies [10]–[12], so silent hallucination during rollout translates into silently incorrect decisions during control.

Despite the increasing fidelity of these models, where an autoregressive rollout will hallucinate, and why, is poorly understood. A natural reading is that hallucination is an architectural problem to be solved by ever-larger backbones and more training compute [13]. However, we argue that hallucination in world models is primarily a data coverage problem: it concentrates in low-coverage regions of the state-action space [14], [15] and is therefore both predictable from signals available at runtime [16]–[18] and preventable by adjusting which data the model is trained on rather than architectural changes. Our experiments reveal that a single underlying cause – coverage gaps – explains failures at every stage of the model pipeline: tokenizer, action-conditioning, and multi-step rollout, and that it manifests as three distinct failure modes shown in Figure 1.

In this work, we set out to characterize different types of hallucinations, predict when they will occur using signals internal to the model, and use that information to close the underlying coverage gap. We explore two concrete ways in which the gap can be closed: (i) using coverage-aware sampling techniques during large-scale pretraining, and (ii) by targeted online data collection using our derived predictors of hallucination. Understanding when and why hallucinations happen requires three resources that no single benchmark currently provides: full control over the training corpus, behaviorally diverse data spanning many tasks and domains, as well as live environments to probe coverage gaps via online interaction. To address this gap we develop MMBench2, a massively multitask dataset for visual world modeling that extends MMBench [19] with \(65.6\)k mixed-quality trajectories equivalent to \(427\) hours of \(224{\times}224\) video at \(15\) fps, complete with ground-truth action and reward labels and live simulators across \(210\) tasks spanning \(10\) domains. Of the \(210\) tasks, \(200\) form the pretraining corpus and \(10\) are held out as entirely unseen transfer tasks, allowing us to probe coverage both within and beyond the training distribution.

Using MMBench2, we train a \(350\)M parameter Dreamer 4 [3] world model and dissect its hallucination behavior. We identify three distinct hallucination modes, each anchored to a different stage of the pipeline: perceptual hallucination in the encoder-decoder pair, action marginalization in the dynamics model, and scene divergence during multi-step rollouts. We develop three predictors of hallucination: tokenizer round-trip residual, flow instability, and inter-seed denoising variance, and use them to detect hallucination at runtime, as well as curiosity reward for online data collection [20] in seen and unseen tasks. Empirically, our hallucination-driven data collection enables adaptation to entirely unseen environments with just \(50\) real trajectories, approaching the effectiveness of human data collection. Together, these results show that hallucination in modern world models is, to a large extent, a coverage problem, and that the same signals that detect it can also be used to mitigate it. We summarize our contributions as follows:

• MMBench2: a large \(427\)-hour dataset for visual world modeling with ground-truth actions and rewards, live environments, and behaviorally diverse data including human play.

• A stage-by-stage characterization of hallucination in generative world models tied to tokenizer, action marginalization, and rollout failures.

• Three hallucination predictors that detect hallucination without labels or additional training.

• A coverage-aware training recipe that reduces hallucination across all three predictors and improves rollout fidelity at no additional cost.

• A framework for targeted data collection that adapts a pretrained \(350\)M world model to unseen environments with as few as \(50\) trajectories.

To support further research on generative world modeling, we release our full dataset, code for training and evaluation, model checkpoints, and a browser interface for open-ended interaction with the world model. See https://nicklashansen.com/mmbench2 for videos and resources.

2 MMBench2: A Dataset for Visual World Modeling↩︎

To fully understand hallucination in large generative world models, it is necessary to have control over the training process as well as access to its training data, but (to our knowledge) no such resource exists today at a satisfactory scale. To address this gap, we propose MMBench2: a massively multitask dataset for visual world modeling. Our dataset consists of \(\mathbf{65{,}600}\) trajectories (\(427\) hours of \(224{\times}224\) video at \(15\) fps for a total of \(23\)M frames) with diverse behaviors (including human play data) across \(\mathbf{210}\) different continuous control tasks, complete with ground-truth action and reward labels as well as live environments and language instructions for every task included in MMBench2. As the name implies, MMBench2 directly builds on top of MMBench which provides a collection of 200 live environments unified under a common interface. Figure 2 shows the breadth of our tasks.

Task overview. MMBench2 spans \(\mathbf{10}\) task domains: DMControl [21], DMControl Extended [19], Meta-World [22], ManiSkill3 [23], MuJoCo [24], MiniArcade [19], Box2D [25], RoboDesk [26], OGBench [27], and Continuous Atari [28]. Together they cover locomotion, dexterous and tabletop manipulation, goal-conditioned navigation, and arcade-style environments, with continuous action spaces of \(1\) to \(16\) dimensions; we zero-pad every action vector to the maximum dimension (\(d_a{=}16\)) and supply a per-dimension validity mask so models can ignore inactive dimensions per task. Of the \(\mathbf{210}\) tasks, \(200\) form the pretraining set and \(10\) (proposed in this work) are held out as unseen transfer tasks.

The 200-task pretraining corpus contains 260 episodes per task. However, episode lengths vary by task, ranging from \(25\) for certain ManiSkill3 manipulation tasks up to \(1{,}000\) for Atari games. As a result, the corpus is highly non-uniform across tasks as shown in Figure 3. Additionally, the visual diversity within a task also differs wildly across the corpus with, e.g., Inverted Pendulum (MuJoCo) varies little whereas tasks like Dungeon Explorer (proposed in this work), Bipedal Walker (Box2D), and Krull (Atari) vary greatly between frames and episodes. Please refer to Appendix 8 for an overview of tasks and Appendix 9 for a detailed description of the data collection process.

3 Training a Large Visual World Model↩︎

We seek to answer the following question: why do world models hallucinate, and can we predict when they are likely to hallucinate? To answer this question, we choose to develop and train a 350M parameter visual world model that largely follows the overall architecture and training recipe introduced in Dreamer 4 [3]. As such, our world model is an action-conditioned generative model that follows an increasingly common two-stage training recipe: (1) first, a video tokenizer is trained via masked auto-encoding of visual observations, and (2) a dynamics model implemented as a block-causal Transformer [29] is subsequently trained via flow-matching over spatial latent tokens produced by the frozen tokenizer with additional conditioning on action tokens. We first pretrain our world model on the MMBench2 training corpus, and then address our specific research questions by conducting a series of targeted finetuning experiments. This section aims to provide an overview of our architecture and training recipe.

Tokenizer. Our video tokenizer is instantiated as a symmetric encoder-decoder Transformer architecture. The encoder ingests \(224{\times}224\) RGB frames “patchified” at stride \(14\) (\(256\) patch tokens per frame), prepends \(64\) learnable latent queries, and projects the latent stream to a \(64\)-dimensional bottleneck bounded by a \(\tanh\) activation, producing a per-frame code \(z \in [-1, 1]^{64\times 64}\). The decoder then reconstructs images conditioned only on latent codes. Training uses a masked reconstruction objective [30] where a fraction of input patches drawn per-frame from \(\mathcal{U}(0, 0.9)\) are replaced by a learned mask token, and the loss (pixel MSE + LPIPS [31]) is computed only on masked positions. We normalize losses by their running RMS to reduce sensitivity to hyperparameters. Encoder and decoder are 50M learnable parameters each.

Dynamics model. Our dynamics model is a 250M parameter block-causal Transformer over packed tokenizer latents, trained on top of the frozen tokenizer introduced above. Per timestep, the input sequence consists of an action token (a \(2\)-layer MLP over the \(16\)-dimensional padded action), a shortcut-conditioning token (encoding noise level \(\sigma\) and step size \(d\)), \(32\) packed spatial latent tokens, \(4\) register tokens, and optional agent tokens used as the readout for the reward and behavior-cloning heads. The model is trained with the shortcut flow-matching objective of [32], which interleaves an empirical one-step regression term with a self-consistency bootstrap that distills two coarser-step predictions into one finer-step target. At inference this enables next-frame sampling in as few as \(4\) Euler substeps.

Training recipe. The training recipe of Dreamer 4 was developed for a setting where only a small percentage of the training data has action and reward labels. However, MMBench2 provides ground-truth actions and rewards for every trajectory in the dataset which we choose to take advantage of. We first pretrain the tokenizer on the full training corpus, and then proceed to also pretrain the dynamics model, conditioning on actions. After this initial pretraining phase, we experiment with additional mid-training and finetuning recipes introduced in the following sections. As part of our experiments, we initialize additional reward prediction and behavior cloning (BC) policy heads after pretraining. The reward predictor is trained via \(L{=}8\) multi-step discrete regression (using \(\operatorname{symlog}\) two-hot encoding) with gradients backpropagated through the dynamics model. The BC policy is a deterministic Gaussian policy trained to predict ground-truth actions via an MSE loss – this is a departure from Dreamer 4 which only considered discrete actions.

4 Hallucination in World Models↩︎

A generative world model imagines a future by chaining three distinct operations: (1) an encoder maps each observation into a latent code, (2) an action-conditioned dynamics head predicts the next latent, and (3) a decoder renders a reconstruction back to pixel space. Each of these stages is a learned function trained on a finite slice of the state-action space, and each can therefore fail independently when asked to extrapolate beyond what it has seen. We use the term hallucination to refer to any such failure: an output that is fluent and visually plausible, yet decoupled from reality. Crucially, because the three stages compose sequentially, a hallucination introduced early (e.g.a corrupted encoding) is propagated and amplified by the stages that follow, so naming which stage produces a given failure is a prerequisite for diagnosing and fixing it. This section explores how hallucinations can be characterized, detected, and finally mitigated.

4.1 Characterizing Hallucination↩︎

We identify three distinct ways in which a generative world model may hallucinate, each tied to a different stage of the imagination pipeline. We illustrate each type of hallucination in Figure 1 and define them as follows:

1. Perceptual hallucination. The tokenizer’s reconstruction of an observation already differs from the observation itself, before any dynamics rollout has been executed. Concretely, the encoder–decoder pair projects out-of-distribution scene structure onto the closest in-distribution exemplar in its learned latent “vocabulary”. For example, an unseen maze layout might be reconstructed with the agent and goal in the correct positions but with the walls of an entirely different layout seen during training. The dynamics head then rolls out against this corrupted scene as if it were ground-truth. This failure mode is a property of the frozen encoder–decoder pair alone and persists even at horizon \(H{=}0\).

2. Action-marginalized (ignored) hallucination. Conditional on a context, the predicted next latent is largely insensitive to the input action. The rollout is visually plausible but collapses onto an action-marginalized future, so the model behaves more like a video generator than a controllable world model. Operationally, we expose this mode by intervening on the action stream at evaluation time, e.g. by randomly shuffling actions within a batch and measuring the resulting change in flow MSE; a model that hallucinates in this way is one whose flow MSE barely moves under the intervention.

3. Scene-diverging hallucination. It is well understood that autoregressive rollouts accumulate compounding error as the prediction horizon increases. However, scene-diverging hallucination is a very specific failure mode where physically implausible events (such as a ball teleporting back into play when scoring in Pong) are predicted. This type of hallucination is most frequent in states with poor data coverage.

These three types of failure modes probe disjoint pieces of the model: the tokenizer, the action-conditioning of the dynamics model, and the multi-step accumulation of dynamics error.

4.2 Detecting Hallucination↩︎

Next, we develop three distinct predictors designed to track model hallucination. Empirically we find that these three metrics, although mechanistically distinct, are all strong predictors of hallucination events. We view this convergence as a feature and not a redundancy, since three mechanistically distinct predictors agreeing with respect to hallucination events is in itself evidence that the underlying signal is real. None of them require any labels nor additional training which makes them especially suitable for runtime detection.

Proposed predictor 1: tokenizer round-trip residual. \(u_r = \|\hat{z} - \mathrm{Encode}(\mathrm{Decode}(\hat{z}))\|\), the latent-space residual of a single decode–encode round-trip of the dynamics-predicted next latent \(\hat{z}\), targets perceptual hallucination by measuring the very symptom that defines it: a predicted latent whose decoded frame falls off the tokenizer’s manifold, e.g.a corrupted scene layout or a fabricated object, does not survive re-encoding and produces a large \(u_r\).

Proposed predictor 2: flow instability. \(u_f\), the flow instability of the dynamics head at a given \((\text{context}, \text{action})\) pair, measures how much the denoiser’s clean-target prediction \(\hat{x}_1\) moves between successive Euler integration substeps, averaged over the second half of substeps. A sharp, well-conditioned dynamics head converges quickly to a stable \(\hat{x}_1\) (low \(u_f\)); a head whose conditioning provides little signal keeps oscillating across substeps (high \(u_f\)).

Proposed predictor 3: inter-seed variance. \(u_s\), the inter-seed variance of the next-latent prediction across \(N\) independent denoising trajectories at fixed past and action, targets scene-diverging hallucination by measuring epistemic uncertainty across noise seeds: regions where seeds disagree are precisely those where multi-step rollouts will fan out.

Controlling for scene motion. In practice, we find that a naive use of the three signals above is confounded by scene activity: high-motion transitions inflate tokenizer residuals, flow instability, and seed-to-seed denoising dispersion alike. To remove this confound we instead report dynamism-normalized variants \(u^{\text{norm}} \doteq u/m\) where \(m\) is the per-step RMS change of the latent representation, computed as a per-task average over the dataset, or as a running estimate when used online for data collection. Normalizing by scene motion means that each predictor tracks model uncertainty relative to how much is happening in the scene. We treat \(u_{r}^{\text{norm}}, u_{f}^{\text{norm}}, u_{s}^{\text{norm}}\) as our primary hallucination predictors in all subsequent analysis.

4.3 Mitigating Hallucination↩︎

The taxonomy and predictors above suggest a single data-centric lens on hallucination: each of the three failure modes are, mechanistically, a consequence of the model having seen too little of some region of the state-action space. A perceptual hallucination is a coverage gap in the tokenizer’s reconstruction distribution, an action-marginalized hallucination is a coverage gap in action-conditional transitions, and a scene-diverging hallucination is a coverage gap along the trajectory the model is asked to imagine. Figure 4 shows the relationship between data coverage and hallucination on four of our tasks. Two interventions follow naturally:

Coverage-aware training. We propose to resample the existing dataset to upweight under-represented regions of the state-action space, then ask whether closing those gaps at training time reduces all three failure modes simultaneously. Because the lens above identifies coverage as the underlying lever for every type of hallucination, a single reweighting recipe is expected to move all three signals in the right direction at once, rather than requiring a separate intervention per type. In practice, we rebalance training data by adjusting sampling to be uniform across tasks rather than frames. We also experiment with loss re-weighting but find interventions in sampling to be superior.

Targeted data collection. When the existing dataset does not have enough coverage of a region for re-weighting alone to help, the hallucination predictors can themselves act as an objective for targeted, curiosity-driven data collection. Specifically, during interaction with a live environment, candidate trajectories are rolled out in the world model, scored by predicted hallucination, and the highest-ranked trajectory is executed in the environment, producing data that, by construction, covers transitions that previously caused the model to hallucinate. In practice, we collect data in a closed-loop manner with a prediction horizon of \(H=32\) and replanning every \(K=16\) steps.

5 Experiments↩︎

We train a 350M parameter action-conditioned world model on the training corpus of MMBench2. Our training data consists of approximately 20M frames at \(224\times224\) resolution across 200 distinct continuous control tasks, and we reserve the remaining \(3\)M frames for testing. Although low-dimensional state information is readily available in MMBench2, we choose to strictly consider RGB observations in this work. We describe our experimental setup below.

Evaluation. We evaluate our world model and its variants across four successive training phases: (1) action-conditioned pretraining without reward labels, (2) coverage-aware "mid-training" that extends the previous phase with our proposed reweighted sampling, (3) action-conditioned world modeling with reward labels, and (4) finetuning on seen and unseen tasks via targeted data collection. Our key evaluation metrics include:

1. Reconstruction PSNR \(\uparrow\) evaluates encoder/decoder quality without considering dynamics.

2. Rollout PSNR gain (dB) \(\uparrow\) evaluates quality of generated rollouts relative to a baseline that repeats the last frame across the entire rollout horizon. Although naive, this baseline can be surprisingly strong depending on the task. We consider a scene divergent when \(\Delta\)PSNR \(\le 0\).

3. Action shuffle ratio \(\uparrow\) evaluates action sensitivity in the dynamics model by measuring one-step teacher-forced flow MSE relative to batch-shuffled actions. We consider actions to be ignored (marginalized) when this ratio \(\le 1.1\).

4. Downstream task performance (normalized score) \(\uparrow\) evaluates how useful the world model is for downstream tasks via closed-loop MPC using the Cross-Entropy Method (CEM) with horizon \(H=32\) and replanning every \(16\) steps. We report normalized score \(s \in [0, 1]\) rather than raw reward since reward scales differ significantly (by several orders of magnitude) across tasks.

Implementation details. The encoder and decoder each have 50M learnable parameters, and the dynamics model (including prediction heads) has 250M parameters. We use \(8\times\) NVIDIA H100 GPUs for training. The tokenizer is pretrained for \(300\)k steps (\(14\) GPU days), and the dynamics model is pretrained for \(180\)k steps (\(24\) GPU days). Our final checkpoints (post-finetuning) are trained for \(380\)k and \(210\)k steps, respectively, for a total of \(58\) GPU days using a context length of \(T=24\). The reward head and BC policy are conditioned on CLIP-ViT/B embeddings of per-task language instructions. See Appendix 11 for more implementation details.

We seek to answer the following questions:

1. Hallucination detection. Do our proposed metrics accurately predict model hallucination?

2. Pretraining. Does coverage-aware training meaningfully reduce hallucination?

3. Targeted data collection. Can model hallucination be mitigated by collecting additional data? Given a fixed budget, which type of data improves the world model the most?

4. Generalization. How can a pretrained world model be adapted to unseen tasks? Does our world model generalize zero-shot or is finetuning necessary? When does it fail to transfer?

5.1 Main results↩︎

Detecting hallucination events. To determine whether our proposed hallucination metrics are predictive of hallucination and by extension model performance, we compute Spearman correlations between open-loop rollout \(\Delta\)PSNR and each of our predictors; results are shown in Figure 5. We find a strong negative correlation (\(\rho\approx0.80\)) across all three predictors, which indicates that they all track the same underlying error. We additionally measure their per-task aggregated AUROC against two binary hallucination labels, action ignored (action shuffle ratio \(\le 1.1\)) and scene divergent (rollout \(\Delta\)PSNR vs. frame-repeating baseline \(\le 0\)); results are shown in Appendix 10. We find that our predictors consistently outperform their unnormalized counterparts (i.e., \(u_f^{\text{norm}}\) is a better predictor than \(u_f\)) as well as baselines that rely on latent scene motion \(m\) or per-task frame count.

r0.6

Coverage-aware training. We evaluate the efficacy of our proposed coverage-aware training by extending the pretraining of our base model by an additional \(30\)k steps for the tokenizer and dynamics model each, while varying the sampling method: uniform sampling across frames (default), and uniform sampling across tasks (ours). Results are shown in Table ¿tbl:tab:coverage-aware-pretraining?. We find that tokenizer and dynamics model both benefit from coverage-aware training, and that adopting it for both achieves best overall performance.

Mitigating hallucination via targeted data collection. Since we observe a strong relationship between data coverage and hallucination, a question that naturally follows is can hallucination be mitigated by filling the gaps in coverage via targeted data collection? To answer this question, we construct a task set that consists of 10 tasks seen during pretraining and 10 completely unseen tasks, and then collect 50 episodes for each of these tasks varying only the behavior policy used. Specifically, we evaluate the following approaches: no-op (all-zero) actions, random actions, expert policies, human play (via keyboard), and curiosity-based exploration [20] using our proposed hallucination predictor \(u_r^{\mathrm{norm}}\) as objective. State densities for different collection policies are shown in Figure 6. We finetune tokenizer and dynamics model for \(50\)k steps and \(30\)k steps, respectively, and report offline performance metrics and downstream task performance via closed-loop MPC (planning with CEM). Results for the 10 unseen tasks are shown in Table 1, and additional results can be found in Appendix 10. Pretraining transfers, to some extent, zero-shot (\(0.276\), \(2.3{\times}\) the \(0.118\) random policy baseline), and collecting just \(50\) trajectories per task using \(u_r^{\mathrm{norm}}\)-based curiosity lifts performance to \(0.325\), within \(\sim 90\%\) of the expert/human oracles (\(0.362\)) despite using no privileged behavior.

5.2 Discussion↩︎

Empirically, we find that hallucination can be mitigated, to a great extent, by targeted data collection, but that not all data sources are equally informative. We believe that there are two concrete reasons for this: (i) world modeling requires broad state-action coverage which, e.g., a random policy does not guarantee, and (ii) downstream tasks are goal-directed and have a much narrower state-action distribution than the general space, so e.g. our closed-loop MPC evaluations mainly measure model accuracy around that narrower subspace. That is, arguably, by design: there are certain behaviors we are more interested in modeling than others, but it raises a broader question about how we can best evaluate world models moving forward.

Do off-the-shelf tokenizers improve perception? One concrete way in which perceptual hallucination can be mitigated on unseen tasks is by leveraging off-the-shelf tokenizers trained on datasets several orders of magnitude larger than the one considered in this work. To investigate the viability of this approach, we compare our trained tokenizers (before and after finetuning on the unseen task set) against four off-the-shelf tokenizers; results are shown in Table 2. We find that off-the-shelf tokenizers such as Wan 2.1 VAE (the strongest) underperform compared to our in-domain tokenizer when evaluated on tasks from the \(200\)-task training corpus, but perform significantly better on the unseen task set unless we finetune our model in which case our in-domain tokenizer wins again. These results suggest that off-the-shelf tokenizers indeed are promising for world modeling but that there are still tangible benefits to in-domain finetuning when possible.

Related work. We include a detailed discussion of related work in Appendix 6, and position MMBench2 vs. existing datasets for world modeling in Appendix 7.

In conclusion, we argue that hallucination in generative world models is, first and foremost, a data-coverage problem, and identify three failure modes: perceptual, action marginalization, and scene divergence. Our predictors track them with \(\rho \approx 0.80\) against rollout \(\Delta\)PSNR and yield two complementary recipes: coverage-aware sampling reduces all three failure modes simultaneously, and using the predictors as curiosity rewards adapts our pretrained model to entirely unseen tasks.

Limitations. Our study operates at the \(350\)M parameter scale across \(210\) simulated control tasks. Whether our findings translate to billion-parameter models or to real robot data with sensor noise and partial observability remains an open empirical question. We also acknowledge the significant computational cost of training large world models, as well as the need for more diverse datasets.

6 Related Work↩︎

World models for control. World models that learn environment dynamics from data have a long history in model-based RL, ranging from abstract latent dynamics models [12], [33]–[35] to high-capacity generative models that render full pixel observations [1], [2], [36]. Recent work scales these models to heterogeneous video corpora [37]–[39] and to real-time, playable neural environments [4], [40], [41]. We build on Dreamer 4 [3], an integrated tokenizer-plus-dynamics architecture with strong action conditioning, but the signals and interventions we propose are largely model-agnostic and apply, in principle, to any modern generative world model. Despite striking visual fidelity, these models continue to hallucinate under distribution shift — a failure mode that we set out to characterize, predict, and mitigate.

Hallucination and uncertainty in generative models. Hallucination has been studied extensively in language models [6], [7] and, increasingly, in image [8] and video generation [9], but the question of where an autoregressive world model will fail along a rollout has received comparatively little attention. A separate but related literature studies uncertainty estimation in deep networks through deep ensembles [16], MC dropout [17], and post-hoc out-of-distribution detectors [42], [43], with applications in offline model-based RL via uncertainty-penalized policies [44], [45]. Closest in spirit, prior work uses ensemble disagreement as an exploration signal in single-task RL [20], [46]. In contrast, our proposed predictors are derived directly from the existing world model and target the three distinct stages at which a generative world model can hallucinate (encoder, dynamics, decoder), without auxiliary networks or labels.

Coverage-aware training and data collection. A growing body of work argues that data scale and composition are first-order levers for generative model performance [13], [15]. In offline RL specifically, data coverage is known to bound policy improvement [14], [47], motivating curated datasets such as ExoRL [48], RL Unplugged [49], and V-D4RL [50]. Curiosity-driven exploration is the natural online counterpart: agents are incentivized to visit states with high prediction error [51], high feature-network novelty [52], or high model disagreement [46], scaled to imagined trajectories in Plan2Explore [20]. Whereas prior curiosity work uses these signals to drive single-task policy exploration, we adapt them to data-collection for generative world modeling in two complementary ways: (i) uniform-task resampling closes a substantial fraction of the hallucination gap at no additional data cost, and (ii) using our proposed predictors as curiosity rewards yields a data-efficient finetuning recipe that generalizes a 350M-parameter world model to entirely unseen environments with as few as 50 trajectories from the target task.

7 Comparison to Existing Datasets↩︎

Table 3 compares MMBench2 to a representative set of prior datasets used for offline reinforcement learning, robot imitation learning, and large-scale generative video and world modeling. The offline RL datasets we compare against include RL Unplugged [49], V-D4RL [50], ExoRL [48], the TD-MPC2 multi-task dataset [12], and the Atari DQN Replay dataset [53]. The robot imitation learning datasets we compare against are RoboNet [54], BridgeData V2 [55], DROID [56], and Open X-Embodiment [57]. For large-scale video pretraining corpora with pseudo-labeled actions, we compare against VPT [58] and NitroGen [59]. Finally, we include MMBench [19], the multi-task benchmark on which MMBench2 builds. Across these datasets, MMBench2 is the only corpus that simultaneously offers ground-truth action and reward labels, live simulators for every task, mixed-quality behavior, and broad coverage across both task domains and embodiments.

8 Task Domains↩︎

We consider 210 tasks across 10 domains. Our task set is comprised of diverse continuous control tasks spanning robot manipulation, locomotion, navigation, arcade games, and classic control problems, each varying in task complexity, time horizon, observation and action space dimensionality, and reward formulation. Table 4 provides an overview of our task domains.

9 Data Collection↩︎

We base our data collection on that of MMBench, a 200-task benchmark designed primarily for online RL; it provides a total of 4k expert demonstrations collected via single-task expert policies, as well as live environments with ground-truth action and reward labels for all tasks. While MMBench is well suited for the multi-task online RL setting it was originally developed for, a dataset that consists solely of expert demonstrations lacks diversity in terms of behavior which, as our experiments show, is a significant source of hallucination in world models.

The overarching goal of MMBench2 is to produce a large, diverse dataset for visual world modeling and research in common failure modes such as hallucination. To do so, we use the single-task expert policies \(\pi^{\star}\) of MMBench as a basis for our data collection, but crucially augment policies to generate diverse behaviors, and also collect data via non-expert policies, including humans. Our methods of data collection can be summarized as follows:

• Random policy. Sample actions uniformly in \([-1, 1]\). Diverse actions; poor task performance.

• No-op actions. Set actions \(\mathbf{a} = \mathbf{0}\). Models dynamics without agent interference.

• Expert actions. Sample from the expert policy \(\pi^{\star}\) without added noise. High task performance; low data diversity despite it being a stochastic policy.

• Transformed expert actions. Sample from \(\pi^{\star}\), then apply any of five transforms with some probability: scale by \(\varepsilon \in [0, 1]\), action dropout, flip sign, all-zero action, repeat previous action. The same transform may be applied for multiple steps. Provides counterfactual transitions.

• Structured noise. Sample from \(\pi^{\star}\), then apply any of three noise types with parameters sampled on a per-episode basis: Gaussian noise, Ornstein-Uhlenbeck noise (temporally correlated), convex mixture with random policy. Mixed performance; improves data diversity.

• Curiosity-driven. Sample trajectories that maximize our developed hallucination predictor \(u_r^{\mathrm{norm}}\), selected by CEM-based planning with the pretrained world model.

• Human play data. Actions are selected by a human interacting with the environment via a keyboard interface. High data diversity; not task-driven.

With the exception of human play data, we frequently switch between any of the above behaviors within a single episode to maximize diversity of our training data. For example, we may apply structured noise at random for a number of steps to visit a less frequently visited part of the state-action space and then switch back to the expert policy to generate recovery behavior. Data is collected across all 210 live environments, and ground-truth actions (the action selected by the behavior policy) and reward labels (task reward obtained from the environment) are logged. We collect pretraining data for 200 tasks, and additional data for targeted finetuning across 10 seen (included in pretraining) as well as 10 unseen tasks, with data partitions separated by behavior policy: random, noop, expert, mixed, curiosity, and human. Figure 9 shows a screenshot of the interface used to collect human play data.

10 Additional Results↩︎

11 Implementation Details↩︎

Our Dreamer 4 world model is a reproduction of the original method as described in [3]. This section provides an overview of our implementation and design choices.

Language embeddings. For task conditioning we use frozen text embeddings from openai/clip-vit-base-patch32 (CLIP; [60]), which produces \(512\)-dimensional continuous embeddings of per-task language instructions. Only the reward prediction and BC policy heads are conditioned on language embeddings.

Block-causal Transformer backbone. Each Transformer is a stack of block-causal layers consisting of (i) space self-attention over the per-frame token sequence, (ii) causal time self-attention along the temporal axis, and (iii) a SiLU-gated MLP with ratio \(4\). Attention uses RoPE [61] on Q/K with a KV-cache-aware offset, QK-normalization [62], and RMSNorm [63] pre-norm with no biases on the normalization layers.

Modality-aware mask. Within space self-attention, we apply a modality-aware mask that depends on the role of each token. In the tokenizer encoder, latent queries attend to all tokens while patch queries only attend within the image modality, preventing patch tokens from mixing across modalities before being bottlenecked through the latents; in the decoder, the directions are reversed so patch queries can read from the latent bottleneck but not from each other directly. In the dynamics model, action, shortcut, spatial, and register tokens are mutually visible while agent tokens (reward head and BC policy) are asymmetrically isolated: agent queries attend to everything, but non-agent queries do not see agent keys.

Spatial packing. The tokenizer produces per-frame latents of shape \((n_L, d_b) = (64, 64)\). Before being fed to the dynamics, these are spatially packed at factor \(k{=}2\) to shape \((n_\mathrm{spatial}, d_\mathrm{spatial}) = (32, 128)\), halving attention cost along the spatial axis at the cost of doubling channel dimension; the inverse unpack is applied at decode time. Concretely, the dynamics model sees the following token layout per timestep: [ACTION x 1, SHORTCUT x 1, SPATIAL x 32, REGISTER x 4, AGENT x 4] where the action token is produced by a \(2\)-layer MLP, the shortcut-conditioning token concatenates two embeddings of the discretized noise level \(\sigma\) and step size \(d{=}1/2^\mathrm{step}\), register tokens [64] are \(4\) learnable “sink” tokens, and agent tokens are initialized from the per-task CLIP embedding broadcast over time.

Loss normalization. Pixel MSE and LPIPS terms in the tokenizer, the empirical and self-consistency branches of the dynamics objective, and reward two-hot cross-entropy are each separately normalized by their own running RMS before weighting. This decouples loss weights from absolute scale and removes the need to tune them when the dataset, resolution, or backbone changes.

Shortcut flow matching. The dynamics model is trained with the shortcut flow-matching objective of [32]. The discretized noise level \(\sigma\) is indexed by an integer in \(\{0, \dots, k_\mathrm{max}\}\) (with \(k_\mathrm{max}{=}64\) in our experiments; \(0\) is pure noise, \(k_\mathrm{max}\) is clean) and the step is indexed by an integer in \(\{0, \dots, \log_2(k_\mathrm{max})\}\) corresponding to step size \(d{=}1/2^\mathrm{step}\). For a fraction \(\rho_\mathrm{self}{=}0.25\) of each batch we apply a self-consistency bootstrap: at \((\sigma, \mathrm{step})\), two coarser half-steps at \(\mathrm{step}{+}1\) are run under no_grad and their averaged velocity is used as a stop-gradient target for the current step’s predicted velocity. The remaining \(0.75\) of the batch uses the empirical one-step regression term at the finest step.

Reward and BC heads. Both heads read from the agent tokens via attention pooling against a learnable query. The reward head predicts \(L{=}8\) multi-step symlog two-hot distributions over \(255\) bins on the range \([-10, 10]\). The BC head is a deterministic Gaussian policy with diagonal covariance trained via an MSE loss on the ground-truth \(16\)-dimensional padded action.

Sampling. At inference we run a shortcut Euler integrator with step size \(d{=}0.125\) (i.e., \(K{=}8\) substeps) for each new frame: starting from \(z\sim\mathcal{N}(0, I)\), we solve \(b = (\hat{x}_1 - z)/(1-\sigma)\) followed by \(z \leftarrow z + b\cdot d\). We apply context corruption to past tokens.

Hyperparameters. We summarize our hyperparameters in Table 10.

References↩︎