# AGI Maze Prediction Datasets: A Compact Benchmark for Learning World Dynamics with Transformers

arXiv: https://arxiv.org/abs/2609.02339 · HTML: https://mlsanity.com/paper/2609.02339/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Latent predictive world models, one of 11 directions in World models.
- Builds on 8 of its 10 references that are in the map.
- First posted on arXiv 2 Sep 2026.

## Placed in

- World models / [Latent predictive world models](https://mlsanity.com/wm/latent-predictive-world-models/)

## Builds on (most cited first, 8 of 8)

- [RT-2: Vision-Language-Action Models Transfer Web Knowledge to Robotic Control](https://mlsanity.com/paper/2307.15818/) (2023)
- [Grokking: Generalization Beyond Overfitting on Small Algorithmic Datasets](https://mlsanity.com/paper/2201.02177/) (2022)
- [LeWorldModel: Stable End-to-End Joint-Embedding Predictive Architecture from Pixels](https://mlsanity.com/paper/2603.19312/) (2026)
- [Neural Networks and the Chomsky Hierarchy](https://mlsanity.com/paper/2207.02098/) (2022)
- [Contrastive Learning of Structured World Models](https://mlsanity.com/paper/1911.12247/) (2019)
- [Length Generalization in Arithmetic Transformers](https://mlsanity.com/paper/2306.15400/) (2023)
- [Executable World Models for ARC-AGI-3 in the Era of Coding Agents](https://mlsanity.com/paper/2605.05138/) (2026)
- [AGI Maze as a Benchmark Framework for World-Modeling Agents](https://mlsanity.com/paper/2607.00627/) (2026)

## Abstract (arXiv)

World modeling requires a predictive model to maintain and update an internal state adequate for reasoning about the consequences of actions. We introduce the AGI Maze Prediction Datasets and Benchmark, a lightweight controlled testbed for studying this capability in Transformers and other predictive models. Derived from procedurally generated, stateful grid worlds, the benchmark comprises per-step transition prediction, fixed-horizon state prediction, and sequential textual-observation prediction. Source-maze-disjoint training and validation splits, together with greedy exact-match evaluation, distinguish learning transferable action-conditioned dynamics from memorizing transitions in familiar layouts. We establish from-scratch byte-level Transformer baselines and compare them with two working-memory-augmented architectures. A generic auxiliary latent-memory Transformer can fit some training sets perfectly but does not consistently improve held-out performance. In contrast, a pseudo-video spatial-memory Transformer initializes a two-dimensional latent workspace from the input map and updates it from action history without receiving intermediate maps, positions, or state labels. Under the same data, objectives, and evaluation protocol, this model reaches perfect validation accuracy on selected fixed-horizon tasks where the byte and unstructured-memory baselines do not, and substantially improves sequential text-trace prediction. These results suggest that structured, task-aligned working memory can be more useful than additional latent capacity alone. More broadly, we argue that language grounding is mediated by persistent data structures and computations over them; the benchmark offers a compact setting for testing architectures that couple textual interfaces to learned structured state.