Research map / Architectures and efficiency / Attention and transformers
Learning Cognitive Maps from Transformer Representations for Efficient Planning in Partially Observed Environments
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes TDB, a transformer with discrete vector-quantized bottlenecks whose active codes form a cognitive map that an external solver uses for shortest-path planning in partially observed, perceptually aliased environments.
Problem
Vanilla transformers and LSTMs trained for next-observation prediction do not learn an explicit, queryable world model and cannot plan efficiently in partially observed environments with perceptual aliasing.
Approach
Add one or more discrete (vector-quantized) bottlenecks on top of a causal transformer, train with multi-step or next-encoding (EMA teacher) prediction objectives, and extract an action-conditioned transition graph from bottleneck indices that an external shortest-path solver can query.
Key results
- On 2D aliased rooms, TDB(S=3, M=1) reaches 99.55% ImpFallback and RatioSP 1.00 while vanilla transformer reaches 29.53% ImpFallback and RatioSP 16.82, with test accuracy ~99%.
- On GINC, TDB(S=1, M=4) achieves higher in-context accuracy than LSTM and vanilla transformer at all context lengths; highest in-context accuracies are around 95%.
- On 3D environments, vanilla sequence models improve over fallback for only 26.60% of problems, whereas multi-bottleneck TDBs reach near-perfect planning performance.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Add next-encoding prediction to the single-step TDB objective (ablation) | TDB (S=1, M=1) · 2D aliased rooms | ImpFallback (%): 6.61~(0.53) -> 99.20~(0.20) | Helped |
| Predict the next 3 observations instead of only the next observation (ablation) | TDB (S=1, M=1) · 2D aliased rooms | ImpFallback (%): 6.61~(0.53) -> 99.55~(0.13) | Helped |
| Use M=4 bottlenecks instead of M=1 (ablation) | TDB (S=3) · 2D aliased rooms | TrainingAbove98: 8300~(284) -> 4000~(0) | Helped |
| Use M=4 bottlenecks instead of M=1 (ablation) | TDB (S=3) · 3D synthetic environments | TrainingAfter1k (%): 28.05~(1.88) -> 99.63~(0.06) | Helped |
| Use M=4 bottlenecks instead of M=1 (ablation) | TDB · GINC | Training convergence steps: 13000 -> 5000 | Helped |
| Replace the standard discrete bottleneck with a clone-structured discrete bottleneck (ablation) | TDB · 3D synthetic environments | Training convergence | Helped |
| Replace the standard discrete bottleneck with a clone-structured discrete bottleneck (ablation) | TDB · In-context room learning | Structure sharing across rooms | Hurt |
| Use vanilla transformer rollouts instead of a learned cognitive map and external solver | Path planning in 2D aliased rooms · 2D aliased rooms | ImpFallback (%): 99.55~(0.13) -> 29.53~(0.75) | Hurt |
| Use a single-step TDB objective without multi-step or next-encoding prediction (ablation) | TDB (S=1, M=1) · 2D aliased rooms | ImpFallback (%): 99.55~(0.13) -> 6.61~(0.53) | Hurt |
Surprising findings
- Single-bottleneck TDBs with the simple next-observation objective fail at path planning because they merge observations with identical neighbors, creating shortcuts in the latent graph.
- Multiple discrete bottlenecks are highly redundant: a logistic regression predicts one bottleneck index from the others with about 85% test accuracy, versus 0.20% if independent.
Limitations
- Only accepts categorical (discrete) observations; continuous images are not handled directly.
- Multiple discrete bottlenecks learn highly redundant, non-disentangled latent spaces.
- Single-bottleneck TDB fails on aliased rooms without the multi-step or next-encoding augmentations.
Experimental setup
Datasets: synthetic 2D aliased rooms aliased cubes simulated 3D environments GINC
Benchmarks: GINC in-context learning test set 2D aliased rooms path planning 3D synthetic environments aliased cubes
Metrics: TestAccu ImpFallback RatioSP NormGED in-context accuracy
Baselines: vanilla transformer LSTM CSCG
Builds on: vector quantization (VQ-VAE) clone-structured causal graphs decision transformer BYOL-style EMA teacher
Closest papers by meaning
- Reasoning in Visual Navigation of End-to-end Trained Agents: A Dynamical Systems Approach (CVPR 2025 · similarity 0.87)
- Sparse Planning in Visual World Models via Cost Gradients (NeurIPS 2026 · similarity 0.86)
- Learning with a Mole: Transferable latent spatial representations for navigation without reconstruction (ICLR 2024 · similarity 0.86)
- Evaluating Vision-Language Models as Evaluators in Path Planning (CVPR 2025 · similarity 0.86)
- World-aware Planning Narratives Enhance Large Vision-Language Model Planner (NeurIPS 2025 · similarity 0.86)
- OneVL: One-Step Latent Reasoning and Planning with Vision-Language Explanation (NeurIPS 2026 · similarity 0.85)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.