Cognitive Maps and Model-Based Learning

This builds on Animal Learning and Cognition, which treated three associative phenomena — blocking, higher-order conditioning, and delayed reinforcement — as reinforcement-learning mechanisms seen in behavior. None of them needs the animal to hold a model of its world. This part takes up the phenomenon that does: the cognitive map, learned without reward and queried by planning, which is model-based reinforcement learning measured in a rat.

Cognitive maps

Model-based reinforcement-learning algorithms use environment models that have elements in common with what psychologists call cognitive maps. Recall from planning and learning that by an environment model we mean anything an agent can use to predict how its environment will respond to its actions — in terms of state transitions and rewards — and by planning we mean any process that computes a policy from such a model. An environment model has two parts: a state-transition part, encoding the effect of actions on state changes, and a reward-model part, encoding the reward signals expected for each state or state–action pair.¹

Whether animals use environment models, and if so what the models are like and how they are learned, has driven a long argument in animal-learning research. Some researchers challenged the then-prevailing stimulus–response (S–R) view of learning — which corresponds to the simplest model-free way of learning policies — by demonstrating latent learning.²

Latent learning

In the earliest latent-learning experiment, two groups of rats were run in a maze. For the experimental group there was no reward during the first stage of the experiment, but food was suddenly introduced into the goal box at the start of the second stage. For the control group, food was in the goal box throughout both stages. The question was whether the experimental rats would have learned anything during the first, unrewarded stage.³

Although the experimental rats did not appear to learn much during the first, unrewarded stage, as soon as they discovered the food in the second stage they rapidly caught up with the control rats. The conclusion was that during the non-reward period the experimental rats were developing a latent learning of the maze, which they were able to use as soon as reward was introduced.⁴ The rats had learned the maze's layout while it earned them nothing; the reward only revealed a knowledge that was already there.

Latent learning is most closely associated with Edward Tolman, who interpreted results like these as showing that animals learn a cognitive map of the environment in the absence of rewards or penalties, and can use the map later when motivated to reach a goal.⁵ A cognitive map could let a rat plan a route to the goal that is different from the route it used in its initial exploration — the signature of a model you can query for paths you never walked. In modern terms, cognitive maps are not restricted to spatial layouts but are more generally environment models, models of an animal's task space. That a rat's spatial map has a physical substrate is now familiar: the hippocampus contains place cells that fire selectively when the animal is at a particular location, a neural correlate of the map Tolman inferred from behavior alone. The cognitive-map explanation of latent learning is the claim that animals use model-based algorithms, and that environment models can be learned even without explicit rewards or penalties, then used for planning once reward appears.⁶

How the map is learned: system identification

Tolman's account of how animals learn cognitive maps is that they learn stimulus–stimulus (S–S) associations by experiencing successions of stimuli as they explore. In psychology this is called expectancy theory: given the S–S associations, the occurrence of one stimulus generates an expectation about the stimulus to come next.⁷ This is much like what control engineers call system identification — learning a model of a system with unknown dynamics from labeled training examples. In the simplest discrete-time version the training examples are $\text{S}-{\text{S}}^{\prime }$ pairs, where $\text{S}$ is a state and ${\text{S}}^{\prime }$ the subsequent state, and the model learns to expect ${\text{S}}^{\prime }$ after $\text{S}$.⁸

The S–S account is the exact contrast to the stimulus–response view it was raised against, and the contrast is what latent learning turns on. An S–R learner caches, for each state, the action that has paid off there: a lookup table from situation to response, built only when reward stamps a response in. Take away reward and there is nothing to stamp in, so an S–R learner should acquire nothing during an unrewarded stage. An S–S learner instead records what follows what — state to next state — a model of the world that reward never enters. It fills in during exploration whether or not anything is paid, and stands ready to be queried once a goal is supplied. The experimental rats behaved like S–S learners: silent during the unrewarded stage, then immediately competent when food appeared, because the map was already built.

Models useful for planning involve actions too, so examples become $\text{SA}-{\text{S}}^{\prime }$: state ${\text{S}}^{\prime }$ is expected when action $\text{A}$ is taken in state $\text{S}$. And to learn where the rewards are, examples take the form $\text{S}-R$ or $\text{SA}-R$, with $R$ the reward signal associated with the state or the state–action pair. These are all forms of supervised learning, by which an agent can acquire cognitive-like maps whether or not it receives any non-zero reward signal while exploring, and that is how the experimental rats could build their map through a stage that paid them nothing.

Planning over the map: a worked maze

The payoff of holding a model rather than cached values is what happens when the world changes. Consider a rat in a maze with two turns and four distinctive goal boxes, each delivering a fixed reward. From the start state $S_1$ the rat chooses left or right to reach $S_2$ or $S_3$, then chooses left or right again to land in one of the four goal boxes.⁹

A model-free rat solves this with stored action values, one number per state–action pair, learned over many runs of the maze. Once those estimates are good enough, it just picks the largest-valued action at each state. A model-based rat instead learns a transition model (the branching structure of the maze) and a reward model (which goal box pays what), and decides by simulating action sequences to find the path with the highest return.

With rewards $0,4,2,3$ across the four goal boxes, both strategies converge on the same choice: go left to $S_2$, then right, for a return of $4$. The difference shows only under change. Suppose the rat is placed directly in the goal box right of $S_2$ and finds the reward there is now $1$ instead of $4$, without ever running the maze. A model-based rat updates its reward model from that single experience; planning then brings the new value to bear on maze-running with no further experience in the maze, changing the policy to right turns at both $S_1$ and $S_3$ for a return of $3$. A model-free rat cannot do this. To change the action its policy specifies at a state, or the value it has cached for a state, it has to move to that state, act from it, and experience the consequences — possibly many times.¹⁰ This is the mechanism behind the outcome-devaluation result from the previous lesson: the goal-directed rat re-plans over its map the instant the goal's value changes, while the habitual rat must relearn by acting.

Trace the replan explicitly. Before the change, the model-based rat backs values up the tree by taking the max at each choice. The right goal box of $S_2$ pays $4$, so $S_2$'s value is $\max (0,4)=4$; the goal boxes of $S_3$ pay $2$ and $3$, so $S_3$'s value is $\max (2,3)=3$. From $S_1$ the rat compares going left ($4$) with going right ($3$) and turns left, then right, for a return of $4$. Now devalue: the rat is dropped into the right box of $S_2$ and finds it pays $1$, not $4$, and it updates that one reward-model entry. Nothing else in the maze was touched, but the backup changes wholesale. $S_2$ is now worth $\max (0,1)=1$; $S_3$ is untouched at $3$; from $S_1$ the comparison is now left ($1$) versus right ($3$), so the plan flips to right, then right, for a return of $3$.

The model-based rat did no maze-running to produce the new plan; it re-solved the tree from one changed number. The model-free rat, holding only the cached action values $⟨S_1,L⟩=4$, $⟨S_1,R⟩=3$, and so on, has no such move. The value it cached for $⟨S_1,L⟩$ was learned from full runs ending at the box that used to pay $4$; a single visit to that box while not running the maze from $S_1$ gives it no update to $⟨S_1,L⟩$ at all. To fix the stale cache it must physically go left from $S_1$, walk to the now-devalued box, take the disappointing reward, and let the error propagate back — and repeat until $⟨S_1,L⟩$ falls below $⟨S_1,R⟩$. Same maze, same single fact learned; the map-holder acts on it at once, the cache-holder cannot act on it until it re-runs.

The connection runs forward as well as back. The same model-based machinery — a learned transition model, a learned reward model, and planning by simulated rollouts — is what modern model-based reinforcement learning scales up with function approximation, learning the model with neural networks and planning with search or imagined trajectories. Tolman's rat, silently mapping a maze it was not paid to explore, was doing in tissue what a model-based agent does in silicon.

The successor representation and evidence for maps

Sutton and Barto leave the model-free/model-based split as a clean dichotomy — cache versus tree, habit versus plan. Work since has filled in the middle and pinned the split to brain and behavior. Three lines are worth stating carefully, each for exactly what its authors showed.

A predictive map between the two systems. Dayan (1993, Neural Computation) introduced the successor representation: instead of caching a scalar value per state, or holding a full one-step transition model to plan over, cache for each state a vector of expected future occupancy — how often, discounted, each other state will be visited if you follow the current policy from here. Value is then a dot product of this predictive map with the reward vector. The successor representation sits between model-free and model-based control. Like a cached value it is fast to read, needing no search; like a model it factors the world's dynamics apart from reward, so a change in reward re-values states without re-running the policy — though a change in the transitions still requires relearning the map, which a full model would not. It is a partial cognitive map: predictive structure without a plannable model.

Stachenfeld, Botvinick and Gershman (2017, Nature Neuroscience) argued that the hippocampus encodes a map of this kind. They showed that place-cell firing fields, and the periodic fields of entorhinal grid cells, are well described as components of a successor representation of the animal's environment — place fields skew and stretch along frequently-taken routes, and grid patterns emerge as a low-dimensional decomposition of the predictive map. This gives Tolman's cognitive map, inferred from behavior in the 1940s, a concrete computational form and a candidate neural substrate: not a static layout but a predictive map of where the animal is headed.

Both systems in one human choice. Daw, Gershman, Seymour, Dayan and Dolan (2011, Neuron) built a two-step Markov task that separates the two controllers in a single stream of choices. A first-stage choice leads probabilistically to one of two second-stage states, and reward there drifts over time. A pure model-free learner repeats whatever first-stage action preceded reward, ignoring how the transition went; a model-based learner uses its knowledge of the transition structure, so after a rare transition it does the opposite. Human choices showed a mixture of the two signatures, and the balance shifted with circumstance — tilting model-free under cognitive load or time pressure. This is the behavioral fingerprint of the habit/goal-directed distinction this lesson maps onto model-free and model-based control, measured in people rather than inferred. Gläscher, Daw, Dayan and O'Doherty (2010, Neuron) added the learning-signal half: alongside the reward-prediction error that trains values, they found a distinct state-prediction error — the surprise when a transition differs from what the model expected — with its own neural correlate, the error signal that would train the model itself rather than its cached values.

Scaled-up maps in machines. The engineering forward-pointer is real systems that learn a world model and plan inside it. Ha and Schmidhuber (2018, World Models) trained an agent to compress its environment into a learned latent dynamics model and then learn a policy largely by dreaming — rolling the model forward instead of acting. Schrittwieser et al. (2020, Nature), in MuZero, learned a model whose only job is to support value-accurate planning, and combined it with Monte-Carlo tree search to reach top play in Go, chess, shogi, and Atari without being told the rules. These are Tolman's map plus plan built at scale: a learned model of consequences, queried by simulated rollouts, to choose actions in situations never directly experienced. The rat mapping a maze it was not paid to explore and the agent planning over a learned latent model are running the same idea at very different sizes.

The one claim

Three phenomena, one point. Blocking is prediction error: an animal learns about a stimulus only to the extent that its outcome is surprising, which is the error term in every update in this course. Higher-order conditioning and conditioned reinforcement are bootstrapping: a learned prediction becomes a reward in its own right, which is what a value function is and what an actor–critic critic supplies. Delayed reinforcement is credit assignment, and the stimulus traces and lengthened goal gradients that Pavlov and Hull proposed for it are eligibility traces and TD-learned values. Cognitive maps are environment models, learned by system identification with or without reward, and queried by planning — model-based reinforcement learning, measured in a rat.

None of this asserts that the brain runs these exact equations. It asserts something narrower and more durable: the structure of the learning problem is the same whether the learner is tissue or silicon, so a theory built to solve it computationally keeps landing on mechanisms that animal-learning researchers had already isolated in behavior. The next module thread follows the contact point one level deeper, into the neuron, where the dopamine signal turns out to behave like the TD error that all of this has been circling.