# Learning Cognitive Maps from Transformer Representations for Efficient Planning in Partially Observed Environments

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

## Key facts

- Cited by 6 papers in the map (data through 14 Oct 2026).
- Placed in Latent predictive world models, one of 11 directions in World models.
- Builds on 27 of its 37 references that are in the map.
- First posted on arXiv 11 Jan 2024.

## Placed in

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

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [Training Verifiers to Solve Math Word Problems](https://mlsanity.com/paper/2110.14168/) (2021)
- [Evaluating Large Language Models Trained on Code](https://mlsanity.com/paper/2107.03374/) (2021)
- [PaLM: Scaling Language Modeling with Pathways](https://mlsanity.com/paper/2204.02311/) (2022)
- [Bootstrap your own latent: A new approach to self-supervised Learning](https://mlsanity.com/paper/2006.07733/) (2020)
- [Neural Discrete Representation Learning](https://mlsanity.com/paper/1711.00937/) (2017)
- [Zero-Shot Text-to-Image Generation](https://mlsanity.com/paper/2102.12092/) (2021)
- [Gaussian Error Linear Units (GELUs)](https://mlsanity.com/paper/1606.08415/) (2016)
- [Emergent Abilities of Large Language Models](https://mlsanity.com/paper/2206.07682/) (2022)
- [Beyond the Imitation Game: Quantifying and extrapolating the capabilities of language models](https://mlsanity.com/paper/2206.04615/) (2022)
- [Transformer-XL: Attentive Language Models Beyond a Fixed-Length Context](https://mlsanity.com/paper/1901.02860/) (2019)
- [Decision Transformer: Reinforcement Learning via Sequence Modeling](https://mlsanity.com/paper/2106.01345/) (2021)
- [GAIA: a benchmark for General AI Assistants](https://mlsanity.com/paper/2311.12983/) (2023)
- [An Explanation of In-context Learning as Implicit Bayesian Inference](https://mlsanity.com/paper/2111.02080/) (2021)
- [Progress measures for grokking via mechanistic interpretability](https://mlsanity.com/paper/2301.05217/) (2023)
- [A Generalist Agent](https://mlsanity.com/paper/2205.06175/) (2022)
- [Towards Automated Circuit Discovery for Mechanistic Interpretability](https://mlsanity.com/paper/2304.14997/) (2023)
- [Emergent World Representations: Exploring a Sequence Model Trained on a Synthetic Task](https://mlsanity.com/paper/2210.13382/) (2022)
- [Leveraging Pre-trained Large Language Models to Construct and Utilize World Models for Model-based Task Planning](https://mlsanity.com/paper/2305.14909/) (2023)
- [Does Circuit Analysis Interpretability Scale? Evidence from Multiple Choice Capabilities in Chinchilla](https://mlsanity.com/paper/2307.09458/) (2023)
- [Evaluating Cognitive Maps and Planning in Large Language Models with CogEval](https://mlsanity.com/paper/2309.15129/) (2023)
- [BYOL-Explore: Exploration by Bootstrapped Prediction](https://mlsanity.com/paper/2206.08332/) (2022)
- [Guaranteed Discovery of Control-Endogenous Latent States with Multi-Step Inverse Models](https://mlsanity.com/paper/2207.08229/) (2022)
- [Plansformer: Generating Symbolic Plans using Transformers](https://mlsanity.com/paper/2212.08681/) (2022)
- [Schema-learning and rebinding as mechanisms of in-context learning and emergence](https://mlsanity.com/paper/2307.01201/) (2023)
- [Graph schemas as abstractions for transfer learning, inference, and planning](https://mlsanity.com/paper/2302.07350/) (2023)

## Built on by (most cited first, 6 of 6)

- [Exploring and Learning Structure: Active Inference Approach in Navigational Agents](https://mlsanity.com/paper/2408.05982/) (2024)
- [Learning Dynamic Cognitive Map with Autonomous Navigation](https://mlsanity.com/paper/2411.08447/) (2024)
- [Cognitive Maps in Language Models: A Mechanistic Analysis of Spatial Planning](https://mlsanity.com/paper/2511.13371/) (2025)
- [A Biologically Interpretable Cognitive Architecture for Online Structuring of Episodic Memories into Cognitive Maps](https://mlsanity.com/paper/2510.03286/) (2025)
- [A Close Look At World Model Recovery In Supervised Fine-Tuned LLM Planners](https://mlsanity.com/paper/2606.03685/) (2026)
- [Differentiable Clone-Structured Causal Graphs for End-to-End Cognitive Map Learning from Image Sequences](https://mlsanity.com/paper/2607.12382/) (2026)

## Abstract (arXiv)

Despite their stellar performance on a wide range of tasks, including in-context tasks only revealed during inference, vanilla transformers and variants trained for next-token predictions (a) do not learn an explicit world model of their environment which can be flexibly queried and (b) cannot be used for planning or navigation. In this paper, we consider partially observed environments (POEs), where an agent receives perceptually aliased observations as it navigates, which makes path planning hard. We introduce a transformer with (multiple) discrete bottleneck(s), TDB, whose latent codes learn a compressed representation of the history of observations and actions. After training a TDB to predict the future observation(s) given the history, we extract interpretable cognitive maps of the environment from its active bottleneck(s) indices. These maps are then paired with an external solver to solve (constrained) path planning problems. First, we show that a TDB trained on POEs (a) retains the near perfect predictive performance of a vanilla transformer or an LSTM while (b) solving shortest path problems exponentially faster. Second, a TDB extracts interpretable representations from text datasets, while reaching higher in-context accuracy than vanilla sequence models. Finally, in new POEs, a TDB (a) reaches near-perfect in-context accuracy, (b) learns accurate in-context cognitive maps (c) solves in-context path planning problems.