Source-linked AI summary
AGI Maze Prediction Datasets: A Compact Benchmark for Learning World Dynamics with Transformers
Alexey Potapov
TL;DR
World modeling requires predictive models to maintain internal state that supports action-conditioned reasoning. This paper introduces a compact maze benchmark and compares textual, auxiliary-memory, and spatial-memory architectures. The spatial-memory model reaches perfect validation accuracy on selected sequence tasks and substantially improves text-trace prediction under matched conditions.
Problem
The paper asks whether predictive models trained from scratch can learn internal representations that support the dynamics of structured environments rather than relying on external or brittle textual state.
Method
The paper introduces procedurally generated maze datasets for per-step, fixed-horizon, and sequential text-observation prediction, evaluated with source-maze-disjoint splits and compared across memory architectures.
Results
The pseudo-video spatial-memory Transformer reaches perfect validation accuracy on selected sequence tasks and substantially improves text-trace prediction, while unstructured auxiliary memory does not consistently improve held-out performance.
Takeaways & Limitations
Structured, task-aligned working memory can be more useful than additional latent capacity alone, motivating textual interfaces complemented by learned working memories.
Takeaways & Limitations
Perfect validation accuracy does not prove that individual latent cells correspond one-to-one with maze cells or objects, and intermediate frames are not trained to reproduce rendered maps.
Abstract
from arXiv · showhide
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.
1. Introduction
The paper frames world modeling as learning persistent internal state for action-conditioned prediction, then introduces a compact maze benchmark to test this capability and structured memory architectures.
- Motivation: World modeling requires representations that preserve current state, update after actions, and support predicting consequences.The motivation connects this capability to model-based reinforcement learning and imagined trajectories.
- Motivation: LLM agents commonly externalize persistent state through prompts, retrieved documents, tools, maps, notes, programs, or databases.Such artifacts are maintained outside the model's forward pass and do not establish an intrinsic internal state.
- Motivation: Textual action-observation histories poorly represent two-dimensional adjacency, obstacles, object locations, inventory, and non-local mechanics.Linear descriptions are costly to update and prone to omission or inconsistency.
- Motivation: The paper hypothesizes that robust world models maintain task-appropriate latent data structures and compute directly over them to ground textual descriptions.The study uses two-dimensional spatial maps as one illustrative structured representation.
- Benchmark: The benchmark provides supervised per-step, fixed-horizon, and sequential text-observation prediction tasks for models trained from random initialization.It targets action-conditioned transitions, multi-step composition, and transfer across related two-dimensional worlds.
- Results: The byte-level Transformer learns elementary transitions, but accuracy and sample efficiency decline as maze size, horizon, and transition complexity increase.The spatial-memory Transformer reaches 100% exact-target accuracy on several tests where the byte baseline remains below 100%.
- Benchmark: Source-maze-disjoint splits and controlled ablations test whether models learn reusable transition rules rather than local correlations.The benchmark varies representation format, recurrence or state slots, supervision targets, data composition, horizon, maze size, and mechanics.
- Contribution: The paper contributes controlled datasets, a benchmark for structured action-conditioned dynamics and transfer, and reproducible byte-level and spatial-memory baselines.It presents the setting as a low-cost diagnostic rather than evidence of general world modeling.
2. Datasets and Task Formulation
The datasets turn procedurally generated maze interaction into three increasingly demanding prediction regimes: local transitions, multi-step state composition, and textual observations grounded in evolving state.
- 2.1 Lightweight Diagnostic Datasets for Predictive Models: AGI Maze Prediction belongs to lightweight diagnostic datasets that test whether models learn rules or procedures rather than training-distribution correlations.Its distinctive focus is an action-conditioned model of a persistent, structured environment.
- 2.2 AGI Maze Prediction Datasets: The release uses procedurally generated ASCII mazes stored as UTF-8 JSON Lines, spanning key and river mechanics across 3x3, 3x4, 4x4, and 4x5 boards.Keys create state-dependent object interactions, while rivers introduce forced movement.
- 2.2 AGI Maze Prediction Datasets: The three formats separate local transition learning, multi-step state composition, and language-level prediction grounded in an evolving world state.The first two use compact map-and-action inputs with exact position targets; the third uses complete trajectories and textual observations.
- 2.2.1 Per-Step Transition Prediction: Per-step records provide a rendered current map and one action, requiring prediction of the resulting zero-based agent coordinate.The task is represented as MAP_t, ACT_t -> POS_t+1.
- 2.2.1 Per-Step Transition Prediction: The per-step collection contains 19,993 examples across five maze families and tests local geometry and mechanics rather than navigation or planning.Random traces yield multiple examples per source maze, with source-maze-disjoint validation sets.
- 2.2.2 Fixed-Horizon Sequence State Prediction: The sequence task is stronger than the per-step case because intermediate maps, positions, and observations are withheld while walls, object interactions, and non-local mechanics must be composed.The per-step task is the h=1 special case without multi-step state-update composition.
- 2.2.3 Sequential Text-Observation Prediction: Text-trace records interleave actions and textual observations, and models predict observation bytes using the map or start description, actions, and prior observations as context.Teacher forcing retains preceding ground-truth observations so later predictions are not confounded by earlier decoding errors.
3. Baseline Models and Initial Experiments
The benchmark evaluates byte-level and memory-augmented Transformers on exact action-conditioned state and text-trace prediction. Byte models learn simple transitions but degrade with scale, horizon, and mechanics, while spatially organized memory improves held-out sequence and trace prediction.
- Evaluation: Greedy exact-target accuracy requires the complete autoregressively generated target span to match ground truth byte-for-byte.An incorrect coordinate or malformed output counts as a failed state prediction.
- Byte-Transformer Baseline: The byte-Transformer uses serialized UTF-8 maps and action blocks, with target-only next-byte cross-entropy for positions or text traces.This tests whether conventional autoregression can learn dynamics without an explicitly structured state representation.
- Byte-Transformer Baseline: The byte baseline solves 3x3-keys perfectly, but larger boards, complex transitions, and longer horizons remain more difficult.Additional data helps 4x5-keys, while 3x4-rivers remains below perfect exact accuracy; fixed-horizon performance also does not uniformly recover perfection.
- Byte-Transformer Baseline: Mixing related maze families increases the 3x4-rivers score by 2%, indicating modest but non-zero transfer.Examples from other maze sizes provide only modest benefit despite identical underlying transition laws.
- Byte-Transformer Baseline: Text-trace accuracy improves with more trajectories but remains below perfect exactness for the byte baseline.With 4,000 trajectories and rollout depth eight, txt_span_exact_accuracy=0.9809 and trace_all_txt_exact_accuracy=0.869.
- Auxiliary Latent-Memory Transformer: The auxiliary latent-memory model can fit training data perfectly yet systematically underperforms the byte baseline on held-out prediction.It reaches training accuracy 1.000 on 3x3-keys but validation accuracy 0.975 versus 0.979 for the byte model, and 0.886 versus 0.951 on 4x4-keys.
- Pseudo-Video Spatial-Memory Transformer: The pseudo-video model initializes a persistent two-dimensional latent frame from the map and updates it after each action using text-derived write contexts.Its spatial arrangement matches the maze geometry and is the principal inductive-bias difference from unstructured auxiliary memory.
- Pseudo-Video Spatial-Memory Transformer: The visual-only model reaches perfect validation accuracy on 3x3-keys four-step and 4x4-keys two-step tasks, while improving 3x4-rivers to 0.973.The byte and auxiliary baselines are not close to 1.000 on the perfect tasks; for 3x4-rivers they reach 0.881 and 0.868, respectively.
Conclusion
The benchmark provides a controlled way to test whether predictive models learn transferable, action-conditioned world dynamics and structured internal state. Results favor task-aligned spatial working memory over unstructured latent capacity, while the benchmark remains a prerequisite-level test rather than sufficient evidence of general world modeling.
- Conclusion: The benchmark converts stateful grid-world interaction into supervised per-step, fixed-horizon, and sequential text-observation prediction.Its procedurally generated, source-maze-disjoint splits and exact-match metrics test transfer to unseen layouts.
- Conclusion: A byte Transformer can learn simpler regimes, but an unstructured auxiliary latent stream does not consistently improve held-out generalization.The auxiliary model can fit training data perfectly while generalizing less well than the byte baseline.
- Conclusion: The pseudo-video spatial-memory Transformer reaches perfect validation accuracy on selected sequence tasks and substantially improves text-trace prediction under the same data and targets.It receives no intermediate maps or state labels, and its results suggest that a dynamically updated two-dimensional workspace can support held-out prediction.
- Conclusion: The paper hypothesizes that language grounding depends on persistent task-appropriate data structures and computations over them.In this setting, a latent two-dimensional visual memory supplies a spatial-state substrate while text carries actions, observations, and predictions.
- Conclusion: Success on these small environments is not sufficient for general world modeling, but exact action-conditioned prediction provides a measurable prerequisite for studying internal state.The authors propose extensions including longer horizons, broader mechanics, partial observability, and probes or interventions on latent memory.