A mechanistic theory of planning in prefrontal cortex

Ring attractors are neural networks that can only encode circular variables (Zhang, 1996; Ben-Yishai et al., 1995). Different neurons have different ‘preferred orientations’, and the network dynamics have a set of fixed points (stable activity patterns) that encode a particular angle. At these fixed points, neurons are only active if their preferred angle is close to the encoded angle (Figure 2A, bottom left). This is achieved by a network where neurons are mutually excitatory if they have similar preferred orientations, and inhibitory otherwise (Figure 2A, top left). The network dynamics infer which angle is most compatible with noisy inputs such as vision and proprioception (Turner-Evans et al., 2017; Skaggs et al., 1995; Kim et al., 2019; Kutschireiter et al., 2023).

Figure 2 with 1 supplement see all

Download asset Open asset

Grid attractors instead encode two-dimensional locations. This is achieved by neurons that are mutually excitatory if they represent nearby locations and inhibitory if they represent more distant locations (Figure 2A, top right; Fuhs and Touretzky, 2006; Burak and Fiete, 2009). The fixed points of the network are hexagonal patterns of activity that resemble canonical grid cells (Figure 2A, bottom right; Hafting et al., 2005). Inputs to the network again determine which location is represented in a particular context (Burak and Fiete, 2009).

Inspired by these known biological attractor networks, we propose the existence of a spacetime attractor in PFC. The STA is a network that infers explicit representations of the future, and its fixed points represent entire paths through space and time (Figure 2B, bottom; ‘Methods’). Similar to ring and grid attractors, this is achieved by connecting STA neurons such that consistent states excite each other and inconsistent states inhibit each other. In the STA, each neuron has both a preferred location and a preferred delay $\delta$, which determines how far into the future its tuning curve is defined. A neuron with $\delta =0$ fires when the agent is currently at the preferred location of that neuron. A neuron with $\delta =3$ fires when the agent expects to be at the preferred location after three actions. In such a network, consistent states can form part of a single trajectory, while inconsistent states cannot. The connectivity between neural populations with preferred delays $\delta$ and $\delta +1$ should therefore correspond to the structure – more specifically the adjacency matrix – of the environment (Figure 2B, top). Unlike ring and grid attractors, the fixed points of the STA depend on tonic inputs. However, the connectivity constrains the fixed points to represent continuous trajectories for any combination of inputs. In this section, we show this empirically. Later, we will see that such connectivity is optimal for planning-as-inference.

To compute a plan, it is necessary to know which states will be rewarding in the future. This reward information is provided as an input to the STA and enables fast adaptation without rewiring the synaptic connections. It alters the fixed points of the recurrent dynamics to only include trajectories that are also associated with high cumulative reward (Figure 3; ‘Methods’). This resembles a previous proposal that PFC can compute stable ‘attraction paths’ from the current state to a distant goal (Duncan, 1990). Importantly, the STA can accommodate problems with time-varying reward structures known to engage PFC (Figure 1D; Shallice and Burgess, 1991; Volle et al., 2011; Carlesimo et al., 2014). This is possible because neural populations that represent different times in the future can receive different inputs. Given a restaurant booking in 2 hours and a cinema ticket in 4 hours, the $\delta =2$ neurons would receive reward input at the restaurant, and the $\delta =4$ neurons at the cinema. The network dynamics would then infer a future trajectory, where the agent is at each location at the appropriate time. There are several possible sources of reward input to a STA in PFC (Appendix 1). We focus on planning under a known reward function and therefore assume access to ground-truth rewards.

Figure 3 with 1 supplement see all

Download asset Open asset

To summarise, the STA has four main components: (i) different neural populations represent different times in the future; (ii) the neurons are connected according to the structure of the environment; (iii) the current location is an input to the ‘present’ subspace; and (iv) the reward structure is an input to all future subspaces. The resulting fixed points (Figure 2B, bottom) represent trajectories that maximise cumulative reward, and they resemble an approximate posterior over the future in a probabilistic formulation of planning (‘Methods’). The network dynamics infer future locations through a gradual relaxation process, where reward inputs bias the representation in each subspace towards high-reward locations, while inputs from neighbouring subspaces constrain their representations to be consistent (Figure 2C and Figure 3).

After computing a plan, it can be used to inform behaviour. In particular, an agent should take an action that leads to the next location on the inferred trajectory (Figure 2C, middle). The entire representation of the future then moves by one ‘action’, so the state that was previously represented by subspace $\delta$ is now in subspace $\delta -1$. The resulting dynamics resemble a ‘conveyor belt’ that allows the STA to execute entire trajectories without recomputing the policy. Such conveyor belt dynamics have been observed experimentally during sequence working memory (El-Gaby et al., 2023). Importantly, the representation remains a fixed point of the STA dynamics, because the location and reward inputs are updated when the agent moves and time passes.

Many planning algorithms apply a ‘world model’ sequentially to simulate the consequences of different actions one at a time. The STA instead instantiates multiple copies of a forward and inverse world model in the synapses between different neural populations (‘Methods’). Recurrent dynamics can then arbitrate between many possible futures in parallel. Entorhinal grid cells also embed a world model in their connectivity (McNaughton et al., 2006), but they only encode a single location at a time (Vollan et al., 2026). Such networks can generate sequences, but the individual elements are represented one by one (Sompolinsky and Kanter, 1986; Kleinfeld, 1986; Widloski et al., 2025). The STA suggests that circuit principles in PFC resemble other cortical areas that use structural knowledge to infer features of the world. The major difference is that PFC instantaneously represents many points in time, which generalises known circuit principles to complex planning.

We designed a set of four tasks that vary in how much the reward changes in space and time (Figure 4A). The tasks are all embedded in Euclidean space for simplicity of exposition, but the underlying principles generalise to other environments with known structure. Task 1 is simple navigation towards a static goal that remains constant across trials. In task 2, the static goal changes between trials. Task 3 is navigation to a goal that also moves within each trial, where each location is only rewarded when the goal is there. Task 4 generalises the idea of time-dependent goals to a non-binary reward landscape. Reward magnitudes are sampled independently between –1 and +1 for each location at each time point in each trial. The objective is to maximise cumulative reward over the trial, which requires balancing immediate reward with the potential for future reward (Figure 4A, right). This is reminiscent of the example from Figure 1D, where the restaurant and cinema are desirable at different times. To better understand how the STA solves these tasks, we compare it to two models commonly studied in systems neuroscience. The first is temporal difference (TD) learning, which has neural correlates in striatum (Sutton, 1988; Schultz et al., 1997). The second is the successor representation (SR), which has neural correlates in hippocampus (Dayan, 1993; Stachenfeld et al., 2017). These models assume fixed rewards across trials (TD) or within a trial (SR). Unlike the STA, they do not generalise well when the reward changes rapidly.

Figure 4 with 2 supplements see all

Download asset Open asset

TD learners gradually propagate value from rewarded locations to all other locations, and optimal policies are only learned when rewards are constant across both time and trials (Figure 4B and C, orange). The SR agent uses the environment structure and trial-specific reward function to compute values, and it can solve tasks where the goal changes between trials (Figure 4B and C, green). The STA can also solve these tasks, but it does so by inferring an entire spatial trajectory to the goal instead of computing a value function (Figure 4B and C, blue). Because the reward is constant throughout a trial, each subspace receives the same reward input. However, the inferred representation differs across subspaces because it is also constrained to satisfy the current agent location and the transition structure of the world. As we will see, the inference of an explicit spacetime representation by the STA increases its flexibility compared to the amortised state values computed by TD and SR agents.

While the SR agent can adapt to rewards that change across trials, the reward structure still needs to remain constant for the duration of the planning horizon. This is because the SR computes time-averaged occupancy and lacks fine-grained information about when the agent is where. When the reward changes within a trial, the SR averages values across time-within-trial and fails to intercept the goal at the correct point in space and time (Figure 4D). The STA can solve the moving goal task because the input to each subspace is the reward at that specific moment in time. The representation is therefore biased towards coinciding with the goal in both space and time, and the network dynamics relax to a future trajectory that correctly intercepts the goal (Figure 4D). The advantage of the STA is exacerbated in the more dynamic reward landscape task, where very different locations can be rewarded at different times (Figure 4E). These results suggest that different algorithms could contribute differently to decision-making. TD learning could drive rapid decisions in environments with stable rewards; the SR would be more flexible when rewards change on intermediate timescales; and the STA could enable adaptive behaviour in environments with rapid changes. Finally, sequential search would facilitate slower planning in novel environments and could be guided by partial plans from a STA.

The representation and connectivity of the STA are inspired by theories of planning as an inference over future locations. However, the proposed weights may not be optimal for this computation. We therefore trained a recurrent network to perform inference over representations of the future, which were constrained to be factorised over time. The exact posterior for each future time step has a known functional form, which combines two factors that reflect constraints from earlier and later times in the plan. These factors can be approximated by recurrent inputs from earlier and later subspaces respectively (Appendix 2). However, we do not know the optimal form of the recurrent weights. We therefore optimised the weights numerically to maximise the similarity between (i) the exact posterior over future locations and (ii) the fixed points of the network dynamics (‘Methods’). The resulting connectivity resembled the handcrafted STA (Figure 5—figure supplement 1). In particular, neurons representing consecutive time steps were connected according to the environment structure. Unlike the handcrafted STA, more distant populations were also weakly connected. This connectivity captured the higher-order structure of possible trajectories.

So far, we have directly imposed an explicit representation of the future. However, this representation requires many neurons. We investigated whether superior representations exist by training RNNs to solve dynamic planning problems efficiently (Mante et al., 2013; Stroud et al., 2023). We will see that regularised RNNs learn to solve these ‘PFC-like’ problems with both representations and connectivity that resemble a STA, which suggests that it is an efficient solution. The RNNs consisted of rectified linear units, and they were trained with supervised learning to generate optimal actions from inputs that indicated the current location and trial-specific rewards (Figure 5—figure supplements 2 and 3; ‘Methods’; Appendix 1). The loss function only constrained the upcoming action and did not impose a representation of the distant future. To avoid biasing the RNN representation, reward information was only provided during initial ‘planning’ prior to the ‘execution phase’ of the task. The objective function also penalised large firing rates and network weights to encourage an efficient solution.

The defining features that allow an STA to plan via inference are (i) it has an explicit representation of the future, (ii) the connectivity embeds the environment structure, and (iii) attractor dynamics constrain the representation. RNNs trained on the reward landscape task exhibit all of these properties. These results are consistent with a prominent theory that PFC resembles a recurrent meta-learner, which learns connections over long timescales that rapidly adapt to new problems through recurrent dynamics (Wang et al., 2018). The STA similarly embeds stable structure in its weights and uses recurrent dynamics for fast planning with different rewards. We show that an STA is the mechanism implemented in the dynamics of meta-learners trained on challenging planning tasks.

STAs compute explicit representations of the future (Figure 5A and B). Critically, these representations generalise over trajectories. In other words, the representation in subspace $\delta$ depends only on the expected location in $\delta$ actions, and not on the rest of the trajectory (Figure 5A, right). To confirm that the RNN learned such a representation, we trained linear decoders to predict the future from its hidden state. To ensure generalisation across trajectories, we trained the decoders while holding out each ‘current location’ and computed the test accuracy only from those held out locations. The RNN had an explicit representation of the future during both planning and execution (Figure 5C and D, Figure 5—figure supplement 4). These representations were high-dimensional and explained substantial variance in the RNN activity (Figure 5—figure supplement 5). Consistent with experimental recordings from PFC (Chiang and Wallis, 2018), the STA and RNN both represented the future more strongly than the past (Figure 5D). In the STA, the same subspace always encodes locations $\delta$ actions into the future. The future representation of the RNN also exhibited such ‘conveyor belt’ dynamics (Figure 5E). The neural representation of the RNN is thus consistent with an STA.

Figure 5 with 7 supplements see all

Download asset Open asset

Intriguingly, RNNs trained on the simpler static and moving goal tasks did not reliably learn explicit future representations (Figure 5—figure supplement 6). Additionally, the RNN trained on the reward landscape task could solve these simpler tasks, while RNNs trained on the simpler tasks failed in the reward landscape task (Figure 5F). Alternative solutions therefore exist in the simpler tasks, and these solutions are favoured by a pressure to be energetically efficient (Figure 5G). We identify a tradeoff between generality and efficiency, where the PFC-like solution is general but needs more neurons and synapses, while simpler algorithms can solve simpler tasks more efficiently. Unlike value coding and related theories, this suggests a reason why PFC occupies such a large cortical territory.

In an STA, subspaces that represent expected future locations are connected according to the environment structure. Testing this prediction in the RNN requires interpretation of the network weights, which is challenging because the functions of individual neurons are unknown. We therefore project the weights into a coordinate system, where the axes are orthogonal directions in neural state space that predict each point in spacetime (‘Methods’). While this does not perfectly recover the true subspaces in the handcrafted STA, it is a good estimate. The projected RNN weights can therefore be interpreted as interactions between representations of different points in spacetime (Figure 6A).

Figure 6 with 6 supplements see all

Download asset Open asset

Remarkably, the recurrent weights between adjacent subspaces closely resembled the adjacency matrix of the environment (Figure 6B, Figure 6—figure supplement 1; mean ± std correlation of 0.91 ± 0.07 vs. 0.72 ± 0.06 for control environments). The RNN thus learned a world model explicitly in its recurrent weights. Consistent with an STA, the $\delta =0$ subspace received input indicating the current agent location, and subspace $\delta$ received input indicating the reward in $\delta$ actions (Figure 6C). The RNN also learned weaker connections between more distant subspaces (Figure 6D and E). These connections reflected higher-order transition statistics, reminiscent of networks optimised explicitly to infer future representations (Figure 5—figure supplement 1). In summary, the RNN learned all of the structural components necessary to infer future actions by using a world model to integrate expected reward across time.

In an STA, structured connections give rise to attractor dynamics with fixed points corresponding to explicit representations of desirable futures. To demonstrate similar dynamics in the trained RNN, we perturbed its neural activity at the end of planning. This resembles previous experiments in biological attractors (Inagaki et al., 2019; Kim et al., 2017; Vinograd et al., 2024). An attractor network should be robust to small perturbations, and it should be more sensitive to perturbations towards other plausible representations than perturbations along random directions in neural state space. We analysed the RNN in a setting with two ‘good’ paths that were close in value (Figure 6F). Like the handcrafted STA, the RNN initially converged to a stable representation of the better path, which persisted under weak perturbations (Figure 6G, left). Stronger perturbation of a single location on the alternate path caused a discrete change to a representation of this entire path, including in subspaces not directly perturbed (Figure 6G, right). This new representation remained after removing the perturbation in the handcrafted STA, suggesting that it is also a fixed point of the unperturbed dynamics (Figure 6—figure supplement 2). In contrast, the RNN representation relaxed back to the original fixed point after removing the perturbation (Figure 6H, Figure 6—figure supplement 2). The unperturbed RNN dynamics therefore only have a single stable fixed point, suggesting that the network may be less susceptible to local minima.

Together, our analyses show that RNNs trained on a dynamic planning task learn to approximate an STA. This was also true across variations in model architecture (‘Methods’; Figure 6—figure supplements 3 and Figure 6—figure supplement 4). These results extend previous findings that explicit spacetime representations are optimal for sequence memory (Appendix 1; Whittington et al., 2023; Dorrell et al., 2026; Wang et al., 2025). Additionally, RNNs with too few hidden units to learn an STA failed to solve the task (Figure 6—figure supplement 5), suggesting that other solutions are not readily learned by gradient descent. The STA is thus an efficient solution to dynamic planning problems, and it is consistent with a prominent theory that PFC can be understood as a recurrent meta-learner.

To study planning in environments with changing structure, we trained RNNs on a variant of the reward landscape task, where the transition structure was different on every trial. The structure was provided to the RNN as a binary input that indicated whether otherwise adjacent states were separated by a wall (‘Methods’). Importantly, the network had to adapt to each new environment through its dynamics while keeping the weights fixed across environments (Wang et al., 2018; Jensen et al., 2024). The generalising RNNs performed nearly as well as networks trained in a single environment (Figure 7B), raising the question of how generalisation is achieved. The neural representation resembled an STA in any single environment (Figure 7—figure supplement 1B). However, the effective connectivity between subspaces changed between environments. Remarkably, the connectivity reflected the transition structure of whichever environment the agent was currently in (Figure 7C and D). This indicates a general solution that allows recurrent networks to use attractor dynamics for planning in any similar environment, without requiring synaptic plasticity.

We can understand how the network adapts to the environment at both the ‘population level’ and at a mechanistic level. At a population level, the directions in neural state space that represent each point in spacetime were more similar on different trials from the same environment than across trials from different environments (Figure 7E, left). This change allows the network to align its representation with the appropriate components of the connectivity matrix (Figure 7E, right). How is this achieved mechanistically? We hypothesised that the network explicitly represents future transitions rather than future locations. In other words, there are directions in neural state space that represent ‘in $\delta$ actions from now, move from state $s_i$ to state $s_j$’, which we denote ${\tau }_{\delta }^{ij}$. This generalises the vanilla STA, which has a representation of future locations that is invariant to the subsequent state.

While it requires more neurons, explicitly representing future transitions allows the network to use input-mediated inhibition to adaptively ‘turn off’ transitions that are not available in any particular environment. The learned structure of the remaining transitions is then used for planning (Figure 7F). Importantly, the scaffold only needs to include transitions that are possible in at least one environment. This model makes three predictions that we verified: (i) the network explicitly represents future transitions (Figure 7—figure supplement 1C), (ii) the representations of future transitions are connected by a world model (Figure 7G), and (iii) input indicating a wall between $s_i$ and $s_j$ specifically inhibits representations of future transitions between those states (Figure 7H). These analyses show that the STA is an efficient planning algorithm in structurally changing environments, and they suggest a general mechanism for adaptive behaviour.

We focus on the problem of planning, which can be defined as computing a sequence of actions that maximises expected future reward. To solve this problem, it is necessary to know (i) the current state $s_t$ of the agent, (ii) the transition structure of the environment, and (iii) which states are likely to be rewarding in the future. We assume that the transition structure remains constant throughout each trial and use $\mathit{𝑨}\in {ℝ}^{N\times N}$ to denote the adjacency matrix of the environment. Here, $N$ is the number of states, and $\mathit{𝑨}$ has elements $A_{ij}=1$ if it is possible to move from state $s_t=j$ to $s_{t+1}=i$, and $A_{ij}=0$ otherwise. We allow the reward function $\mathit{𝑹}\in {ℝ}^{T\times N}$ to vary over the course of a trial, such that $R_{\delta i}$ indicates the reward associated with visiting state $s_{t+\delta }=i$ in $\delta$ actions, and $T$ is the fixed length of the planning horizon. The challenge of estimating future rewards from past data is beyond the scope of this paper, and we simply assume access to a ground-truth reward function. Instead, we show how a recurrent neural network can embed $\mathit{𝑨}$ in its weights and take $s_t$ and $\mathit{𝑹}$ as inputs to compute a trajectory $\{s_t,s_{t+1},\dots ,s_{t+T}\}$ with high cumulative reward $\sum _{\delta =0}^T{R}_{\delta }(s_{t+\delta })$.

The spacetime attractor is a recurrent neural network with an exponential nonlinearity and block-wise normalisation. We define $z_{\delta i}$ as the ‘potential’ and $r_{\delta i}$ as the ‘firing rate’ of a neuron that represents a plan to ‘be at location $s_{t+\delta }=i$ in $\delta$ actions’. Each iteration of network dynamics then involves three steps:

Here, $\tau =50$ network iterations is the ‘time constant’ governing the exponential decay of the neural potential in the absence of inputs. $A_{ij}$ are weights corresponding to the adjacency matrix, and $R_{\delta i}$ is an input that reflects the reward available at location $s_{t+\delta }=i$ in $\delta$ actions. To anchor the plan at the current agent location, we set $R_{0i}=20$ if $s_t=i$ and 0 otherwise. $\beta$ is a temperature parameter that determines how strongly the agent values reward, and we set it to $\beta =9$ unless otherwise noted. We take the policy of the agent to be greedy in the representation in subspace $\delta =1$ of locations accessible from the current state: $a={\text{argmax}}_{i\in \mathcal{N}(s)}\left[r_{1i}\right]$. The action $a$ specifies which location to move to next. We used a maximum planning horizon of $T=6$ for all analyses in the main text, a planning horizon of $T=16$ in Figure 4—figure supplement 1B, and a planning horizon of $T=32$ in Figure 4—figure supplement 1C. The normalisation in line (2) can be implemented through a subspace-specific inhibitory neuron with a logarithmic nonlinearity ($I_{\delta }=\text{log}[z_{I_{\delta }}]$), which receives input from all excitatory neurons in the subspace (${\dot{z}}_{I_{\delta }}=-z_{I_{\delta }}+\sum _i{r}_{\delta i}$) and provides reciprocal inhibition. In practice, we normalised analytically for simplicity of implementation.

The fixed points of the STA dynamics have the form

At convergence, the representation is therefore biased towards high-reward states (first term), while also being constrained to represent continuous trajectories (second and third terms).

Here we briefly outline how the STA dynamics are inspired by previous models of ‘planning-as-inference’. Following Botvinick and Toussaint, 2012, we formulate planning as an inference problem in a Markovian probabilistic graphical model. We assume a uniform prior over possible actions, such that $p(s_{t+1}=i|s_t=j):=T_{ij}=\frac{A_{ij}}{\sum _k{A}_{kj}}$. In line with Solway and Botvinick, 2012 and Levine, 2018, we define observations $o\in [0,1]$ as binary ‘optimality variables’ with emission probabilities

where we define ${\mathcal{O}}^{\ast }:=\bigcap _{\delta =0}^T[o_{t+\delta }=1]$ as the condition that $o_{t+\delta }=1$ for all $\delta$. States with high reward are therefore more likely to have $o=1$ than states with low reward, and the probability that all observations take a value of 1 is monotonic in the reward associated with the entire future trajectory. We can now condition on ${\mathcal{O}}^{\ast }$ to compute a posterior distribution over future states:

The posterior marginals can be computed with the forward-backward algorithm, but that requires three separate neural populations for each time in the plan to represent forward messages, backward messages, and their product. Instead, we define a factorised distribution $Q(\{s_{t+\delta }\})=\prod _{\delta }{q}_{\delta }(s_{t+\delta })$ and optimise it to resemble the true posterior as closely as possible. We assume that each component of this distribution is encoded in the firing rates ${\mathit{𝒓}}_{\delta }\in {ℝ}^N$ of a single population of neurons, such that $q_{\delta }(s_{t+\delta }=i):=r_{\delta i}$.

It is common to factorise the true posterior distribution over trajectories as

However, we can instead ‘split’ the transition terms across consecutive factors in a symmetrised factorisation:

We can therefore define a ‘factor’ $f_{\delta }(s_{t+\delta },s_{t+\delta -1},s_{t+\delta +1}):=p(o_{t+\delta }=1|s_{t+\delta })p(s_{t+\delta }|s_{t+\delta -1}{)}^{\frac{1}{2}}p(s_{t+\delta +1}|s_{t+\delta }{)}^{\frac{1}{2}}$ for each time in the plan. While the approximate $q_{\delta }$ only depends on $s_{t+\delta }$, $f_{\delta }$ also depends on $s_{t+\delta -1}$ and $s_{t+\delta +1}$. We want our approximate posterior to resemble the product of true factors as closely as possible:

Loosely inspired by expectation propagation (Minka, 2013), we choose to locally minimise the discrepancy between the true and approximate posteriors. In other words, we define an objective function for component $\delta$ as the KL divergence between (i) the factor $f_{\delta }(s_{t+\delta },s_{t+\delta -1},s_{t+\delta +1})$, while approximating the distributions over $s_{t+{\delta }^{\mathrm{\prime }}\ne \delta }$ by $q_{{\delta }^{\mathrm{\prime }}}(s_{t+{\delta }^{\mathrm{\prime }}})$; and (ii) the distribution $Q(\{s_{t+\delta }\})$ implied by the firing rates $\left\{{\mathit{𝒓}}_{\delta }\right\}$. We also introduce a normalisation constant $Z_{\delta }$ to the first distribution. This gives rise to a loss function of the form

Here, $T_{ij}^{\frac{1}{2}}$ indicates an element-wise square root of each term in the transition matrix. We introduce a Lagrange multiplier λ for normalisation and take derivatives to minimise this objective:

This local optimisation of a factorised approximate posterior leads to the ‘mode-seeking’ dynamics seen in Figure 3B. The fixed points have the form

We now recall the STA fixed points:

The difference between the two is that the STA uses the adjacency matrix $\mathit{𝑨}$ in place of the element-wise square root of the transition matrix, ${\mathit{𝑻}}^{\frac{1}{2}}$. The STA therefore has the form of a dynamical system with fixed points that locally approximate the posterior marginals over trajectories, conditioned on ‘optimality’ in the planning-as-inference formalism. It also has very similar dynamics to another neural implementation of planning-as-inference by Donnarumma et al., 2025. However, the local approximation with marginals may be quite poor as it neglects parts of the transition matrix. For this reason, Figure 5—figure supplement 1 and Appendix 2 directly optimise the dynamics of a recurrent network for inference.

Here we note a few deviations from a simple interpretation of the STA dynamics as approximate inference.

First, we use the adjacency matrix instead of the transition matrix or its square root, which differs from a probabilistic interpretation. We do this because the true posterior distribution favours trajectories that pass through low-degree states, which have higher probability under the prior. In contrast, the recurrent weights between RNN subspaces representing consecutive future actions in Figure 6 were generally more similar to the adjacency matrix than the transition matrix. One way to interpret the adjacency matrix is as an action-optimised transition matrix, $A_{ij}={\text{max}}_{a}p(s_{t+1}=i|s_t=j,a_t=a)$. This generalises the idea to stochastic environments, but it does not optimise a well-defined objective to the best of our knowledge. All results in the paper hold qualitatively when using $T_{ij}^{\frac{1}{2}}$ in place of $A_{ij}$ for an implementation that is more faithful to planning-as-inference.

Second, we add a small amount of noise ${\eta }_{ij}\sim \mathcal{U}(-0.025,-0.015)$ to each element of the weight matrix instead of using the exact adjacency matrix. Noise is added because it is unreasonable to expect biological systems to exactly recapitulate the transition structure of the environment. It is restricted to be negative because the STA can be susceptible to discontinuous plans if noisy weights imply a finite probability of transitioning between two distant states (Appendix 1). We similarly added white noise $\xi \sim \mathcal{N}(0,\sigma =0.1)$ to the potential of all neurons at each iteration of the network dynamics.

Third, we threshold both the neural potentials ($z_{\delta i}$), the forward messages ($\log \sum _j{A}_{ij}{r}_{\delta -1,j}$), and the backward messages ($\log \sum _k{A}_{ki}{r}_{\delta +1,k}$) at a minimum value of $ϵ=-100$ for numerical stability.

Finally, the STA representation updates automatically when the agent moves, since both the location and reward inputs change. However, an additional ‘feedforward component’ in the form of the identity matrix was added to the connectivity between subspaces $\delta$ and $\delta -1$ during the first 100 network iterations (2 time constants) after every action. This is inspired by the ‘update neurons’ of the fruit fly head direction circuit (Turner-Evans et al., 2017) and stabilises the ‘conveyor belt’ dynamics, but it is not necessary for any of the main results in the paper (Appendix 1). We ran the STA dynamics for 400 iterations before each environment update to ensure convergence to a stable representation.

In this task, a full goal ‘trajectory’ was sampled on each trial. The goal trajectory was sampled as a restricted random walk that could only turn around if it reached a dead end. The start location of the agent was restricted to not coincide with the start location of the goal. The trial terminated when the agent was at the same location as the goal at a given moment in time, or after a maximum of $T=6$ actions. The reward input to the agent, $\mathit{𝑹}\in {ℝ}^{T\times N}$, was a matrix indicating the location of the goal at every point in the future. The reward input consisted of ‘+0.6’ for the goal location and ‘–0.6’ for all other locations at each moment in time. In other words, $R_{\delta i}$ was ‘+0.6’ if the goal would be at location $s_{t+\delta }=i$ after $\delta$ actions, and -0.6 otherwise.

This task was identical to the moving goal task, except that the goal remained stationary for the duration of each trial. The static goal task was implemented in two different variants – one where the goal remained fixed across all trials, and one where the goal was resampled at random in every new trial. The handcrafted models were evaluated in both of these tasks, while the RNNs were only trained with a goal that changed between trials.

In this task, every element $R_{ti}$ of the reward function $\mathit{𝑹}\in {ℝ}^{T\times N}$ was sampled uniformly and independently between –1 and +1 at the beginning of each trial. The complete future reward structure was provided as an input to the agent as described above. Every trial finished when the agent had taken $T=6$ actions.

For most performance comparisons, we computed the probability of choosing the optimal first action in each trial. We define the optimal first action as the first action of the trajectory that maximises cumulative reward over the entire trial. We use this metric rather than the cumulative reward for two reasons. (i) We are interested in the process of planning, whereby an agent balances immediate and long-term reward. This is most challenging for early actions, while the greedy policy is optimal for the last action. (ii) The probability of choosing an optimal action is readily interpretable as a number between 0 and 1. All results were qualitatively similar if we instead used the probability of choosing an optimal action averaged over the entire trial as a performance metric, or if we used the average empirical reward.

For the analyses in Figure 4—figure supplement 1, we analysed whether the entire trajectory implied by the STA representation was correct. A representation was considered correct if the location with the highest firing rate in each subspace (i) was adjacent to the previous location, (ii) was closer to the goal than the previous location, and (iii) had a firing rate of at least 0.3.

We implement vanilla TD learning, which computes a value function by iteratively applying the update

We set the temporal discount factor to $\gamma =1$, since we are interested in maximising the non-discounted cumulative reward. We allowed the TD agent to interact with the environment for 4000 trials with a learning rate of $\alpha =0.05$ before analysing its performance.

In the SR, the value function is decomposed as

Here, $\mathit{𝒓}$ is a vector of the average reward associated with each state, and ${\mathit{𝒎}}_s^{\pi }$ is a vector of the expected discounted future occupancy of state $s^{\mathrm{\prime }}$ if the agent starts in state $s$ and follows policy $\pi$:

The full matrix ${\mathit{𝑴}}^{\pi }$, constructed from stacking the ${\mathit{𝒎}}_s^{\pi }$ corresponding to all states $s$, is denoted the ‘successor matrix’, and it allows us to write down a vector of expected rewards from all states as

${\mathit{𝑴}}^{\pi }$ can be learned using a TD-like algorithm as above. However, since the STA algorithm also has explicit access to the transition structure of the environment, we simply computed the exact successor matrix as the geometric series ${\mathit{𝑴}}^{\pi }=\mathit{𝑰}+\gamma {\mathit{𝑻}}^{\pi }+{\gamma }^2{\left({\mathit{𝑻}}^{\pi }\right)}^2+\mathrm{\dots }={\left(\mathit{𝑰}-\gamma {\mathit{𝑻}}^{\pi }\right)}^{-1}$ with $\gamma =0.95$. When computing the successor matrix, we take $\pi$ to be the diffusion policy, and ${\mathit{𝑻}}^{\pi }$ is therefore the diffusion matrix. This is similar to how the adjacency matrix serves as a global structural prior in the STA.

In Figure 5, we compare representations to an agent that computes a value function in a state space consisting of space and time. This agent uses dynamic programming to compute the value of being at location $s$ at time $t$ under an optimal policy for every combination of $s$ and $t$:

At time $t$, the optimal policy is greedy in the value of accessible locations at time $t+1$. For the decoding analyses in Figure 5, we took the ‘neural representation’ to be the concatenation of (i) the flattened value function ${\mathit{𝒗}}_f\in {ℝ}^{TN}$, (ii) a one-hot representation of the current location, and (iii) a one-hot representation of the current time-within-trial. These three quantities are sufficient to compute an optimal policy.

For each analysis in Figure 3, we ran the STA for 1200 network iterations (24 time constants) with recurrent noise $\xi \sim \mathcal{N}(0,\sigma =2.0)$ and no structural noise added to the weights. To quantify the frequency of convergence to different fixed points, we repeated each simulation 1000 times from a uniform initial representation. A simulation was considered to converge to a particular reference representation if the appropriate location in each subspace had a final firing rate of at least 0.9. For the analysis in Figure 3F (bottom), we set the initial neural potential in subspace $\delta =2$ to $z_{2i}=0$ for the bottom-right location and $z_{2i}=-100$ for all other locations. The remaining subspaces were initialised to uniform representations. For the simulation in Figure 3E, we initialised the neural potentials to the final state from Figure 3B, set all reward inputs to zero, and ran the network dynamics for 1200 iterations.

To decode future locations from neural activity (Figure 5C and D), we used scikit-learn (Pedregosa et al., 2011) to train L2-regularised logistic regression models that predicted location at time $t_L$ from neural activity at time $t_N$ with an inverse regularisation strength of $C=1.0$. We performed this analysis in crossvalidation across locations at time $t_N$. In other words, we trained a decoder on data where the agent was in any state $s_{t_N}\ne s$ to predict all locations at time $t_L$, and we then tested this decoder in trials where the agent was in state $s$ at time $t_N$. We repeated this analysis across all test locations $s$ and averaged performance across the resulting 16 folds. We did this to test whether the RNN had a generalisable representation of future location, rather than an encoding of, for example, current location and neighbouring values. In Figure 5C (left), we plot the performance of decoders trained on $t_N=-1$ to predict location at any $t_L\ge 1$. In Figure 5D, we trained decoders on every pair of $t_N\in [0,5]$ and $t_L\in [0,5]$. We then averaged the performance across every ‘delay’ $t_L-t_N$.

To investigate the generalisation of decoders in Figure 5E, we trained a single decoder using neural activity at time $t_N=1$ to predict location at time $t_L=3$. We then applied the same decoder to neural activity at all times $t_N^{\mathrm{\prime }}$ and quantified how well it predicted location at all times $t_L^{\mathrm{\prime }}$. This analysis was again performed in crossvalidation. A separate decoder was trained while holding out each location $s_1=s$, and then tested only on trials where $t_N^{\mathrm{\prime }}=s$. Performance was averaged across all held-out locations and across 5 independently trained RNNs. Figure 5E shows the time at which the average predictive performance was highest for each $t_N^{\mathrm{\prime }}$. Figure 5—figure supplement 4 shows the full generalisation behaviour of the decoder.

To predict the time at which the agent would be at a particular location in Figure 5C (right), we analysed every state $s$ separately, and then averaged over all $s$. For each $s$, we identified trials where the RNN passed through $s$ exactly once. We trained a decoder to predict the time at which $s$ was visited and computed the test performance as a function of the true time at which $s$ was visited. As before, all decoders were trained in crossvalidation across current agent location.

In Figure 5F and G, we compare three classes of RNNs, which were trained on either the reward landscape task, the moving goal task, or the static goal task with a goal that changed between trials. The performance of each RNN was quantified in each of the three tasks as the probability of choosing the optimal first action. We also computed the average parameter magnitudes of all the networks, $|\theta {|}_2^2$. Finally, we computed the average firing rate of each network in the static goal task as ${𝔼}_{\tau \sim \pi }\left[{\sum }_t{\left|{\mathit{𝒓}}_t\right|}_2^2\right]$. All other RNN analyses apart from Figure 5—figure supplement 6 used networks trained on the reward landscape task.

Prior work has investigated the extent to which working memory subspaces are orthogonal (Xie et al., 2022; Dorrell et al., 2024). However, we are primarily interested in the dynamics between subspaces, which are easier to interpret in an orthonormal coordinate system. We therefore asked whether there exists a set of orthogonal subspaces that predict the location of the agent at every time in the future. To do so, we first simulated 6,000,000 trials and collected pairs of neural activity at time $t$ and location at $t^{\mathrm{\prime }}(\delta )$ separately for each $\delta$. We did this in two different ways. To estimate ‘planning’ subspaces, subspace $\delta$ was defined as a decoder that predicts location at time $t^{\mathrm{\prime }}(\delta )=\delta$ from neural activity at times $t\in \{-2,-1\}$. To estimate ‘execution’ subspaces, subspace $\delta$ was defined as a decoder that predicts location at $t^{\mathrm{\prime }}(\delta )=t+\delta$ from neural activity at any $t\ge 0$.

We then defined a predictive distribution parametrised by ${\varphi }_{\delta }=\left\{{\mathit{𝑪}}_{\delta }\in {ℝ}^{N\times N_{\text{rec}}},{\mathit{𝒃}}_{\delta }\in {ℝ}^N\right\}$ for each $\delta$:

Finally, we minimised an objective function that combines the accuracy across δs and the overlap between subspaces:

${\mathit{𝒄}}_{{\delta }_{k}i}$ indicates the parameter vector that predicts being at location $s_i$ at a delay of ${\delta }_k$. $(1-{\delta }_{ij}{\delta }_{{\delta }_k{\delta }_l})=0$ if both $i=j$ and ${\delta }_k={\delta }_l$, and 1 otherwise. The parameters were optimised using ADAM with a learning rate of $5\times {10}^{-3}$ until convergence or for a maximum of 2000 iterations. We used ${\alpha }_1={10}^{-4}$, ${\alpha }_2={10}^{-3}$, and ${\alpha }_{\text{orth}}$ was annealed from 0 to $2\times {10}^{-3}$ over 500 iterations. These hyperparameters were chosen because they resulted in a good approximation to the ‘true’ subspaces in the handcrafted STA.

For the analyses in Figure 5—figure supplement 5, we considered neural activity from our 6,000,000 simulated trials separately for each time-within-trial. We then performed PCA on neural activity across all trials for each time point. We analysed both how much variance was explained by the top 80 PCs, and how much variance was explained by each of the five ‘planning’ and ‘execution’ subspaces defined by ${\widehat{\mathit{𝒄}}}_{\delta i}$ for $\delta =[0,1,2,3,4]$.

To compute the effective connectivity between representations of different points in spacetime, we projected the learned network parameters into a coordinate system defined by the parameter vectors $\mathit{𝑪}\in {ℝ}^{NT\times N_{\text{rec}}}$. Each row of $\mathit{𝑪}$ is a normalised vector ${\widehat{\mathit{𝒄}}}_{\delta i}$ that predicts a particular point in spacetime. In this coordinate system, the input weights are given by ${\mathit{𝑾}}_{\text{in}}^{\text{eff}}=\mathit{𝑪}{\mathit{𝑾}}_{\text{in}}$; the output weights by ${\mathit{𝑾}}_{\text{out}}^{\text{eff}}={\mathit{𝑾}}_{\text{out}}{\mathit{𝑪}}^T$; and the recurrent weights by ${\mathit{𝑾}}_{\text{rec}}^{\text{eff}}=\mathit{𝑪}{\mathit{𝑾}}_{\text{rec}}{\mathit{𝑪}}^T$. To analyse weights between ‘adjacent’ subspaces, we averaged the blocks of ${\mathit{𝑾}}_{\text{rec}}^{\text{eff}}$ that corresponded to weights from any subspace $\delta$ to $\delta +1$ and from any subspace $\delta$ to $\delta -1$. To avoid our analyses being biased by the fact that the networks were trained with supervised learning to predict optimal locations that could only be adjacent, we did not include the weights between subspaces 0 and 1 in this average.

To quantify the similarity between the recurrent weights in this spacetime coordinate system and different order adjacency matrices for the environment, we computed point-biserial correlations. The Δ^th order adjacency matrix ${\mathit{𝑨}}_{\mathrm{\Delta }}\in {ℝ}^{N\times N}$ was defined as a binary matrix with elements equal to 1 for pairs of states that can be reached from one another in Δ actions, and 0 for all other pairs of states. The zeroth-order adjacency matrix was defined as the identity matrix.

For the perturbation analyses in Figure 6F–H, we constructed an environment where the reward function had elements (i) $R_{ts}=1.0$ for $(t,s)\in \{(0,0),(1,1),(2,2),(3,6),(4,10),(5,10),(6,10)\}$, (ii) $R_{ts}=0.7$ for $(t,s)\in \{(1,4),(2,8),(3,9)\}$, and $R_{ts}=-1.0$ for all other points in spacetime. The RNN reliably converged to a representation of the optimal path. We first ran the network dynamics for 10 environment iterations without perturbation after the end of the normal planning period. We then continued to run the network dynamics for 10 environment iterations with a bias term defined by ${\mathit{𝒃}}_{\text{rec}}^{\text{stim}}={\mathit{𝒃}}_{\text{rec}}+\alpha {\widehat{\mathit{𝒄}}}_{2,8}$, where $\alpha$ is the stimulation strength. The perturbed representation was defined as the representation at the end of this perturbation period. The two example representations in Figure 6G used $\alpha =0.3$ (‘weak’) and $\alpha =10$ (‘strong’). The quantification in Figure 6G used a range of $\alpha$ from 0 to 10 for the RNN, and from 0 to 500 for the handcrafted STA. The control analyses (grey lines in Figure 6G, Figure 6—figure supplement 2C) were performed by running the same analysis on the RNNs, but with stimulation of the same magnitude in a random direction in neural state space. For the analysis in Figure 6H, we also removed the perturbation and ran the network dynamics for 10 environment iterations after the end of the perturbation period. All analyses in the main text focused on representations in the space of implied future trajectories, $p_{{\varphi }_{\delta }}(s_{t^{\mathrm{\prime }}}=s_i)$. The ‘representational change’ was quantified as the L1 norm of the difference from the spacetime representation at the end of the normal planning period. See Figure 6—figure supplement 2 for an analysis of the raw firing rates.

For the analyses in Figure 7, we trained another set of RNNs in a version of the reward landscape task where the transition structure changed between trials. The locations of all walls in the environment were provided as a binary input to the agent. For the performance comparison in Figure 7B (top), we evaluated the performance of these networks in the environments that the ‘single structure’ networks had been trained on.

The effective connectivity in Figure 7C and D was computed as in the RNNs trained on a single structure. For these analyses, we identified the subspaces from 1,000,000 trials in a single environment, and we repeated this analysis across 30 different environments. We computed similarities between (i) the effective connectivity estimated in an environment and the adjacency matrix of the same environment, and (ii) the effective connectivity and the adjacency matrix from a different control environment. In Figure 7E, the ‘subspace similarity’ was defined as the correlation between the set of vectors that define the subspaces, averaged over all points in spacetime. We computed the similarity between (i) subspaces identified from two sets of independent trials from the same maze, and (ii) the same number of trials from two different mazes. Recall that the effective recurrent weights between subspaces are given by ${\mathit{𝑾}}_{\text{rec}}^{\text{eff}}=\mathit{𝑪}{\mathit{𝑾}}_{\text{rec}}{\mathit{𝑪}}^T$. By using different subspaces in different environments, the RNN changes $\mathit{𝑪}$ between environments, which changes the effective connectivity between pairs of future subspaces (Figure 7E).

For the analyses in Figure 7G and H, we repeated the subspace identification procedure, but with two notable differences. First, we used trials across many environments to find generalised directions in neural state space that predict the future in any environment. Second, we defined the objective function in terms of future transitions ${\tau }_{\delta }^{ij}$ instead of locations. Figure 7G quantifies the effective connectivity between future transitions in consecutive subspaces (${\tau }_{\delta }^{ij}$ and ${\tau }_{\delta +1}^{kl}$) that are either ‘consistent’ ($k=j$), ‘adjacent’ ($j\ne k$ but $s_i$ and $s_k$ are adjacent in an environment with no walls), or any other transitions. In Figure 7H, we first projected the input corresponding to a given wall location $w_{ij}$ onto each subspace and normalised the projection within each subspace. We then calculated the dot product between this projection and the normalised directions in neural state space that predict either (i) transitions between $s_i$ and $s_j$, (ii) transitions from $s_i$ or $s_j$ to some other state $s_k$, or (iii) transitions that do not originate at $s_i$ or $s_j$. Projection magnitudes were averaged over transitions within each of these three groups, then across subspaces, and finally the mean and standard deviation were computed across 5 RNNs.