GridWM: Grid-Like Spatial Indexing for Long-Horizon World Models

Long-horizon world modeling in partially observable environments requires compact memories that preserve spatial structure despite accumulated localization drift. Inspired by the multi-periodic coding of entorhinal grid cells, we propose GridWM, which adopts a layer-wise organization of local temporal RoPE and multi-periodic spatial RoPE, together with content-conditioned address offsets, to capture recent temporal context and long-range spatial structure. A fixed-capacity memory keeps KV storage and per-step attention cost bounded as the sequence grows. On MemoryMaze, GridWM outperforms strong world-model baselines within the training sequence length and generalizes substantially better to prediction horizons far beyond those seen during training, while using a compact model. The temporal-spatial encoding remains effective across deeper Transformer stacks, with consistent gains as depth increases. These results show that combining local temporal context with structured spatial indexing provides an effective inductive bias for efficient long-horizon world modeling.