Giving Transformers a Cognitive Map for Smarter Planning

A breakthrough in transformer modeling?

To keep track of visited locations, building a cognitive map is useful. But how can this be tacked in a transformer architecture? Image created by author, Source code: GitHub

Transformers have revolutionized AI by excelling at pattern recognition and sequence modeling, from language to images. However, vanilla transformers lack an explicit “world model” — they don’t internally build a map of their environment that can be queried for planning routes or navigation​. This becomes a problem when an AI agent needs to plan its actions, like finding the best path to a goal. Without a mental model of the world, a transformer-based agent is essentially “blind” beyond the next prediction, making long-term decision-making hard.

In contrast, humans and animals navigate by constructing cognitive maps — internal representations of the world that allow flexible planning. Endowing AI systems with similar cognitive maps could let them plan routes, imagine outcomes, and adapt to new environments on the fly. This is especially important in partially observed environments (POEs), where the agent only sees limited, ambiguous observations. In such settings, different locations can look identical (perceptual aliasing), so the agent must remember where it’s been to infer where it actually is. Consider a maze with many similar corridors: an agent’s camera view might not tell it which corridor it’s in unless it has built a memory of the path taken. Path…