ArXiv

Concept-Guided Spatial Regularization for World Models in Atari Pong

Authors
Yukuan Lu, Zaishuo Xia, Weyl Lu...
Categories
cs.AI, cs.LG
arXiv
https://arxiv.org/abs/2607.15142v1
PDF
https://arxiv.org/pdf/2607.15142v1

Brief

The paper diagnoses failures of five visual world models in Atari Pong by freezing reproduced models and testing them with closed-loop rollouts and pixel-space zero-shot MBRL. It finds pervasive visual/dynamical errors and large policy-performance drops (e.g., DreamerV3 from -5.5 to -20.9). The authors propose Concept-Guided Spatial Regularization (CGSReg), which boosts rollouts and zero-shot policy quality for several agents but is not a complete fix.

Why it matters

Reproduced five visual world-model agents (DreamerV3, DIAMOND, TWISTER, Simulus, STORM) on Atari Pong, froze each learned world model, and evaluated them with closed-loop rollouts driven by separately trained policies.

Key details

  • All five frozen models showed clear simulation failures (ball disappearance, incorrect ball motion, invalid ball–paddle interactions). Pixel-space zero-shot MBRL policies trained entirely inside frozen models underperformed their original agents—DreamerV3 mean return dropped from -5.5 to -20.9 (near Pong minimum -21).
  • Introduced Concept-Guided Spatial Regularization (CGSReg), an auxiliary pixel-reconstruction loss applied to segmented concept regions; CGSReg improved closed-loop rollouts and zero-shot MBRL for DreamerV3, DIAMOND, and TWISTER, but effects varied across models and metrics.
Source evidence

Abstract

World models are usually evaluated as components of model-based reinforcement learning (MBRL) systems, while the world models themselves are rarely studied in isolation. We examine five representative visual world-model agents in Atari Pong: DreamerV3, DIAMOND, TWISTER, Simulus, and STORM. After reproducing their training pipelines and matching the reported agent performance, we freeze the learned world models and evaluate them with a closed-loop rollout diagnostic: a policy trained separately from the corresponding MBRL agent interacts with each frozen model, and the generated video trajectories are inspected for visual and dynamical errors. Across all five models, the rollouts contain clear failures, including ball disappearance, incorrect ball motion, and invalid ball-paddle interactions. Beyond visual trajectories, we further evaluate them with pixel-space zero-shot MBRL, where a new policy is trained entirely inside a frozen world model and then evaluated in the real environment. Across all five models, the resulting policies substantially underperform those produced by the corresponding original MBRL training pipelines. The gap is particularly large for DreamerV3, whose mean return drops from -5.5 to -20.9, near the minimum Pong return of -21. We hypothesize that insufficient modeling of task-critical concepts, such as the ball in Pong, may contribute to these failures. We therefore propose Concept-Guided Spatial Regularization (CGSReg), an auxiliary pixel reconstruction loss applied to segmented concept regions. Experiments show that CGSReg improves both closed-loop rollouts and pixel-space zero-shot MBRL in DreamerV3, DIAMOND, and TWISTER. Its effects vary across the remaining models and evaluation metrics, indicating that CGSReg alone does not address all world-model bottlenecks.

Comment: 16 pages, 5 figures