1 Introduction↩︎

World models are commonly evaluated through the reinforcement-learning performance of the complete model-based reinforcement learning (MBRL) agent. However, strong agent performance does not guarantee that the learned world model is a reliable simulator on its own. In Atari Pong [1], we reproduce five representative visual world-model agents: DreamerV3, DIAMOND, TWISTER, Simulus, and STORM [2]–[6]. Their agent performance matches the original reports, yet their frozen world models are unreliable when evaluated on their own. Their generated trajectories contain failures such as ball disappearance and incorrect rebounds, and policies trained from scratch inside these models transfer poorly to the real environment. The contrast is especially large for DreamerV3: its mean Pong return falls from \(-5.5\) for the original agent to \(-20.9\) when policy learning is performed inside the frozen world model, near the minimum return of \(-21\).

Many of these rollout failures involve the ball and its interactions, despite the ball occupying only a few pixels. This leads us to investigate whether task-critical concepts receive insufficient learning signal. We propose Concept-Guided Spatial Regularization (CGSReg), which adds reconstruction supervision to image regions corresponding to concepts such as the Pong ball. It improves both closed-loop rollouts and pixel-space zero-shot MBRL in DreamerV3, DIAMOND, and TWISTER, but does not clearly improve Simulus or STORM. This contrast shows that concept-focused reconstruction addresses one important source of world-model error, but not every bottleneck.

A world model predicts future observations from past observations and actions. Depending on the task, an observation may be an image or a state describing the environment. In MBRL, the learned model generates future trajectories in which the policy can be trained.

Many MBRL agents follow a common training pattern: the current policy collects real-environment transitions, the world model is updated on the replay buffer, and the policy is improved using model-generated trajectories. These steps are repeated throughout training, so the replay distribution, world model, and policy evolve together. Following the Dyna framework [7], [8], we refer to agents trained in this way as Dyna-style agents. 1 summarizes this process.

We study five visual world-model projects that follow this Dyna-style training pattern: DreamerV3, DIAMOND, TWISTER, Simulus, and STORM. They span several model backbones, including recurrent state-space models, diffusion models [9], Transformers [10], and RetNet [11], as well as different world-modeling paradigms: latent-state prediction, direct pixel-space generation with diffusion models, and autoregressive next-token prediction. 1 summarizes these differences.

We focus on Atari Pong, an environment from the Arcade Learning Environment [1] and the Atari100K benchmark [12]. Pong contains two vertical paddles and a ball moving across the screen. The agent controls the player paddle on the right, while the game controls the opponent paddle on the left. The agent can move its paddle up or down, remain still, and combine these actions with FIRE, which starts or resumes play. A point is scored whenever one side fails to return the ball. An episode ends when either side reaches 21 points, so the return lies in \([-21,21]\).

Pong is visually simple, but its dynamics depend on accurately tracking a small ball and modeling its interactions with the paddles and walls. This makes world-model errors easy to inspect and directly relevant to control. In our closed-loop rollouts, the most common failures involve the ball or ball-related events: the ball disappears, changes direction without a valid interaction, or responds incorrectly to paddle collisions.

Pong’s action space contains six actions: NOOP, FIRE, RIGHT, LEFT, RIGHTFIRE, and LEFTFIRE. NOOP leaves the paddle unchanged, while FIRE starts or resumes play. In Pong, RIGHT moves the player paddle upward and LEFT moves it downward. RIGHTFIRE and LEFTFIRE combine the corresponding movement with FIRE. In our figures, we abbreviate these actions as \(-\), F, R, L, RF, and LF, respectively.

We evaluate the frozen world models with two complementary diagnostics. In the closed-loop rollout diagnostic, a policy trained separately from the corresponding MBRL agent interacts with the frozen world model, allowing us to inspect the generated video trajectories. Across all five models, these rollouts contain clear failures involving the ball, including disappearance, incorrect motion, and invalid ball–paddle interactions.

Beyond visual trajectories, we introduce pixel-space zero-shot MBRL. Zero-shot MBRL was introduced by MoSim [13] as a challenging evaluation for state-space world models: a new policy is trained entirely inside a frozen world model and then evaluated in the real environment. We extend this setting to visual world models, where the policy learns directly from generated pixel observations. Across all five models, the resulting policies perform substantially worse than the policies produced by the corresponding original MBRL training pipelines.

Together, these observations motivate our hypothesis that insufficient modeling of task-critical concepts, particularly the ball in Pong, contributes to the frozen-model failures. We therefore propose CGSReg, which adds reconstruction supervision to the corresponding image regions.

To isolate the effect of CGSReg from the co-evolving policy and replay distribution in Dyna-style training, we adopt offline world-model training. For each project, we train matched baseline and CGSReg world-model checkpoints on the same fixed replay dataset, with policy learning disabled and the world-model update budget held fixed. CGSReg improves both closed-loop rollouts and pixel-space zero-shot MBRL in DreamerV3, DIAMOND, and TWISTER. Simulus and STORM show less consistent changes across the two evaluations, revealing bottlenecks not fully addressed by concept-region reconstruction.

Our contributions are:

• We isolate and directly evaluate the frozen world models learned by five reproduced Dyna-style visual MBRL agents. Closed-loop rollouts reveal recurring visual and dynamical failures, while pixel-space zero-shot MBRL shows that these models provide weak standalone environments for policy learning.

• We propose CGSReg, a spatial regularizer that adds reconstruction supervision to task-critical concept regions. In Atari Pong, we instantiate it on the segmented ball region. It improves both closed-loop rollouts and pixel-space zero-shot MBRL in DreamerV3, DIAMOND, and TWISTER.

2 Diagnosing Frozen Visual World Models↩︎

We use Pong as a controlled setting for diagnosing frozen visual world models. The visual scene is simple, but the dynamics are sensitive to small objects and precise ball-paddle interactions.

2.1 Closed-Loop Rollout Diagnosis↩︎

We drive all five reproduced frozen world models with the same external controller: a pixel-based PPO [14] policy trained with Stable-Baselines3 [15] in the real Atari Pong environment. We use the publicly available PongNoFrameskip-v4 checkpoint from Stable-Baselines3¹. This controller is independent of the policies used to collect the reproduced Dyna-style replay and is only used to probe the frozen simulators. It supplies actions, while the frozen world model predicts the next observation, reward, and termination signal.

2 shows representative failure cases from the five frozen world models. Although the rollouts often preserve the global Pong scene, they frequently break task-critical objects and interactions. Across these examples, we observe recurring failure modes including ball disappearance, paddle disappearance, spurious ball direction change, and incorrect ball-paddle interaction. These failure modes are not mutually exclusive and can appear across multiple world models. Beyond these object-level failures, TWISTER also exhibits local visual artifacts around the ball. These errors directly affect rewards and policy learning because the ball and paddles determine rebounds and scoring events. Furthermore, STORM exhibits an action-response failure, where the generated dynamics do not reliably follow the supplied actions. For example, the generated paddle can return toward its original position without an intervening action that would cause such a movement.

2.2 Frozen-Model Zero-Shot MBRL Evaluation↩︎

Rollout videos reveal visual and dynamical errors, but they do not directly measure whether a frozen world model can support policy learning. We therefore compare two settings for each reproduced project.

First, we evaluate the final policy of the reproduced Dyna-style agent. This policy is obtained from our reproduction of the project’s original joint training pipeline, and its performance matches that reported in the corresponding paper.

Second, we extract and freeze the final world model from the same reproduced agent, discard its policy and value components, and train a new policy from scratch entirely inside the frozen model. We refer to this setting as pixel-space zero-shot MBRL, because the new policy learns from pixel trajectories generated by the frozen world model and is then evaluated in the real environment. We use each project’s native policy-learning procedure, unlike the unified protocol introduced later in 4.3. World-model updates and additional real-environment data collection are disabled throughout policy training.

As shown in 3, the frozen-WM policies perform far below the corresponding reproduced Dyna-style policies across all five projects. Both policies are evaluated with deterministic action selection over 20 reset seeds in the real ALE Pong environment. 2 reports the exact scores and compares our reproduced Dyna-style agents with the scores reported in the original papers.

This diagnostic shows that a world model learned within a strong Dyna-style agent may still be a weak standalone environment for policy learning. Together with the closed-loop rollout failures, this motivates our hypothesis that task-critical concepts, particularly the Pong ball and its interactions, are insufficiently modeled.

3 Concept-Guided Spatial Regularization↩︎

CGSReg is an auxiliary regularization term that adds extra image reconstruction supervision to pixels corresponding to important concepts.

For the visual world models studied in this paper, we decompose the original world-model objective into an image-space loss and the remaining nonvisual objectives:

\[L_{\mathrm{wm}} = L_{\mathrm{img}} + L_{\mathrm{nonvisual}}, \label{eq:base-wm-objective}\tag{1}\]

where \(L_{\mathrm{img}}\) is the image reconstruction or prediction loss and \(L_{\mathrm{nonvisual}}\) collects the remaining objectives, such as dynamics, reward, and continuation losses.

CGSReg adds an additional masked image loss:

\[L_{\mathrm{wm}} = L_{\mathrm{img}} + \lambda_{\mathrm{CGSReg}}L_{\mathrm{CGSReg}} + L_{\mathrm{nonvisual}}, \label{eq:cgsreg-objective}\tag{2}\]

where \(\lambda_{\mathrm{CGSReg}}\) controls the regularization strength.

Specifically, we represent the pixels of an important concept with a binary mask \(m\). Given a target image \(x\) and a world-model output \(\hat{x}\), we define

\[L_{\mathrm{CGSReg}} = \frac{ \sum_p m_p(x_p-\hat{x}_p)^2 }{ \sum_p m_p }, \label{eq:cgsreg-loss}\tag{3}\]

where \(p\) indexes individual image values. For RGB images, the mask is applied to all color channels. Normalizing by the mask area ensures that each concept contributes according to its reconstruction error rather than its pixel size.

For Pong, we use the ball as the main task-critical concept and obtain its mask with SAM2 [16]. Paddle masks, multi-mask training, and loss-scale matching are described in 9.

4 Experiments↩︎

This section evaluates whether CGSReg improves frozen visual world models. For each project, we compare two matched offline world-model checkpoints and additionally evaluate the reproduced-agent WM, extracted from our reproduction of the original Dyna-style agent.

We adopt offline world-model training introduced in 1 to isolate the effect of CGSReg from policy learning and replay-buffer evolution in the Dyna-style loop. Under this setting, the baseline and CGSReg checkpoints are trained using the same fixed replay data and matched world-model update budgets, making them the controlled comparison for measuring the effect of CGSReg. Full training details are provided in 11.

The reproduced-agent WM serves a different role. It connects the controlled offline experiments to the frozen world models diagnosed in 2 and shows how the offline-trained models compare with a world model learned through each project’s native Dyna-style pipeline. It is included only as a reference and is not part of the controlled CGSReg comparison.

After freezing all checkpoints, we evaluate them using the two diagnostics from 2: closed-loop rollout diagnosis and frozen-model zero-shot MBRL. The rollout evaluation follows the setup in 2.1. For zero-shot MBRL, we use one shared policy-learning protocol across all projects and checkpoints, enabling direct comparison.

4.1 Compared World-Model Checkpoints↩︎

For each project, we evaluate three frozen world-model checkpoints:

• Reproduced-agent WM: the world model extracted from our reproduction of the original Dyna-style agent.

• Offline baseline WM: the \(\lambda_{\mathrm{CGSReg}}=0\) special case of the offline training objective, in which CGSReg is disabled.

• Offline CGSReg WM: the corresponding offline world model trained with a selected nonzero \(\lambda_{\mathrm{CGSReg}}\).

The two offline checkpoints form the controlled comparison. They are trained on the same fixed 100k-step Pong replay dataset collected from our DIAMOND reproduction and use the same world-model update budget; their only difference is the value of \(\lambda_{\mathrm{CGSReg}}\).

We use \(\lambda_{\mathrm{CGSReg}}=0.1\) for DreamerV3, \(0.01\) for DIAMOND, \(1.0\) for TWISTER, \(0.1\) for Simulus, and \(0.01\) for STORM. All main experiments use the ball mask. Weight sweeps are reported in [app:weight-ablation].

4.2 Qualitative Closed-Loop Rollouts↩︎

We evaluate all frozen checkpoints using the same external controller and horizon-512 rollout setup as in 2.1. We compare the offline baseline and CGSReg rollouts using the same qualitative criteria: object visibility, ball motion, collision behavior, and action response.

4 shows DreamerV3, for which CGSReg produces the clearest qualitative improvement. The ball remains visible through the displayed interaction and follows a plausible rebound trajectory.

Across the inspected rollouts, DIAMOND and TWISTER also show improved ball modeling and dynamics, although occasional visual artifacts and ball-prediction errors remain. Simulus does not show a clear rollout improvement, while STORM continues to exhibit the action-response failure identified in 2.1.

4.3 Unified Pixel-Space Zero-Shot MBRL↩︎

We apply one shared pixel-space zero-shot MBRL protocol to all three checkpoint types. For each frozen world model, we initialize a new image-based on-policy actor-critic and train it from scratch for 20k updates. As in 2.2, each final policy is evaluated with deterministic action selection over 20 reset seeds in the real ALE Pong environment.

The main comparison is between the matched offline baseline and CGSReg world models. The reproduced-agent WM is included only as a reference. Full protocol details are provided in 12.

The resulting real-environment Pong scores are shown in 5.

Compared with the matched offline baselines, CGSReg raises the mean real-environment Pong return for DreamerV3 \((-21.00 \rightarrow -11.90)\), DIAMOND \((-13.90 \rightarrow -5.80)\), TWISTER \((-21.00 \rightarrow -1.90)\), and Simulus \((-15.80 \rightarrow -4.10)\). STORM remains at \(-21.00\).

4.4 Architecture-Dependent Outcomes↩︎

Simulus and STORM reveal different limitations of concept-region reconstruction. In Simulus, CGSReg directly regularizes image-tokenizer reconstruction but not the subsequent token-dynamics model. Its higher zero-shot return therefore does not necessarily indicate better long-horizon token prediction, consistent with the lack of a clear rollout improvement.

For STORM, the main observed failure is its unreliable response to supplied actions. When the generated dynamics do not correctly follow the controller’s actions, improving reconstruction of the ball region cannot address the dominant error.

Overall, CGSReg improves both closed-loop rollouts and pixel-space zero-shot MBRL in DreamerV3, DIAMOND, and TWISTER. The mixed Simulus result and the unchanged STORM result indicate that concept-focused reconstruction addresses only some world-model bottlenecks.

5 Ablations↩︎

The main experiments fix the offline dataset and compare matched \(\lambda_{\mathrm{CGSReg}}=0\) and CGSReg checkpoints. We keep the ablation study separate from this claim. Appendix 13 reports additional sweeps over \(\lambda_{\mathrm{CGSReg}}\) and the offline replay dataset.

6 Limitations and Discussion↩︎

6.0.0.1 CGSReg Alone Is Not Sufficient.

Although CGSReg improves pixel-space zero-shot MBRL in several models, the resulting policies still do not solve Pong. As shown in 3, their mean returns remain far below the winning return of \(21\). Their performance also varies substantially across evaluation episodes, particularly for TWISTER and Simulus, whose standard deviations reach \(21.26\) and \(15.78\), respectively. Thus, improving the reconstruction of the ball region is not sufficient to produce a reliable policy-training simulator. Such a simulator also requires long-horizon consistency, accurate reward prediction, reliable action response, and robustness to the policy-induced state distribution.

6.0.0.2 Evaluation scope.

We focus on Atari Pong and five implementations. This setting is useful because the key concept is visible and the failures can be inspected directly. It is also deliberately controlled: we have not systematically varied larger datasets or more complex tasks.

6.0.0.3 The concepts are manually specified.

CGSReg assumes that the important concept is known and can be mapped to a pixel mask. This is a strong human prior. It is reasonable for the Pong ball, but it does not directly cover concepts such as rules, object relations, or latent game states.

A more general approach should discover task-relevant concepts automatically and ground them in their predictive and control consequences. This direction is related to Sutton’s view of world knowledge as empirically grounded predictions [17] and to Horde, which learns many predictive questions in parallel from sensorimotor experience [18]. Following Sutton’s The Bitter Lesson [19], such a method should rely less on manually specified concepts and benefit from increased data and computation.

6.0.0.4 Weak frozen models, strong Dyna-style agents.

The main unresolved question is why these world models perform poorly as standalone policy-training simulators after freezing, while the corresponding Dyna-style agents can still learn strong Pong policies. 3 shows this gap for the reproduced checkpoints: after freezing the trained world models, we train new policies from scratch inside these frozen models. These policies perform far below the original Dyna-style agents. 5 shows a related problem for the offline-trained checkpoints under the unified zero-shot protocol. Thus, the issue is not tied to one policy-learning protocol: frozen visual world models are poor standalone policy-training simulators, even though Dyna-style MBRL can use them effectively during joint training.

One possible explanation is that continual world-model updating in Dyna-style training lets the policy benefit from a model that is not a strong standalone simulator. In the Dyna-style method, the policy and model are updated together: the policy mainly uses model rollouts near its current behavior, while new real environment transitions are continually added to replay to update the model. In this setting, a world model can improve RL by augmenting the collected real environment data, without becoming a simulator that generalizes well after it is frozen. We tested this by freezing the world model at different points during Dyna-style training and then continuing policy learning (15). The result is mixed: DreamerV3 and DIAMOND show evidence consistent with this explanation, TWISTER partly supports it, Simulus does not, and STORM is non-monotonic. Continual world-model updating is therefore important for some architectures, but it does not fully explain why frozen visual world models remain weak policy-training simulators.

7 Conclusion↩︎

World models are central to model-based reinforcement learning, but strong agent performance does not guarantee that the learned model is a reliable standalone simulator. We directly evaluate five frozen visual world models from reproduced Dyna-style agents in Atari Pong. Their closed-loop rollouts contain visual and dynamical errors, and policies trained from scratch inside these models perform poorly when transferred to the real environment.

We propose CGSReg to target one possible source of these failures: insufficient modeling of task-critical concepts. Under controlled offline world-model training, CGSReg improves both closed-loop rollouts and pixel-space zero-shot MBRL in DreamerV3, DIAMOND, and TWISTER. Its effects are less consistent in Simulus and STORM. These results indicate that concept-focused reconstruction can improve some frozen world models, but does not resolve all world-model bottlenecks.

Future work should generalize CGSReg to broader environments and world-model architectures, and reduce its reliance on manually specified concepts by discovering task-relevant concepts automatically.

8 Exact Quantitative Scores↩︎

9 CGSReg Details and Project Implementations↩︎

Section 3 defines CGSReg as a masked image-space loss. This section provides additional details about loss reduction, multiple concept masks, and project-specific implementations.

9.1 Loss Reduction↩︎

Image losses in visual world models use different reduction conventions. A mean-reduced loss averages over image elements, whereas a sum-reduced loss accumulates their errors. Because CGSReg uses a region-normalized MSE, we scale it to match the reduction convention of the original image loss.

For a single concept mask \(m\), the region-normalized loss is

\[\ell(m) = \frac{ \sum_p m_p(x_p-\hat{x}_p)^2 }{ \sum_p m_p }.\]

If the original image loss uses mean reduction, we directly use

\[L_{\mathrm{CGSReg}}=\ell(m).\]

If the original image loss uses sum reduction, we rescale the region-normalized loss by the number of elements in the full image:

\[L_{\mathrm{CGSReg}} = |\Omega|\ell(m),\]

where \(\Omega\) denotes the full image support. For an image with shape \((H,W,C)\),

\[|\Omega|=HWC.\]

This scaling makes the magnitude of CGSReg comparable to the original image loss under different reduction conventions.

9.2 Multiple Concept Masks↩︎

CGSReg can be applied to multiple concept regions. Let \(\mathcal{M}_{\mathrm{enabled}}\) denote the set of enabled binary masks. We first compute the unscaled aggregate loss

\[L_{\mathrm{CGSReg}}' = \sum_{m\in\mathcal{M}_{\mathrm{enabled}}} \ell(m).\]

Each mask is normalized by its own region size before aggregation. Different concept regions therefore contribute according to their reconstruction errors rather than their spatial sizes.

The reduction matching described above is then applied to the aggregate loss:

\[L_{\mathrm{CGSReg}} = \begin{cases} L_{\mathrm{CGSReg}}', & \text{for mean reduction}, \\[2mm] |\Omega|L_{\mathrm{CGSReg}}', & \text{for sum reduction}. \end{cases}\]

In the main experiments, the ball mask is the only active concept region.

9.3 Implementation in Each Project↩︎

CGSReg is defined on image-space losses and does not require architectural changes. For each project, we apply it to the image prediction or reconstruction objective of the corresponding visual learning module. 4 summarizes the implementation details.

10 SAM2 Segmentation Pipeline↩︎

We augment the 100K-step Atari Pong replay dataset with concept masks generated by a guided SAM2 [16] video-segmentation pipeline. Each replay frame is paired with its corresponding object masks, which are stored as additional dataset fields and used as spatial supervision during world-model training.

Directly applying SAM2 to every replay trajectory would require manually providing object prompts for each trajectory, which is impractical for a large replay dataset. Instead, we create a small set of representative guiding frames. Each guiding frame is an auxiliary frame with manually specified object point prompts. These guiding frames are inserted into replay trajectories and provide reusable prompts for SAM2 video propagation. The same object annotations can therefore be transferred to many replay trajectories without additional manual labeling.

Because applying SAM2 to long trajectories can exceed GPU memory, we process each replay trajectory in smaller temporal chunks. For each chunk, we insert a set of corresponding guiding frames, run SAM2 video segmentation, remove the inserted frames from the output, and align the remaining masks with the original replay frames.

For Atari Pong, we generate three object masks:

• mask1: ball;

• mask2: opponent paddle;

• mask3: player paddle.

The final replay dataset stores the original observations together with their aligned concept masks. The main CGSReg experiments use only mask1.

11 Offline World-Model Training↩︎

The matched offline baseline and CGSReg checkpoints are trained on the same fixed replay dataset containing 100k Atari Pong steps collected from our DIAMOND reproduction. Only world-model components are optimized; policy, actor, and critic updates are disabled.

Because the original projects report training budgets in different units, we align the offline runs by the number of gradient updates applied to world-model components. 5 summarizes the matched budgets and updated components.

For DreamerV3, 100k reproduced environment steps correspond to 25k world-model updates. For DIAMOND, 1000 epochs correspond to 409.6k denoiser updates. For TWISTER, Simulus, and STORM, we directly count world-model updates while excluding policy and critic updates.

12 Unified Pixel-Space Zero-Shot MBRL Protocol↩︎

This section describes the unified zero-shot protocol used in 4.3. For each frozen world-model checkpoint, we train a new policy entirely inside the frozen model and evaluate the resulting policy in the real ALE Pong environment. The world-model parameters remain fixed, and no additional real-environment data are collected.

At the beginning of each episode, the frozen world model receives an initial context sampled from the fixed replay dataset. It then generates future observations, rewards, and terminal signals conditioned on the policy actions. The policy receives only pixel observations: \(64 \times 64 \times 3\) RGB images, without access to latent representations or internal environment states.

We use the same policy-learning algorithm and hyperparameters for all world models. The policy is a discrete-action on-policy actor-critic agent with an image encoder, a categorical policy head, and a value head. Each update collects 15 rollout steps from 64 parallel world-model environments, followed by one optimization step. We train for 20k updates, corresponding to

\[64 \times 15 \times 20000 = 19.2\mathrm{M}\]

world-model environment steps.

The rollout horizon inside each frozen world model is 512. Reaching this horizon is treated as truncation rather than termination. Rewards follow the Atari Pong convention \(-1,0,+1\). For models that produce continuous rewards, we apply a threshold of \(0.1\) before policy learning. After training, we evaluate the learned policy over 20 real-ALE reset seeds with deterministic action selection. Pong returns lie in \([-21,21]\), where higher scores indicate better performance.

13 Ablations↩︎

The main experiments select one \(\lambda_{\mathrm{CGSReg}}\) for each project and use a fixed replay dataset. This section studies two factors that may affect the effectiveness of CGSReg: the regularization strength and the offline replay dataset. All ablation experiments use the unified pixel-space zero-shot MBRL protocol from 4.3.

We sweep \(\lambda_{\mathrm{CGSReg}} \in \{0,0.01,0.1,1.0\}\) on two fixed replay datasets. The main experiments use the dataset collected from our DIAMOND reproduction. The second dataset is collected from our STORM reproduction [6].

The ablation results show that CGSReg improves DreamerV3, DIAMOND, and TWISTER on both replay datasets, although the optimal \(\lambda_{\mathrm{CGSReg}}\) depends on the model and dataset.

14 Representative World-Model Projects↩︎

This section provides an overview of the five representative visual world-model agents evaluated in this work. These projects explore different architectural designs and prediction paradigms, including recurrent latent dynamics models, pixel-space diffusion models [9], Transformer-based latent sequence models [10], and discrete token-based predictive models.

DreamerV3, DIAMOND, TWISTER, and STORM are trained end to end, jointly optimizing visual representation learning and future dynamics prediction. Simulus instead adopts a two-stage framework that first learns a visual tokenizer and then trains a separate token-dynamics model.

The evaluated agents also differ in their prediction spaces. DIAMOND directly models future observations in pixel space using an action-conditioned diffusion model. DreamerV3, TWISTER, and STORM predict future continuous latent states. Simulus performs next-token prediction in a discrete visual token space using a RetNet-based autoregressive model [11]. Despite these differences, all five agents include visual learning objectives that encourage consistency with future observations.

14.1 DreamerV3↩︎

DreamerV3 [2] is a representative latent world-model-based reinforcement learning framework. It uses a recurrent state-space model (RSSM) to model environment dynamics in a compact latent space. The RSSM state contains a deterministic recurrent state and a stochastic latent state represented by discrete categorical variables. Conditioned on previous latent states and actions, the RSSM predicts future latent states, while decoder networks reconstruct observations and predict auxiliary signals such as rewards and continuation probabilities.

DreamerV3 performs reinforcement learning inside the learned world model through latent imagination. Because dynamics prediction occurs in latent space, policy optimization uses latent-state trajectories rather than pixel observations. DreamerV3 stabilizes latent world-model training through design choices including discrete categorical latent variables, KL balancing, and symlog-based prediction targets.

These design choices enable DreamerV3 to perform strongly across diverse reinforcement learning benchmarks. It is competitive on Atari 100K, scales to the full Atari benchmark, and applies to continuous-control tasks such as the DeepMind Control Suite and long-horizon environments such as Minecraft.

14.2 DIAMOND↩︎

DIAMOND [3] is a pixel-space visual model-based reinforcement learning agent that uses a diffusion model for action-conditioned future prediction. Given previous observations and actions, DIAMOND predicts future frames by iteratively denoising noisy images. It therefore models future pixels directly rather than learning an intermediate latent dynamics model. This design enables high-fidelity visual rollouts, strong Atari 100K results, and an interactive neural game engine trained from Counter-Strike: Global Offensive gameplay.

14.3 TWISTER↩︎

TWISTER [4] adopts a latent-space world-model architecture similar to STORM, using a Transformer-based dynamics model to predict action-conditioned future latent states. It follows the Dreamer-style latent-state modeling paradigm while replacing recurrent dynamics with Transformer sequence modeling. Its main contribution is the action-conditioned contrastive predictive coding (AC-CPC) objective, which trains the model to predict and distinguish action-conditioned future latent states and preserve information useful for long-horizon prediction.

14.4 Simulus↩︎

Simulus [5] is the successor to REM [20] and adopts a token-based visual world-modeling paradigm. In REM, the visual world model is trained in two stages: a VQ-VAE tokenizer [21] first represents image observations as discrete visual tokens, after which a RetNet-based autoregressive dynamics model [11] learns next-token prediction conditioned on actions. Simulus extends this framework with multi-modal tokenization and several training improvements. Unlike the other evaluated agents, it models future visual evolution through next-token prediction in a discrete token space.

14.5 STORM↩︎

STORM [6] follows the Dreamer-style latent world-model design, where each latent state contains a deterministic representation and a stochastic latent variable. Instead of the recurrent state-space model used by DreamerV3, STORM employs a Transformer-based sequence model to predict action-conditioned future latent trajectories. Its successor, OC-STORM [22], further incorporates object-centric representations into the framework.

15 Diagnostic Study of World-Model Freezing in Dyna-Style Training↩︎

The strong RL performance of Dyna-style agents contrasts with the poor standalone performance of their frozen world models. Although these agents can achieve high returns during joint training, the corresponding frozen world models often produce incorrect closed-loop rollouts and support weak pixel-space zero-shot MBRL performance.

One possible explanation is that Dyna-style agents benefit from continual world-model updating. During Dyna-style training, the policy and world model are optimized together, while newly collected real-environment data continually update the world model. The policy may therefore benefit from a model that is useful near its current data distribution without requiring the model to become a strong standalone simulator after freezing.

This hypothesis predicts that freezing the world model during Dyna-style training should reduce final RL performance, especially when freezing occurs earlier. To test this hypothesis, we freeze the world model at different training stages and continue the remaining training process. For each project, we follow its native Dyna-style training pipeline. After freezing, real-environment interaction, data collection, policy learning, and evaluation continue, while only world-model updates are disabled.

Let \(T\) denote the nominal reproduction training budget. We compare four settings: no freeze, freeze at \(0.5T\), freeze at \(0.75T\), and freeze at \(1.0T\). All runs continue until \(1.5T\). This diagnostic follows each project’s native training and evaluation protocol and is therefore separate from the unified pixel-space zero-shot MBRL protocol in 4.3.

7 shows the performance trajectories after freezing world-model updates. 8 reports the final real-environment returns.

The results provide partial support for the hypothesis. DreamerV3 and DIAMOND show lower post-freeze trajectories than their no-freeze controls, suggesting that continual world-model updating contributes to their Dyna-style performance. TWISTER shows a weaker effect, while Simulus does not degrade and STORM exhibits non-monotonic behavior. Continual world-model updating can therefore be important for some architectures, but it does not fully explain why frozen world models remain weak standalone policy-learning simulators.

The cause of the frozen-world-model gap remains unclear. Other factors, including robustness to policy shifts, long-horizon consistency, and action-conditioned prediction, may also contribute and remain important directions for future work.

References↩︎

---

1. https://huggingface.co/sb3/ppo-PongNoFrameskip-v4↩︎