Neural sampling from cognitive maps enables goal-directed imagination and planning

Details of our model for spatial goal-directed imagination

Generating a cognitive map based on grid cells

We used a simple cognitive map model (Fig. 1a) based on the functioning of the grid-cell system in the entorhinal cortex. Furthermore, in our applications, we derive place cells by combining grid-firing patterns.

Figure 1 shows that these experimental data can be reproduced with a standard model for the 2D spatial map that is provided by the grid-cell system of the rodent^6,7,8. Furthermore, an inverse model that maps differences between grid states to movement commands that reduce this difference can be learnt by a simple local rule for synaptic plasticity. By adding noise to this action-selection mechanism, the cognitive map becomes a generative model that produces goal-directed movement plans qualitatively resembling those decoded from the rodent brain¹⁸, even rerouting around novel barriers, as demonstrated in the neural recordings of ref. ²¹. Furthermore, the samples generated by the model generalize to never-experienced trajectories, as observed in rodents^20,36.

The model shown in Fig. 1a consists of grid cells, place cells and action basis neurons. All neurons were linear, with each neuron’s activation equal to the sum of its inputs. Our model is based on the core idea that grid cells can perform ideal path integration based on a velocity input^53,54. We regard grid-cell activity itself as the state space of our model, in which both target state s* and its estimated current state \({\widehat{{\bf{s}}}}_{t}\) are expressed (Fig. 1b). We approximate post-training grid-cell activity in their model analytically, by a deterministic function \({\bf{x}}\,\mapsto \,{\bf{s}}\in {{\mathbb{R}}}^{1,000}\) that outputs the firing rates of 1,000 grid cells at any location in a 2D arena. Each grid cell is initialized independently and modelled as the sum of three positive planar cosine gratings \(r({\bf{x}})={\sum }_{j=1}^{3}(1+\cos ({{\bf{k}}}_{j}\,\cdot \,{\bf{x}}+{\phi }_{j}))\), whose wavevectors k_j are 60° apart and share magnitude ∣k_j∣ = 2π/λ. The grid scale was drawn once per cell from λ ≈ U(0.05, 8) and each phase independently from ϕ_j ≈ U(0, 2π). To eliminate directional bias, a global rotation θ₀ ≈ U(0, 2π) was applied, giving \(\arg ({{\bf{k}}}_{j})={\theta }_{0}+\mathrm{j\pi }/3\).

Under this setting, PCA applied to the grid-cell representations of the 10% largest-scale (λ) grid cells revealed a nearly planar 2D manifold (Fig. 1b). Because larger spatial periods generate progressively flatter interference patterns, this low-curvature embedding is well suited for goal-directed navigation.

Grid-cell firing in the brain is updated by path integration of locomotion activity. A velocity signal is provided in our model by the activity of four action basis neurons, whose sum provides a continuous velocity vector. We constrained a fixed step size η_a = 0.05 by normalizing the length of the resulting velocity vector always to η_a.

To simulate continuous movement, time was discretized, with each time step representing 10 ms. The resultant movement (Δx, Δy) during each time step is computed by the cumulative contributions of all four action basis neurons:

This provides the forward model for the cognitive map in this section.

As in ref. ⁵⁴, grid-cell activity can predict place-cell activity. In our model, the place fields are established by one-shot learning from grid-cell activity. The incoming weight of each place cell is assigned in a single update using a simplified rule that captures the core idea of the LTP component of behavioural time-scale synaptic plasticity^55,56.

During 500 s of random exploration, a plateau-potential gating signal occurs every 500 ms. At each event, one postsynaptic neuron i is randomly selected from the neuron pool of 1,000 neurons, and its incoming weight vector P_i (row i of P) is overwritten with the current presynaptic grid-cell activity pattern s_t:

This centres the place field of neuron i at the location (x_t, y_t). Because the learning happens only at visited places, comprehensive exploration is essential; regions that are never visited cannot recruit the corresponding place fields. We used grid-cell activity, rather than place-cell firing, for generating imagined trajectories, particularly for representing the current state of an imagined trajectory. Hence, this imagination process was not affected by the lack of place cells for the unexplored region shown in Fig. 1c. Our motivation for leaving part of the 2D arena unexplored was to demonstrate that the unexplored region does not hinder grid cells from generating imagined trajectories that include segments in the unexplored region. However, to be able to plot the spiking activity of place cells during imagined trajectories (Fig. 1e), we assumed that place cells had formed in the meantime also for the initially unexplored region.

Using the cognitive map for generating imagined paths

An inverse model in the form of a weight matrix W learned in a self-supervised manner to infer actions from differences in grid-cell states which caused these differences, using the same simple Hebbian learning rule in equation (14).

We simulated here random exploration for 500 s, discretized by 10 ms within a 2D bounded environment ([0, 4] × [0, 4]), excluding the lower-left quadrant (x < 2, y < 2) shown in Fig. 1c. The agent started within the permitted region and executed a continuous-time random walk at constant speed (0.5 unit per second) with Gaussian turn noise. Invalid movements (entering forbidden zones or crossing boundaries) triggered resampling of the direction. We consistently use the matrix W as learned inverse model for a cognitive map throughout this paper.

We used this inverse model for (virtual) action selection during the stochastic generation of an imagined trajectories to a given goal s* as follows:

where ϵ is the noise that introduced stochasticity into the model, and η_a = 0.05 is the step size. Each dimension of ϵ was sampled from a Gaussian distribution \({\epsilon }_{i} \sim {N}(0,{\sigma }_{{\rm{n}}})\), where σ_n represents the noise level depicted in Fig. 1d. We assumed four basis action neurons defining an action vector \({{\bf{a}}}_{t}=[{({{\bf{a}}}_{t})}_{{\rm{up}}},{({{\bf{a}}}_{t})}_{{\rm{down}}},{({{\bf{a}}}_{t})}_{{\rm{left}}},{({{\bf{a}}}_{t})}_{{\rm{right}}}]\).

Spike generation (Fig. 1e) was modelled by interpreting the continuous activation of each place cell i as a firing probability p_i(t) within a discrete time window of 10 ms. Spikes were then sampled from a Bernoulli distribution:

The firing probability at each time step was computed by applying the softmax function to the normalized activation ℓ_i(t) of all place cells:

Obstacle avoidance in imagined trajectories

For generating imagined trajectories to a given goal in the presence of obstacles (Fig. 1f), we augmented equations (3) and (1) with additional repulsive forces originating from each obstacle (x_o, y_o) towards the agent (x, y). The information encoded by the vector difference, δ = (x_o − x, y_o − y), may be represented, in the brain, by object vector cells⁵⁷. Each wall was discretized into three point obstacles (the center and two end-points), each exerting a repulsive force f_i. The total repulsive force F was the vector sum of these:

The denominator \(| {{\boldsymbol{\delta }}}_{i}{| }_{2}^{2}\) ensures that the strength of each repulsive force diminishes with increasing obstacle distance. The total force F, scaled by factor γ, modified the movement of the agent during one simulation time step (10 ms), from equation (1) to

For all four navigation tasks in the presence of obstacles in Fig. 1f, we set γ = 0.2 and noise scale σ_n = 0.3.

This provides a substantially more parsimonious model for planning 2D navigation in the presence of obstacles than the model of ref. ⁵⁸, where first, the outer product of grid and object vector cells has to be created through an offline learning process. Our model is much simpler, and predicts that no new map has to be learnt when an obstacle is moved.

Details of the CML algorithm

The CML algorithm²⁶ maps observations and actions into the same high-dimensional space. Specifically, each observation o_t is mapped using a learned embedding matrix \({\bf{Q}}\in {{\mathbb{R}}}^{{n}_{{\rm{s}}}\times {n}_{{\rm{o}}}}\):

where n_s is the dimensionality of the state space, and n_o is the dimensionality of the observation o_t. The target state (goal) is obtained as the embedding of a target observation:

Similarly, all possible actions a are embedded into the same high-dimensional space via a learned embedding matrix \({\bf{V}}\in {{\mathbb{R}}}^{{n}_{{\rm{s}}}\times {n}_{{\rm{a}}}}\), where n_a is the dimensionality of the action space. The objective of learning these embeddings is ensuring that the predicted next state \({\widehat{{\bf{s}}}}_{t+1}\) is given by

Learning aims at ensuring that \({\widehat{{\bf{s}}}}_{t+1}\) becomes a good approximation of the actual next state s_t+1 = Qo_t+1:

The prediction error \(({{\bf{s}}}_{t+1}-{\widehat{{\bf{s}}}}_{t+1})\), which the learning rules for Q and V aim to minimize (Fig. 2c), therefore, has the following form:

where η_v and η_q are the learning rates. These plasticity rules, known as delta rules, approximate gradient descent to minimize the prediction error. If this learning process is successful and \({\widehat{{\bf{s}}}}_{t+1}\) becomes a good approximation of the actual next state s_t+1, one has a basis for generating a sequence of virtual actions in a goal-directed manner, without receiving observations (Fig. 2e).

In our implementation, the CML additionally learns an inverse model during exploration in the form of a matrix \({\bf{W}}\in {{\mathbb{R}}}^{{n}_{{\rm{a}}}\times {n}_{{\rm{s}}}}\). This matrix is acquired through Hebbian learning and maps state differences s_t+1 − s_t onto the actions a_t that caused them (Fig. 2d):

where η_w is the learning rate. To every possible action a_t, W assigns a utility that scores its estimated usefulness (or value) for reaching s* from the current state s_t:

Each dimension stores the utility of a specific action, forming a utility vector u_t of length n_a. The estimation of action utility is referred to as principle 2 in CML²⁶.

This mechanism relies on the hypothesis that the geometry of the learnt cognitive map provides a sense of direction when choosing the next action. Although W is trained only on local state differences s_t+1 − s_t, it is applied during planning to imaginary state differences s* − s_t that can be substantially larger. Linearity of the inverse model supports this generalization. These utility estimates are analogous to value estimates in RL^23,59. In contrast to RL, however, the learned matrix W provides a universal value function that assigns values for reaching any possible state, and these estimates do not depend on an a priori chosen policy.

As in RL, utility estimates must be modulated by an affordance factor that vetoes actions that cannot be executed in the current state, or that increases or decreases the utility of actions depending on their difficulty in more complex environments. The utility values are masked by the learned affordance factor \({\widehat{{\bf{g}}}}_{t}\in {{\mathbb{R}}}^{{n}_{{\rm{a}}}}\), which aims to approximate the real affordance g_t defined by the constraints of the environment, a one-hot vector that indicates which actions are executable (1 for executable and 0 otherwise). This gating factor prevents the GCML from proposing infeasible actions. We denote the product of the utility of an action (for reaching s*) with this affordance factor as its current eligibility. Action selection at time t then consists of choosing the action a_t with the highest eligibility.

After executing this action, the CML transitions to the next state s_t+1, represented by the embedding of the next observation o_t+1 (s_t+1 = Qo_t+1). From this new state, the same action-selection procedure is iterated until the goal state is reached.

After learning Q, V and W, the CML can be tasked with reaching an arbitrary goal state s* = Qo* that results from a desired observation o*. Such a target o* may correspond to a spatial goal location or, more generally, to a problem to solve in an abstract state space (Fig. 2a and ref. ²⁶). The learned matrix W is then used to select actions that are likely to be useful for reaching this goal state s*.

Mathematical description of the GCML and its relation to biology and other theoretical models

Cognitive maps for non-spatial task domains

Experimental data from neuroscience and cognitive science suggest that the brain—especially the human brain—uses cognitive maps not only for spatial navigation but also for navigation in more abstract concept spaces in which spatial locations are replaced, for example, by images or by combinations of learnt ranks of an item in different linear orders^{2,12,13,14,15,16}. We show here that the principles that enable according to the preceding section goal-directed imagination for 2D spatial cognitive maps generalize to cognitive maps for abstract concept spaces, and enable generic problem-solving (Fig. 2a shows an illustration of a non-spatial problem-solving task).

The role of the grid-cell system in supporting the generation of cognitive maps and conceptual navigation in abstract concept spaces is now well documented across a growing body of studies, including non-invasive neuroimaging experiments in humans (functional magnetic resonance imaging/magnetoencephalography^60,61,62), functional magnetic resonance imaging studies in monkeys⁶³ and intracranial recordings in monkeys⁶⁴. However, the computational mechanisms that would allow grid-cell systems to scale to abstract spaces—particularly those that are not well captured by 2D Euclidean geometry and may require higher-dimensional representations—remain more debated⁷.

We used distinct maps to navigate physical space and abstract problem spaces. However, it is important to note that our approach to ‘sampling from cognitive maps’ is general and can be applied across both physical and abstract spaces, provided that the cognitive map is adapted to the task and its dimensionality. Spatial navigation and foraging occur in approximately Euclidean spaces, requiring simple 2D maps. Such low-dimensional maps can also potentially be used to address simple navigation and foraging tasks in 2D abstract spaces, such as the one studied in ref. ⁶⁰.

However, 2D maps are insufficient for planning in non-planar random graphs and for our compositional problem-solving tasks. Hence, a system like the GCML, capable of learning in high-dimensional spaces, is useful. The use of distinct maps in our tasks is, therefore, not motivated by the difference between physical and abstract spaces, but by the dimensionality of the problems.

Crucially, by highlighting the differences between learning maps of simpler (Euclidean and 2D) versus more complex (higher-dimensional) spaces, we are not ruling out the possibility that the grid-cell system also generalizes to maps in higher dimensions required for challenging abstract tasks. For example, ref. ⁷ proposes that the grid-cell system can generalize to non-Euclidean environments with complex transition structures, using methods called ‘mixed modular coding’ or ‘map fragmentation’ (refs. ^5,65,66 discuss other approaches).

By providing a general approach to sampling from cognitive maps, our method is agnostic as to whether the maps are learned through the grid-cell system (an idea that is increasingly empirically supported) or with the aid of other neural systems.

Sampling from cognitive maps for non-spatial task domains by the GCML

From a biological perspective, the mechanism used by the GCML—the stochastic generation of multiple paths and their subsequent selection based on expected reward—could be associated with the neural circuit formed by the hippocampus and the ventral striatum. Various studies suggest that this circuit supports goal-directed navigational planning and ‘vicarious trial-and-error’, with the hippocampus generating candidate future trajectories and the ventral striatum assessing the reward associated with those trajectories^67,68,69,70.

The GCML uses the same cognitive maps as those that are learned in a self-supervised manner by the CML presented in ref. ²⁶. However, the agent cannot directly access the current state during imagination from the observation. Therefore, the imagined state \({\widehat{{\bf{s}}}}_{t}\) is used as a substitute for the actual state s_t when computing the utility:

The computation of the imagined state will be detailed in the following paragraphs.

Another difference is that GCML cannot get the environment input on affordance, but the agent needs to learn and imagine it. The affordance in GCML is computed using the affordance gating matrix \({\bf{G}}\in {{\mathbb{R}}}^{{n}_{{\rm{a}}}\times {n}_{{\rm{s}}}}\) that receives the (imagined) state vector as its input. The matrix is updated during the learning process according to the following rule:

where η_g represents the learning rate.

The resulting vector, after applying the affordance gating, is termed the eligibility:

where ϵ is the noise is a crucial term in GCML that turns the CML into a probabilistic generative model. This injected noise enables GCML to explore actions beyond the one with maximum eligibility. Hence, in contrast to CML’s deterministic selection, GCML stochastically generates diverse trajectories.

Finally, the selected action a_t is the one with the highest eligibility:

where WTA(⋅) denotes the winner-take-all operator, selecting the action with the highest eligibility from all possible actions. It can be approximated by a simple neural circuit with lateral inhibition⁷¹.

After selecting a virtual action a_t at the estimated state \({\widehat{{\bf{s}}}}_{t}\), the estimated next state, t + 1, is calculated using the bootstrapping process:

By iteratively applying the trajectory generation process described earlier, a virtual goal-directed trajectory is generated.

By repeatedly generating trajectories in the presence of inherent noise in equation (18), different trajectories are generated. This results in a set of possible solutions from the initial state s₀ to the goal s*.

A downstream network could select a specific one from this repertoire of imagined trajectories based on a given criterion. For example, when rewards are associated with specific states, the estimated rewards \({\widehat{r}}_{t}\) of each state can be accumulated over an entire imagined trajectory, and the one with the largest sum of predicted rewards can be selected (Fig. 3f).

GCML suggests a model for neural sampling

On a more abstract algorithmic level, the GCML provides an intriguing paradigm for using the inherent noise of biological and physical computing systems as a computational resource⁷². An important feature is that this noise enables goal-directed sampling, for a goal that is provided in the form of a synaptic input, rather than just sampling from a fixed probability distribution that has been programmed into the weights of synaptic connections. This endows the neural sampling method with a brain-like flexibility in adjusting to new goals or contingencies, and to solve new problems that were never encountered before. From a more general theoretical perspective, our model for goal-directed sampling is related to previous models for probabilistic inference in the brain for perception and action selection through sampling^{33,52,73,74,75,76}. However, in contrast to this preceding work, it provides a neural network paradigm for sampling from a marginal distribution, relative to problem specifications that are defined through the current synaptic input to the network. A theoretical basis for that had been provided in ref. ⁴⁵. It was shown there that a stationary distribution exists for this sampling from a marginal distribution even if the synaptic connections of the network are not symmetric, which is always the case if one has synaptic connections between excitatory and inhibitory neurons. It also exists if the neurons are spiking, and exhibit a brain-like diversity of firing properties.

From a biological perspective, our approach is supported by experimental evidence, showing a dual role of (hippocampal) cognitive maps, for learning states in an associative mode and supporting sequential activity in a predictive mode³⁰—and provides a possible mechanistic explanation of this finding. Novel empirical predictions can be addressed in future studies. In particular, our model proposes that together with each cognitive map, the brain also learns an inverse model that maps state differences to possible actions for reducing this state difference. Note, however, that an inverse model for a cognitive map could also be implemented in the brain in a somewhat different way, where the current and goal states are provided separately as inputs, instead of their difference. However, as long as the inverse model is approximately linear, it supports local actions with ‘foresight’, preferring actions that move into the direction of the goal.

Predictions of model tests through biological experiments

Future studies could align this proposal with empirical findings showing that goal and reward information shape trajectory representations and the geometry of cognitive maps in the hippocampus and cortex^{77,78,79,80,81,82,83,84}. The assumption that goal-directed behaviour is driven by minimizing the discrepancy between current and goal states is shared with frameworks such as cybernetics^85,86,87 and active inference⁵².

A key prediction of the GCML model for future biological experiments is that noise in the selection of possible actions during imagination is directly related to the resulting diversity of action plans, and implicitly also on the difficulty of solving a given task. This could be tested experimentally.

Furthermore, our model predicts that the speed of the generation of the next step of a multistep solution plan depends only mildly on the current distance to the goal. The reason is that the next step is generated by a feedback circuit (Fig. 2e) whose processing speed is independent of the distance to the goal. However, the generation speed of the next step is predicted to grow with the number of action alternatives that are available, since this is likely to affect the speed of the underlying WTA computation for each action selection.

Finally, our model predicts that current goals can be decoded from brain areas that are involved in action selection, as already partially shown in ref. ⁸³. By contrast, RL models predict that the values of adjacent states are the most relevant for action selection, and the identity of the current goal has no direct influence.

GCMLs and RL

Sampling from a marginal distribution provides an alternative to RL for generating goal-directed behaviour. Models based on RL often require methods from machine learning for training, such as backpropagation and backpropagation through time, whose biological viability is debated. Furthermore, except for methods based on successor representations, RL approaches require retraining for each new goal that is to be reached. The approach of ref. ⁶⁵ used the successor representation for RL in combination with a cognitive map. We are not aware of efforts to sample from this data structure. It is also an open problem whether this approach can be applied to tasks such as the compositional tasks that we consider, where planning needs to visit states that were not encountered during exploration.

This also holds for single-goal-conditioned contrastive RL, which was recently proposed^88,89. However, this approaches focuses, like the CML and GCML, on learning state representations that support fast action selection. In fact, actions are selected there also through a WTA operation applied to a dot product of vectors, one of which represents the goal. The focus of single-goal-conditioned contrastive RL is on producing efficient exploration strategies, however, for a fixed goal. An open question is whether this method could also be used to enhance goal-invariant exploration in cognitive-map-based approaches.

Apart from the versatility advantage of cognitive-map-based methods, there are also some recent experimental data that suggest that dopamine signals support the learning of various types of prediction in the brain^90,91,92, not only reward predictions. Hence, they might also support the learning of cognitive maps by minimizing prediction errors for sensory inputs, which provides the basis for the CML and GCML model. Altogether, the diversity of data on dopamine signals suggests that the brain could use not just one but a multitude of mechanisms and strategies for learning goal-directed behaviour, as suggested by ref. ⁹¹.

GCMLs and other theoretical models

There are various models for how the brain could learn cognitive maps, with a particular emphasis on the hippocampal–entorhinal system. Some approaches, such as Tolman–Eichenbaum machine⁶⁶ and others⁹³, emphasize a separation between learning the general structure of the environment (putatively via a grid code)—for example, the transitions between spatial locations, which generalize across different mazes—and learning the content of specific locations (putatively via place cells), which is maze specific. Other models dispense with this distinction^3,5 but are still able to learn, offline, the transition structure of complex, non-Euclidean spaces.

The GCML, by contrast, can learn cognitive maps in high-dimensional spaces through a simple online, predictive learning rule. A unique feature of the GCML (and its predecessor, the CML) is that the learned embeddings provide a sense of direction toward goals—a capability not explicitly addressed in the aforementioned methods or related approaches.

Model predictive control (MPC) is another already-existing framework for action selection based on a learned dynamics model. Compared with the GCML, it has the advantage that it can also handle continuous state and action spaces. Apparently, however, it lacks the capability of the GCML to choose goal-directed actions with short latency, which result from the embedding of states and actions into a high-dimensional space. A nice challenge for future research will be to combine the best of both approaches in a hybrid model.

Further details for the application of GCML for goal-directed imagination for generic problem-solving

The task setup was a random graph with 32 nodes, where the degree of each node was uniformly sampled between 2 and 5. Observations o_t corresponded to nodes, and actions a_t represented traversing edges in a specific direction. Hence, each edge is modelled by two possible actions. Both o_t and a_t were encoded as unique one-hot vectors. The training dataset consisted of 200 randomly generated trajectories, each of length 32. During the planning process, GCML samples a set of trajectories from a given starting node to a target node.

The mappings G transform the estimated high-dimensional state s_t into the corresponding affordance factors. The affordance gating vector is a binary vector indicating available actions, where ones represent feasible actions and zeros indicate restricted actions. A linear transformation is sufficient in this setting, as the vectors s_t are approximately orthonormal.

Exploration was regulated by adding a Gaussian noise vector ϵ—with each dimension sampled from \({\mathcal{N}}(0,{\alpha }_{\epsilon })\)—to a normalized unit-length utility vector. The noise amplitude α_ϵ was set to a default value of 0.1.

Initial values for model parameters were drawn from Gaussian distributions: \({\bf{Q}} \sim {\mathcal{N}}(0,1)\) and \({\bf{V}},{\bf{W}},{\bf{G}} \sim {\mathcal{N}}(0,0.1)\). Smaller initial values for V improved performance. Learning rates were set as η_q = 0.1 and η_v = η_w = η_g = 0.01. GCML demonstrates robustness to variations in learning rates, similar to the original CML, with stable training performance for η_g in the range of [0.001, 0.05].

Each node was assigned a reward drawn from a uniform distribution \({\mathcal{U}}(-5,5)\).

To demonstrate that GCML provides a good approximation of the top-k-shortest-path problem (Fig. 3c,d), 40 trajectories were generated from a given starting node to a goal node under different noise levels in the GCML. The number of distinct trajectories for each possible trajectory length (that is, the number of edges in the path from the starting node to the goal node) was computed using a breadth-first search algorithm. For each trajectory length, the proportion of distinct trajectories generated by a given noise-scale GCML relative to the total number of possible distinct trajectories serves as a measure of how well GCML approximates the top-k-shortest-path algorithm.

Details of goal-directed sampling in compositional task domains

This domain is arguably more transparent than natural language, but it still exhibits a core feature of compositional computing, its computational complexity. In fact, this decomposition task is NP-hard, which implies that every known deterministic algorithm requires computational resources that grow exponentially with the problem size^40,47. Recent experimental data suggest that the human brain solves it by addressing possible BBs in a sequential manner²². Their recordings from the human brain showed that during planning, participants considered the possible BBs that could be removed from the silhouette in a sequential manner, starting with the most obvious candidates. Furthermore, the authors suggested that replay sequences supported a hypothesis testing process for silhouette decomposition.

Other experimental studies demonstrate, more generally, that replay from spatial and non-spatial cognitive maps provides a neural substrate for goal-directed imagination and compositional computing. But few studies have, so far, advanced mechanistic explanations of these phenomena^{34,35,58,94,95}. The silhouette decomposition problem involves finding a set of BB (tiles) that completely cover a given silhouette without gaps or overlaps. This process is called the decomposition of the given silhouette. GCML generates candidate decompositions through goal-directed sampling, similar to the previously described navigation tasks. In the simplest version of the problem, the goal node o* is defined as the empty silhouette, whereas the starting node o₀ is the silhouette to be decomposed. An action is defined as removing a BB from a particular position in the remaining silhouette. Details of this process are provided below.

BBs

We used the same set of BBs as considered in the experiments of ref. ²², except that we removed one of the BBs (an atomic square tile) that would make every silhouette decomposable. Hence, there were eight predefined BBs. The first four BBs were categorized as corners, each with four possible orientations and a width of 1 pixel. The remaining four BBs were bars, including two longer types (1 × 3 and 3 × 1 pixel²) and two shorter types (1 × 2 and 2 × 1 pixel²). These eight BBs were derived from the human silhouette decomposition experiment of ref. ²². However, we excluded the 1 × 1 BB since each presence makes each given silhouette decomposable. GCML was tasked with imagining possible decompositions of given silhouettes, each composed exactly of five BBs.

The observation of a silhouette is encoded by a binary vector that has a 1 at each pixel position in a 2D grid space that corresponds to black in the silhouette.

The simplest possible compositional embedding Q of silhouettes is the identity map over these binary vectors. This provides already decent performance of the GCML. However, its performance increases substantially, and its cognitive map is endowed with a better sense of direction if the embedding Q slightly departs from the identity map and assigns greater weight to black pixels that are surrounded in a silhouette by white pixels. This embedding increases the utility of removing BBs that cover black pixels which protrude from a silhouette. This heuristic is beneficial since the choice of these BBs is more constrained and, therefore, less error prone. It also creates a cognitive map of the compositional domain that supports the same geometric heuristic for action selection as in the previously discussed GCML applications: choose an edge from the current node that best points into the direction of the goal (Fig. 4d). Further details and the full algorithm are provided in the Methods and Supplementary Section C.

Generation of training and test datasets

Observation of a silhouette was encoded by a binary vector that has a 1 at each pixel position in a 2D grid space that corresponds is black in the silhouette. Each silhouette in the dataset was created by randomly selecting five BBs and placing them within a 10 × 10 image at random positions. In total, 18,000 training samples and 2,000 test samples were generated. Any silhouette appearing in the training dataset was excluded from the test dataset, regardless of its position. Each sample represented a decomposition trajectory, in which BBs were sequentially removed from an initial silhouette. The silhouette generation process followed these rules: all BBs had to remain within image boundaries without overlapping. An initial BB was placed randomly within the image with equal probability for each position. Subsequent BBs were added adjacent to the existing silhouette, with at least one pixel adjacency to increase task difficulty. Longer adjacencies increased the number of possible removal actions, further challenging the model. In the training dataset, the five BBs composing each silhouette were removed sequentially in random order to generate training trajectories.

Details of actions

In the silhouette decomposition problem, each action a_t corresponds to removing a specific BB from a specific valid position within the 10 × 10 grid, resulting in 664 possible actions. Each action is uniquely encoded as a one-hot vector. This formulation aligns with human experiments²², where participants manipulate BBs using a mouse and keyboard to cover the silhouette. Since moving the same BB to different positions requires distinct movements, each BB at a different position is treated as a separate action.

Details of observations

In the silhouette decomposition problem, the observation \({{\bf{o}}}_{t}\in {{\mathbb{R}}}^{10\times 10}\) directly represents the silhouette image, rather than being encoded as a one-hot vector as in previous experiments.

High-dimensional embedding details

The simplest possible compositional embedding Q of silhouettes is the identity map over these binary vectors. This provides the already decent performance of the GCML. However, its performance increases substantially, and its cognitive map is endowed with a better sense of direction if the embedding Q slightly departs from the identity map and assigns greater weight to black pixels that are surrounded in a silhouette by white pixels. This embedding increases the utility of removing BBs that cover black pixels that protrude from a silhouette. This heuristic is beneficial since the choice of these BBs is more constrained and, therefore, less error prone. It also creates a cognitive map of the compositional domain that supports the same geometric heuristic for action selection as in the previously discussed GCML applications: choose an edge from the current node that best points into the direction of the goal (Fig. 4d). The state representation s_t = Qo_t lies in a 100-dimensional space that matches the dimensionality of the observation. In our implementation, however, this effect is achieved entirely through the context-dependent affordance term (see the paragraph describing \({{\mathsf{g}}}_{2}\)), which already assigns higher influence to protruding pixels through its convolution-based computation. As a result, no explicit modification of Q is required, and we maintain Q = I. Further algorithmic details are provided in Supplementary Section C.

Details of learning the cognitive map

The action embedding matrix V is learned according to equation (12). The learning rate of V was set to η_v = 0.01, with values between 0.0001 and 0.1 also demonstrating stable performance. In this silhouette decomposition task, each action corresponds to removing a specific BB at a specific valid location, and the corresponding action embedding V is trained to predict the resulting silhouette after the removal of that block (equation (11)).

The inverse model W is learned in a self-supervised manner from successful decomposition trajectories using equation (14) with a learning rate η_w = 1, and the weights in W saturate at 1:

As a result, W becomes a binary matrix. Both matrices V and W were initialized using Gaussian distributions \({\mathcal{N}}(0,0.001)\). The method remained effective across a broad range of initial standard deviations, from 0.0001 to 1.

One can easily see that W outputs for a given state difference for each possible action (BB) the number of pixels in the overlap between the state difference and the BB, provided that the BB is contained in the state difference (otherwise, it outputs 0).

Affordance details

The affordance gating mechanism in the silhouette decomposition problem comprises two components. The first component, \({{\mathsf{g}}}_{1}\), ensures that the agent selects only BBs present in the silhouette. We assumed GCML could directly access \({{\mathsf{g}}}_{1}\) from the environment during all the imagined steps. This assumption aligns with biological experiments²², where subjects directly observe the screen as they solve the silhouette decomposition problem. Thus, affordance information is assumed available at all steps, requiring no inference from \({\widehat{{\bf{s}}}}_{t}\).

The second component, \({{\mathsf{g}}}_{2}\), represents context-dependent affordance determined by the estimated state \({\widehat{{\bf{s}}}}_{t}\). Each pixel’s affordance value is set as the number of non-empty adjacent pixels (top, bottom, left and right). This is modelled as a convolution operation with the kernel:

using a 3 × 3 kernel size, padding of 1 with value 1, and stride of 1. The context-dependent affordance function \({{\mathsf{g}}}_{2}\) transforms the delta vector \({\boldsymbol{\Delta }}={{\bf{s}}}^{* }-{\widehat{{\bf{s}}}}_{t}\) into

where * denotes the convolution, and ⊙ represents element-wise multiplication. The delta vector Δ contains unremoved pixels with a value of −1 and all other pixels with a value of 0. The expression (1 + Δ) flips these values, converting −1 to 0 and 0 to 1; therefore, the subsequent convolution counts the number of adjacent pixels with a value of 0. Finally, since only the affordance values at unsolved pixels are relevant, an element-wise multiplication between the minus delta vector (−Δ) and the affordance value k⋅(1 + Δ) is performed as the second component of the affordance check.

Trying to decompose a silhouette

This task was formulated as a navigation problem, where the embedding of the empty silhouette, o* = 0, represents the target node in the high-dimensional state space (cognitive map), whereas the embedding of the initial silhouette, o₀, served as the starting node. In the presence of noise, GCML selected the action with the highest utility value from the set of currently feasible actions, determined by the environment’s affordance factors \({{\mathsf{g}}}_{1}(\cdot )\) and \({{\mathsf{g}}}_{2}(\cdot )\). This selection process followed the directional intuition of the cognitive map:

where the noise term ϵ followed a Gaussian distribution \({\mathcal{N}}(0,0.1)\).

After selecting an action a_t, the next state was estimated as \({\widehat{{\bf{s}}}}_{t+1}={\widehat{{\bf{s}}}}_{t}+{\bf{V}}{{\bf{a}}}_{t}\), and GCML proceeded to the next iteration to attempt the removal of another BB. The algorithm terminated when no feasible actions remained. If all pixels were removed, GCML successfully identified a valid decomposition; otherwise, the attempt was deemed unsuccessful.

The four paths from start to goal that are plotted in Fig. 4d with solid, dotted, dashed and dash–dotted line segments correspond to four trajectories that solve the problem. Their intermediate states are indicated in Fig. 4c (right). Note that all of these correct trajectories result from actions that move into the direction of the goal in this cognitive map. The top part of the left plot in the panel depicts three silhouettes that result from removing a ‘wrong’ BB from the starting silhouette that does not lead to a full decomposition of the initial silhouette. Note that they lie relative to the starting point on the opposite direction of the goal. Another dead-end configuration that arises if one does not move towards the goal node in the second removal step of the paths indicated by solid and dotted lines is shown at the bottom. Additional examples of cognitive maps and generated trajectories to decompose silhouettes are shown in Supplementary Fig. 4.

Details of performance evaluation

The performance of GCML was evaluated based on the proportion of test cases in which it successfully solved the given silhouette in the test dataset.

As baseline comparisons, we consider a random strategy that removes as many BBs as possible from the silhouette in a random manner, as well as the CML algorithm²⁶.

Decomposition to a given non-empty goal silhouette

In the non-empty goal silhouette decomposition problem, GCML removes BBs from a given silhouette o₀ to transform it into various non-empty target silhouettes o* ≠ 0. When generating the test dataset, the initial silhouettes contained six BBs, generated by following the same method described earlier. Any initial silhouette that appeared in the training dataset was removed, regardless of its position. The goal silhouettes contained two BBs, generated by randomly removing four BBs from the six BBs used to generate the initial silhouettes. GCML could plan trajectories for non-empty goal silhouettes using equation (24) without requiring retraining. The performance was evaluated using 1,000 test silhouette pairs.

Combining sihouette decomposition and composition

To address the silhouette rebuilding problem, GCML began with a given silhouette and iteratively modified it by adding or removing BBs to reconstruct the goal silhouette. This task was solved using the same cognitive map learned from the silhouette decomposition problem introduced earlier, leveraging mathematical symmetry.

GCML followed a two-stage approach. In the first stage, it removed BBs as long as possible from the given silhouette using a straightforward method. Here GCML adopted for that a sub-goal s^† = 0, and the delta vector was computed as Δ = s^† − s_t. Then, it switched to the second stage, where it aimed at constructing the given target silhouette by sequentially adding BBs. More precisely, the goal was the final silhouette s* = Qo*, and the delta vector was given by Δ = s* − s_t. Since adding BBs was the inverse operation of removing them, GCML selected actions based on

where the affordance factor \({{\mathsf{g}}}_{1}^{{\prime} }\) ensured that GCML only added BBs contained within the goal s*. The function \({{\mathsf{g}}}_{2}^{{\prime} }({\boldsymbol{\Delta }})=({\bf{k}}* ({\bf{1}}-{\boldsymbol{\Delta }}))\odot {\boldsymbol{\Delta }}\) was designed following the same principle as in equation (23). The negative transformation of the sign of Δ in equation (23) facilitated the addition of BBs as the inverse of their removal. Here \({{\bf{a}}}_{t}^{{\prime} }\) represented the action of adding a BB corresponding to a_t, which removed a BB.

For the testing dataset, the initial and goal silhouettes were generated using the five BBs following the previously introduced generation process. To test the robustness of our model, a randomly selected noise pixel was added to the same empty location in both start and goal silhouettes. Any initial or goal silhouette that appears in the training dataset was removed, regardless of its position. A total of 1,000 initial and goal silhouette pairs were generated to test the performance.

Comparison of performance of the GCML with competing methods for solving the k shortest path and tiling tasks

We compared the performance of the GCML on the k-shortest-path problem with that of three other algorithms that have been proposed for this task: K*, mA* and bidirectional edge labelling A* algorithm (BELA*) algorithms. The results are shown in Fig. 6a–c. In Fig. 6d, we compare the performance of the GCML for tiling tasks with that of an RL algorithm and the MPC algorithm. We found that the performance of the GCML is fairly competitive, in spite of the fact that it is an online algorithm: it immediately generates the first heuristic step towards a solution, without generating a complete solution path, or even generating and comparing many possible solutions paths in an offline manner, as the competing methods do (we are not aware of competing online methods for solving this task, except for the random method examined in Fig. 6d). Our comparisons show that the GCML requires altogether much less computational effort and latency. We also expect that it is superior when fast adaptation to a variation of the task is desired.

Detailed description of the competing algorithms and details of our comparisons are given below. We also discuss the differences in the amount of working memory that is required by different approaches.

k-shortest-path problem

As shown in the main text, GCML is an approximate k-shortest-path algorithm. We compare it with three k-shortest-path algorithm baselines: the classical K*⁹⁶, mA*⁹⁷ and more advanced BELA*⁹⁸. No heuristic form of the algorithms is used, since nodes in our random graphs environment are equivalent.

Details of K* algorithm

K* is an on-the-fly algorithm for enumerating the k shortest paths in a non-decreasing cost order⁹⁶. It first runs an A*-style search to obtain the shortest path and to incrementally reveal sidetrack edges (that is, detours) relative to the current shortest path tree. These detours are organized into a path graph on which a Dijkstra-like procedure is executed to repeatedly extract the next best detour combination, reconstruct the corresponding full path and output it. The two searches are interleaved: when the auxiliary search exhausts currently known detours, K* resumes A* to expand more of the original graph and updates the auxiliary structure, continuing until k solutions are produced.

Details of mA* algorithm

mA*⁹⁷ extends A*⁹⁹ from returning a single shortest path to enumerating the k shortest paths in non-decreasing path length. It maintains a priority queue of candidate nodes initialized with the start node and repeatedly selects the most promising node according to an evaluation that combines the accumulated path length with a heuristic estimate of the remaining distance to the goal (that is, the remaining distance is fixed at zero when no heuristic information is provided). The selected node is expanded to generate successors (that is, child nodes). The algorithm discards any successor that has already been generated k times; otherwise, the successor is inserted into the priority queue with its updated cost. Unlike standard A*, it does not terminate after the first goal-reaching path; each time a complete start-to-goal path is found, it is output as one solution and the search continues until k solutions are produced.

Details of BELA* algorithm

Compared with mA*, BELA*⁹⁸ does not enumerate the k shortest path by repeatedly pushing many alternative path-specific copies of the same state into the priority queue. Instead, BELA* keeps the forward search closer to a single A*-like exploration that builds an explored graph with sufficient predecessor information, and it delegates the top-k enumeration to a separate layer: it stores each newly certified deviation from the shortest-path tree as a centroid (that is, a compact representation of a family of solutions) and later reconstructs multiple complete paths at once by combining compatible prefixes and suffixes from the explored structure. Consequently, BELA* emphasizes the implicit representation of many solutions and outputs them only when their ordering among the next shortest solutions is guaranteed.

Measure of performance

To quantify the performance L_k of a k-shortest-path algorithm, we compute the relative error in the total length of the returned k paths. Let p denote the summed path length of the different k shortest paths produced by the algorithm, and let p_k denote the ground-truth summed path length. We define L_k as

Lower L_k is better.

We also report the latency to obtain the first shortest path to the goal, measured as the number of visited nodes required to generate the first trajectory. For the baselines, this metric counts how many times a node is removed from the candidate queue and expanded by generating successors and inserting them into the queue. It is a standard measure of search effort and an implementation-robust proxy for computational cost. Baseline methods typically visit far more nodes than the length of the returned trajectory because they must explore the graph. For GCML, each decision selects an affordance-feasible action and transitions to the next state, so the visited-node count equals the sampled trajectory length.

Comparison of performance

Here we set k = 5 for all experiments. We evaluate all methods on ten graphs for each node size (N ∈ {32, 64, 96, 128, 160}) and 100 random start and goal pairs per graph. The graph structure and the start and goal pairs remain unchanged for all algorithms. For GCML, to ensure a well-trained cognitive map at larger N, we fix the number of trajectories to 1,000 and the state dimensionality to 5,000; all other settings are provided in this section.

Results are shown in Fig. 6a,b. As shown in Fig. 6a, the three baselines are exact and, therefore, achieve L_k = 0. GCML degrades slightly as N increases. However, as shown in Fig. 6b, GCML visits far fewer nodes to generate the first potential trajectory of the solution than the baselines, benefiting from the directional guidance provided by the cognitive map, which allows it to avoid extensive exploration of the graph as baseline algorithms.

Comparison of replanning latency

We also measured replanning latency when the start or goal node changes, keeping the graph fixed. Replanning latency is defined as the wall-clock time an algorithm needs to explore the graph after the start or goal query is changed whereas the underlying graph remains fixed. All experiments related to wall-clock time in this section were run on a workstation equipped with dual Intel Xeon E5-2690 v4 CPUs (2.60 GHz), providing 28 physical CPU cores and 56 hardware threads in total, and 6 NVIDIA Quadro P6000 GPUs connected via PCIe, with 24 GB of on-board memory per GPU. As shown in Fig. 6c, the replanning time of the baseline algorithms increases with the number of nodes. By contrast, because GCML relies on the learned cognitive map, it incurs negligible additional replanning cost.

Comparison of working memory consumption

We define the working memory of the algorithms as the transient memory required during online trajectory generation. The three search-based baselines impose a substantially heavier online memory burden than GCML, although for different reasons. K* must preserve structures for sidetrack management and path enumeration. mA* must maintain a dynamically growing search frontier, a record of explored states and multiple predecessor references to support the recovery of alternative shortest paths. BELA* also requires a record of labelled explored state and auxiliary structures called centroids for combining prefixes and suffixes of candidate paths. As the number of requested solutions increases, these intermediate representations of the baseline algorithms can accumulate.

By contrast, GCML performs online generation by following a learned cognitive map prior and sampling directly towards the goal; therefore, its transient memory is dominated by the current local decision context and the partial trajectory being produced, rather than by a global search over the graph. Consequently, GCML’s online working memory is typically much closer to a local, near-constant-size inference buffer, whereas the baselines require graph-scale, search-dependent memory that grows with problem difficulty and the number of requested paths. In conclusion, GCML may incur a fixed static memory cost for storing model parameters and the learned cognitive map, whereas the working memory of the search baselines grows dynamically during online computation.

Tiling problem

We implemented two mature and widely used baselines on the same tiling environment setting used by our method (that is, five BBs in training silhouettes and eight BBs in the test silhouettes): an RL baseline and an MPC baseline.

Details of RL algorithm

For the RL algorithm, the duelling double-deep Q-network learning algorithm is adopted. The Q-function uses a duelling architecture¹⁰⁰ with double-deep Q-network target selection¹⁰¹. The definition of observation and action remains unchanged to GCML. When selecting actions, the action feasibility is enforced by an affordance mask as in GCML (invalid actions are masked before selection by affordance). We tuned the learning rate over {10⁻³, 10⁻⁴, 10⁻⁵, 10⁻⁶} and selected the value with the best validation performance. We also varied the depth of the encoder’s linear layers, the value and advantage heads, adopting the best-performing architecture.

The network model consists of a convolutional encoder, and two duelling heads for value and advantage estimation. The encoder has (1) a convolutional stack with channels 1 → 32 → 64 → 64, using 3 × 3 kernels with stride 1 and padding 1, followed by rectified linear unit activations, and (2) a feed-forward stack of 6,400 → 1,024 → 1,024 → 512 with LayerNorm and Gaussian error linear unit activations, producing a 512-dimensional state representation s. On top of this representation, the value head is an multilayer perceptron with layer sizes of 512 → 512 → 256 → 128 → 1 that outputs V(s), and the advantage head is a multilayer perceptron \(512\to 512\to 256\to 128\to | {\mathcal{A}}|\) that outputs A(s, a); both heads use LayerNorm and Gaussian error linear unit activations. \({\mathcal{A}}\) denotes the set of all possible actions.

We train the network with a duelling critic and experience replay. The Q-network is optimized with Adam (learning rate of 10⁻⁵) using mini-batches of 512 transitions sampled from a replay buffer of capacity 2 × 10⁵. We use a discount factor of γ = 0.99. Training is run for 10,000 episodes. Exploration follows an ϵ-greedy policy with ϵ linearly annealed from 0.30 to 0.02 over 1,200 episodes. During evaluation, to generate different trajectories, we keep a stochasticity with ϵ_eval = 0.3 (optimized among {0.05, 0.1, 0.15…0.4}). With probability ϵ_eval, it performs policy-guided stochastic exploration: affordance-masked action scores from the trained Q-network are converted into a probability distribution via softmax, and a valid action is sampled from this distribution. The per-step reward is designed to reward the valid decisions: \({r}_{t}=2\cdot {{\mathbb{I}}}_{{\rm{done}}}\), where \({{\mathbb{I}}}_{{\rm{done}}}\) indicates task completion.

Details of MPC algorithm

To choose an action, MPC¹⁰² plans an optimal action sequence under a learned policy model, executes only the first action and replans at the next step by discarding the remaining actions. In our implementation, MPC uses the cross-entropy method (CEM)¹⁰³ as its policy optimization (planning) procedure. CEM maintains a sampling distribution over action sequences and updates it iteratively over time. At each iteration, it draws candidate sequences, rolls them out through the learned dynamics to approximate the induced state trajectories and returns, and scores each candidate accordingly. A fixed number of the best-scoring sequences are then treated as elites, and their statistics are used to refit the distribution, yielding an updated action distribution for the current decision step.

Further details of the MPC algorithm are given below. The observation and action definition is same as GCML. At real step t, MPC optimizes an action sequence

for i = 1…N sampled rollouts, where H is the planning horizon and N = 96 is the sample count. We set the horizon H to 5 to stay consistent with training on five-BB silhouettes, and then test performance on eight-BB silhouettes.

Sampling is performed from per-time-step categorical distributions:

with affordance like GCML. Here π_τ(⋅) is categorical policy at planning step τ. After simulating a candidate sequence, we compute the CEM score defined as

where \({r}^{(i)}=\parallel {{\bf{s}}}_{t+H}^{(i)}-{{\bf{o}}}^{* }{\parallel }_{1}\) is the terminal mismatch, \({\chi }^{(i)}={\mathbb{I}}[{r}^{(i)}=0]\) is the success flag, \({{\bf{s}}}_{t+H}^{(i)}\) is the remaining silhouette at step t + H, o* is goal silhouette and β = 100 is terminal success bonus. Higher J⁽ⁱ⁾ is better.

For each planning step τ, the elite empirical distribution is

where c_τ(a) is the elite count of action a at step τ, and δ = 10⁻⁶ is the minimum action probability. Then, distribution smoothing is

followed by clamping π_τ(a) ≥ δ and renormalization. Here α = 0.3 is the smoothing factor.

CEM optimization runs for I = 4 iterations, then only the first action of the best sequence is executed:

One BB is removed following \({a}_{t}^{\star }\), horizon shifts by one step, and planning repeats until success or no more BBs could be removed.

Comparison of performance

The success rates (defined in the same way as for GCML) for the random baseline, RL, MPC, CML and GCML are shown in Fig. 6d. All algorithms are evaluated on the same 2,000 test samples as described here.

For the RL baseline, the network architecture is much more complex than those of the other methods. It also requires backpropagation to train, which makes deployment on neuromorphic hardware challenging. In addition, the reward is difficult to design, and the sparsity of the reward signal makes training unstable.

For the MPC baseline, there is no policy training, and it achieves higher success rate than the RL baseline. However, compared with RL and GCML, MPC incurs higher planning cost because it performs an iterative sampling-and-refinement procedure at every decision step, rather than amortizing computation across episodes. In our implementation, MPC uses CEM to optimize a horizon H action sequence by repeatedly rolling out N candidate sequences for I iterations, which results in a per-step workload proportional to O(I⋅N⋅H). This issue is exacerbated in compositional tiling tasks in which early placements can create irreversible geometric constraints; therefore, short-horizon optimization may fail to anticipate long-term feasibility. Consequently, when the sampling budget is reduced (smaller N or fewer iterations I), MPC’s distribution updates become noisy and can collapse to suboptimal modes, leading to a pronounced drop in success rate. By contrast, GCML leverages a cognitive map that provides a global embedding of the observations and the learnt utility, which guides sampling towards regions that are more likely to yield feasible completions. Therefore, under comparable or even smaller online budgets, GCML can maintain higher success and better sample efficiency than MPC in the current setting. Overall, we emphasize that MPC remains a strong general-purpose control baseline, but its per-step compute and horizon-limited look ahead makes it less attractive for the highly combinatorial environments like the tiling problem. Same as the RL baseline, it also needs to define a rollout score equation (28) to measure the quality of a candidate action sequence.

Our GCML achieves the best overall success in this task across different sample numbers and requires algorithmically simpler local updates compared with RL and MPC.

Reporting summary

Further information on research design is available in the Nature Portfolio Reporting Summary linked to this article.