Concept-Guided Spatial Regularization for World Models in Atari Pong¶

🕒 Published (v1): 2026-07-16 15:46 UTC · Source: Arxiv · link

Ask a follow-up

Open an assistant pre-loaded with this paper's context.

TL;DR¶

This paper diagnoses frozen visual world models in Atari Pong—finding that all five studied models (DreamerV3, DIAMOND, TWISTER, Simulus, STORM) fail qualitatively and as standalone policy-training simulators—then proposes Concept-Guided Spatial Regularization (CGSReg), an auxiliary masked reconstruction loss on task-critical pixel regions, which improves world-model fidelity in three of the five architectures.

Problem¶

World models in MBRL are evaluated almost exclusively through end-to-end agent performance, masking whether the learned model is a reliable standalone simulator. Strong Dyna-style agents can use weak world models effectively during joint training, so the true fidelity of frozen world models is unknown. No systematic diagnostic existed for isolating and directly evaluating these models, nor a targeted remedy for concept-level modeling failures.

Method¶

Diagnosis: Five reproduced Dyna-style agents (DreamerV3, DIAMOND, TWISTER, Simulus, STORM) are trained to match reported performance; their world models are then frozen. Two evaluations are applied: (1) closed-loop rollout diagnosis—an external PPO policy drives each frozen model for 512-step horizons and generated trajectories are inspected; (2) pixel-space zero-shot MBRL—a new policy is trained entirely inside the frozen model, then evaluated in the real ALE environment.

CGSReg: An auxiliary masked image reconstruction loss is added to the standard world-model objective: $$L_{wm} = L_{img} + \lambda_{\text{CGSReg}} L_{\text{CGSReg}} + L_{\text{nonvisual}}$$ where $$L_{\text{CGSReg}} = \frac{\sum_p m_p (x_p - \hat{x}_p)^2}{\sum_p m_p}$$ with binary mask $m$ selecting pixels of a task-critical concept (the ball, segmented via SAM2). Normalization by mask area makes the loss concept-size-invariant. To isolate CGSReg from Dyna-style co-evolution, all comparisons use offline world-model training on a fixed 100k-step replay dataset with the policy loop disabled.

Key Contributions¶

First direct, controlled evaluation of frozen visual world models from five representative MBRL systems, using both qualitative rollouts and quantitative zero-shot MBRL.
Introduction of pixel-space zero-shot MBRL as a diagnostic: train a fresh policy inside a frozen visual world model, evaluate in the real environment.
Proposal of CGSReg, a concept-region auxiliary reconstruction loss that improves world-model fidelity independent of the downstream policy loop.
Empirical characterization of per-architecture bottlenecks (action-response failure in STORM, tokenizer-vs-dynamics mismatch in Simulus) showing CGSReg addresses only a subset of failure modes.

Results¶

Baseline gap (zero-shot vs. Dyna-style): DreamerV3 drops from $-5.5$ (Dyna-style) to $-20.9$ (frozen WM); DIAMOND from $19.7$ to $-9.6$; TWISTER from $20.8$ to $-13.3$; Simulus from $17.7$ to $-11.6$; STORM from $18.7$ to $-21.0$.
CGSReg unified zero-shot MBRL (offline baseline → CGSReg):
DreamerV3: $-21.00 \to -11.90$
DIAMOND: $-13.90 \to -5.80$
TWISTER: $-21.00 \to -1.90$
Simulus: $-15.80 \to -4.10$
STORM: $-21.00 \to -21.00$ (no change)
Qualitative rollouts improve for DreamerV3, DIAMOND, TWISTER (better ball visibility, plausible rebounds); Simulus shows no clear rollout improvement; STORM's action-response failure persists.

Limitations¶

CGSReg-trained policies still do not solve Pong; mean returns remain well below $+21$, with high variance (TWISTER std $= 21.26$, Simulus std $= 15.78$).
Concepts are manually specified via a human-labeled pixel mask (SAM2 on the ball); does not generalize to latent or relational concepts.
Evaluation is restricted to Atari Pong and five implementations; generalization to more complex environments is untested.
CGSReg improves the image decoder/tokenizer but does not directly regularize token-dynamics models (relevant for Simulus) or action-conditioning (relevant for STORM).
The core question—why Dyna-style joint training succeeds where frozen-model training fails—is only partially resolved; continual model updating explains some but not all architectures.

Relevance to Agentic AI / LLM Agents¶

World models are increasingly considered as internal environment simulators for planning-capable agents; this paper surfaces a critical gap: an agent's world model can be "good enough" for Dyna-style RL without being a reliable standalone simulator, which directly matters for any agent architecture that plans in imagination (e.g., DreamerV3-style or token-prediction-based agents). The pixel-space zero-shot MBRL diagnostic is a transferable evaluation protocol for assessing whether a world model could support offline or zero-shot policy search—relevant to sample-efficient agentic systems. CGSReg's principle—amplifying supervision signal on semantically important regions—connects to broader work on grounded perception and concept-aware representation learning in agents operating over visual observations.