P REDICTION AND G ENERALISATION O VER D IRECTED A CTIONS BY G RID C ELLS
A BSTRACT
Knowing how the effects of directed actions generalise to new situations (e.g. moving North, South East and West
or turning left right , etc.) is key to rapid generalisation across new situations. Markovian tasks can be characterised by a state space and a transition matrix and recent work has proposed that neural grid codes provide an efﬁcient representation of the state space, as eigenvectors of a transition matrix reﬂecting diffusion across states, that allows efﬁcient prediction of future state distributions.
Here we extend the eigenbasis prediction model, utilising tools from Fourier analysis, to prediction over arbitrary translation-invariant directed transition structures (i.e. displacement and diffusion), showing that a single set of eigenvectors can support predictions over arbitrary directed actions via action-speciﬁc eigenvalues.
We show how to deﬁne a "sense of direction" to combine actions to reach a target state (ignoring task-speciﬁc deviations from translation-invariance), and demonstrate that adding the Fourier representations to a deep Q network aids policy learning in continuous control tasks.
We show the equivalence between the generalised prediction framework and traditional models of grid cell ﬁring driven by self-motion to perform path integration, either using oscillatory interference (via Fourier components as velocity-controlled oscillators) or continuous attractor networks (via analysis of the update dynamics).
We thus provide a unifying framework for the role of the grid system in predictive planning, sense of direction and path integration: supporting generalisable inference over directed actions across different tasks.
I NTRODUCTION
A "cognitive map" encodes relations between objects and supports ﬂexible planning (Tolman [40]), with hippocampal place cells and entorhinal cortical grid cells thought to instantiate such a map (O’Keefe and Dostrovsky [32]; Hafting et al.
Each place cell ﬁres when the animal is near a speciﬁc location, whereas each grid cell ﬁres periodically when the animal enters a number of locations arranged in a triangular grid across the environment.
Together, this system could support representation and ﬂexible planning in state spaces where common transition structure is preserved across states and tasks, affording generalisation and inference, e.g., in spatial navigation where
Euclidean transition rules are ubiquitous (Whittington et al.
Recent work suggests that place cell ﬁring provides a local representation of state occupancy, while grid cells comprise an eigenbasis of place cell ﬁring covariance (Dordek et al.
[15]; Stachenfeld et al.
[38]; Sorscher et al.
[37]; Kropff and Treves [26]).
Accordingly, grid cell ﬁring patterns could be learned as eigenvectors of a symmetric (diffusive) transition matrix over state space, providing a basis set enabling prediction of occupancy distributions over future states.
This “intuitive planning" operates by replacing multiplication of state representations by the transition matrix with multiplication of each basis vector by the corresponding eigenvalue (Baram et al.
[2]; Corneil and Gerstner [13]).
Thus a distribution over state space represented as a weighted sum of eigenvectors can be updated by re-weighting each eigenvector by its eigenvalue to predict future state occupancy.
Fast prediction and inference of the common effects of actions across different environments is important for survival.
Intuitive planning, in its original form, supports such ability under a single transition structure, most often corresponding to symmetrical diffusion (Baram et al.
Here we show that a single (Fourier) eigenbasis allows representation and prediction under the many different directed transition structures corresponding to different “translation invariant" actions (whose effects are the same across states, such as moving North or South or left or right in an open environment), with predictions under different actions achieved by action-speciﬁc eigenvalues.
We deﬁne a “sense of direction" quantity, i.e., the optimal combinations of directed actions that most likely lead to the goal, based on the underlying translation-invariant transition structure (e.g., ignoring local obstacles).
We then show how this method could be adapted to support planning in tasks that violate translation invariance (e.g. with local obstacles), and show how adding these Fourier representations to a deep
RL network improves performance in a continuous control task.
We propose that the medial entorhinal grid cells support this planning function, as linear combinations of Fourier eigenvectors and therefore eigenvectors themselves, and show how traditional models of grid cells performing path integration are consistent with prediction under directed actions.
Hence we demonstrate that the proposed spectral model acts as a unifying theoretical framework for understanding grid cell ﬁring.
“I NTUITIVE P LANNING " WITH A S INGLE T RANSITION S TRUCTURE
Intuitive planning represents the occupancy distribution over the state space as a weighted sum of the eigenvectors of a single transition matrix (usually corresponding to symmetric diffusion), so that the effect of one step of the transition dynamics on the distribution can be predicted by reweighting each of the eigenvectors by the corresponding eigenvalue.
And this generalises to calculating the cumulative effect of discounted future transitions (Baram et al.
Speciﬁcally, consider a transition matrix, T ∈ R N × N , T ss (cid:48)
=
P ( s t +1 =
s (cid:48) | s t = s ) where s t encodes the state at time t and N is the number of states.
Then, T n is the n -step transition matrix, and has the same set of eigenvectors as T .
Speciﬁcally, the eigendecomposition of T and T n are:
[Equation]
where each column of the matrix Q is an eigenvector of T and Λ = diag ( σ P ( T )) , where σ P ( T ) is the set of eigenvalues of T .
Similarly, any polynomial in T , p ( T ) , shares the same set of eigenvectors as T and the set of eigenvalues σ P ( p ( T ))
=
p ( σ P ( T )) .
Hence:
[Equation]
The resolvent form (Eq. 2) is an inﬁnite discounted summation of transitions, which under a policy and transition structure corresponding to diffusion, is equivalent to the successor representation (SR, Fig.
1E) with discounting factor γ (Dayan [14]; Stachenfeld et al.
See Mahadevan and
Maggioni [29] for a related spectral approach using Fourier decomposition of T for estimating the value function.
The SR has been shown to be useful for navigation via gradient ascent of the future probability of occupying the target state, and has a linear relationship with the true underlying
Euclidean distances in spatial tasks (hence "intuitive planning", see Fig.
1 and Fig.
2D-E).
[Figure] [Equation] Example transition matrix [Equation]
B
[Equation]
P ( s t +
[Equation]
s (cid:48)
[Equation]
s t = 5)
; C
[Equation]
D : same are shown for predicted distribution over the next three time steps
E the form/SR (Eq. 2) computed from the eigenbasis of the transition matrix; F : SR values for state 5 , which can used for navigation.
The eigenvectors of the diffusion transition matrix generally show grid-like patterns, suggesting a close relationship to grid cells.
However, intuitive planning is restricted to predictions over a single transition structure, hence cannot ﬂexibly adjust its predictions corresponding to the effects of arbitrary directed actions (i.e., variable asymmetric transition structure), hence cannot support the presumed role of grid cells in path integration.
Moreover, predictions over different directed actions would require different eigendecompositions, hence incurring high computational costs that undermines its biological plausibility.
In Section 3 we unify the prediction and path integration approaches by exploiting translation invariant symmetries to generalise across actions, using a single common eigenbasis and cheaply calculated updates via action-dependent eigenvalues.
F LEXIBLE PLANNING WITH DIRECTED TRANSITIONS
Updating state representations to predict the consequences of arbitrary directed actions is an important ability of mobile animals, known as path integration and thought to depend on grid cells (McNaughton et al.
To generalise the intuitive planning scheme to simultaneously incorporate arbitrary directed transition structures, we consider the transition dynamics corresponding to translation (drift) and Gaussian diffusion with arbitrary variance (including 0 , equivalent to plain translation).
Our assumption that the transition structure is translation invariant (implying periodic boundary conditions), leads to circulant transition matrices.
Consider a 2D rectangular environment with length L and width W where each state is a node of the unit square grid, then the transition matrix can be represented by T ∈ R LW × LW , with each row the vectorisation ( vec ( · ) ) of the matrix of transition probabilities starting from the speciﬁed location, i.e.,
T [ j, :] =
vec [ P ( s t +1 | s t = j )] , where T is constructed by considering the 2D state space as a 1D vector and concatenating the rows ( j = xL + y for ( x, y ) ∈ [0 , W − 1] × [0 , L − 1] ), see Fig.
2A.
The transition matrix is circulant due to the translation invariance of the transition structure (see
Appendix Prop. A.1), and takes the following form:
[Equation]
The normalised eigenvectors of the circulant matrix T ∈ R N × N ( N = LW ) are the vectors of powers of N th roots of unity (the Fourier modes):
[Equation]
where ω k = exp( 2 πiN k ) , for k = 0 , . .
, N − 1 , and i =
√− 1 .
Hence the matrix of eigenvectors (as the columns), F = ( q 0 , q 1 , . .
, q N − 1 ) , is just the (inverse) discrete Fourier transform matrix (Bracewell [4]), where F kj = ω kj for 0 ≤ k, j ≤ N − 1 .
The Fourier modes projected back onto the
L × W 2D spatial domain are plane waves, as shown in Fig.
2G, with wavevector determined by the value of k that speciﬁes the direction and spatial frequency of each plane wave (see Appendix B).
We can immediately compute the corresponding eigenvalues for the eigenvectors in Eq.
4 (equivalent to taking the discrete Fourier transform (DFT) of the ﬁrst row (or column) of T , see Bracewell [4]):
[Equation]
where { T 0 , . .
, T N − 1 } are the N unique elements that fully speciﬁes the circulant matrix T (Eq. 3).
We can then utilise tools from Fourier analysis for efﬁcient updating of the eigenvalues whilst leaving the universal eigenbasis unaffected.
For a transition matrix T v corresponding to an arbitrary action (translation velocity) v = ( v x , v y ) , each row of T v is again a circulant, but shifted version of the corresponding row vector of the symmetric transition matrix corresponding to zero drift velocity, T 0 . Speciﬁcally, the ﬁrst rows of the two matrices are related as follows: [Equation]
[Figure]
i t i
[Figure]
e
[Figure]
i [Figure] Given the eigenvalues for T 0 , Λ 0 = (cid:2) λ 00 , λ 01 , . . .
, λ 0 N − 1 (cid:3) ∈
C N (via the DFT of the ﬁrst row of T 0 ,
Eq. 5) we can immediately derive the eigenvalues of T v , Λ v , via a one-step update based on the
Fourier shift theorem (Bracewell [4]) without recomputing the eigendecomposition:
[Equation]
for k
[Equation]
[Equation]
[Equation]
for arbitrary v ,
[Equation]
This allows path integration by reweighting the common set of eigenvectors at each timestep by the updated eigenvalues corresponding to the current drift velocity (Eq. 7).
Note that additionally, T 0 can include diffusion, thus reweighting by the eigenvalues of the diffusive transition matrix also allows tracking of increasing uncertainty.
Utilising the ﬁxed eigenbasis (Eq. 4) and the respective eigenvalues (Eq. 7) for arbitrary transition structures, we can make efﬁcient prediction for the distribution of future state occupancy with respect to arbitrary action (see Figs.
2B-C).
Adding translation to the translation-invariant transition matrix does not change the set of eigenvectors - allowing one set of eigenvectors (Fourier modes) to support prediction for actions in all directions (or plain diffusion), hence prediction of effects of directed actions can be efﬁciently generalised across environments.
Sense of Direction.
We deﬁne a "sense of direction", θ ∗ , as the angle of the transitions (or the linear combinations of the available actions in a non-spatial setting) that maximise the future probability of reaching the target state given an initial state, which is modelled by the SR matrix.
[Equation]
where γ is the discounting factor, D j , j = 1 , · · · , LW are the eigenvalues for the symmetric diffusion transition matrix, k j , j = 1 , . .
. , LW are the wavevectors for the j -th Fourier components, x 0 , x G are the coordinates of the start and goal states, and v θ = ( v cos( θ ) , v sin( θ )) represents the velocity (with speed v and head direction θ ).
We see that the "sense of direction" supports generalisation of predictions of effects of actions across all environments with the same translation-invariant transition structure, i.e., such predicted effects ignore any local deviations from translation invariance.
See
Appendix B for the derivation of Eq. 8.
Note that here we assume that the goal state s G is known a priori, e.g., we consider a problem where the animal is navigating towards a previously visited location.
The derived analytical expression for the sense of direction can be retrieved via a lookup table when the state space is small and discrete, whereas in large or continuous state spaces, it can be computed either via optimisation algorithms, or modelled by a non-linear function approximator that represents Eq. 8.
See Bush et al.
[11] for neural network approaches to ﬁnding goal directions from grid representations.
We thus propose that a computational role for the neural grid codes: generating a “sense of direction" (capturing the transition structure of the state space, ignoring the obstacles and boundaries) that reﬂects a global sense of orientation which allows generalisation to completely new environments.
Flexible Planning & Application Beyond Translation-Invariant Structures.
The proposed model can be applied to ﬂexible planning under arbitrary drift velocity as demonstrated in Fig.
3 (A-E).
An agent is trying to navigate towards a goal state in a windy grid world.
The navigation is performed by following the ascending "gradient" of the SR for occupancy of the target state (the resolvent metric,
Eq. 2).
The SR computed from the transition matrix including the effects of diffusion and wind (Fig.
3 B) based on our analysis (eq. 7) leads straight to the target (Fig. 3 C).
Given the analytical expression of the SR (Eq. 2), we could efﬁciently adjust the SR matrix to accommodate local changes in the state space, e.g., insertion of a barrier, using the Woodbury inversion formula to update the parts of the SR matrix affected by the local obstacles (see Appendix A.3 for derivations [34]); and again in this case, the agent correctly adjusts for the wind as well as taking the shortest path around the inserted wall (Fig.
3 D-E).
We note, however, that the proposed model is also able to solve tasks without periodic boundary conditions, by considering the original task state space, S 0 , being embedded into a larger, periodically bounded pseudo state space S p , at least twice as large in each dimension as S 0 (Fig. 3 F).
We again follow the previous procedures, utilising the Fourier modes, this time computed on S p , to perform predictions in S 0 (Fig.
3 F-G), and the performance is unaffected.
Note that under such formulation, the underlying transition structures can be applied to environments with both periodic and non-periodic boundary conditions - allowing sense of direction planning in either case.
Path Integration.
We can also use our model for path integration (see also Section 4) in S 0 , by taking velocity inputs (given any path in the grid world) to update the state occupancy distribution (Eq. 7).
The path integration performance is strongly correlated with the degree of uncertainty (i.e., the diffusion strength caused by self-motion noise in addition to translations).
This is indeed captured by our model (Fig. 3H), with perfect path integration when the uncertainty is low up totime steps (the discretisation of state space means that uncertainty below 0.075 has no effect ), and monotonically increasing path integration error when the uncertainty is higher.
3. N EURAL I MPLEMENTATION FOR D EEP R EINFORCEMENT L EARNING
Our proposed framework supports prediction and planning under arbitrary direction actions and path integration.
To further demonstrate its utility in non-spatial tasks, we propose GridCell-DQN (gcDQN), a neural network implementation based on a modiﬁed version of the classic Deep-Q Network (DQN, Mnih et al.
The architecture of gc-DQN is designed so that quantities corresponding to combinations of Fourier modes weighted by action-dependent values are explicitly available to action-value computation in addition to value estimates based on the state-inputs alone.
This should enable the network to predict future state occupancy and thus facilitate the learning of Q values.
We evaluated the performance of gc-DQN on the CartPole task (Barto et al.
[3]) and compared with plain
DQN.
We restrict our introduction of the gc-DQN architecture based on the evaluation on the CartPole task.
The state space of the CartPole task is 4 -dimensional, corresponding to the cart position and velocity, pole angle and angular velocity, and there are 2 possible actions ( 0 and 1 corresponding to cart movement left or right).
The Q-values are learnt using the standard temporal-difference rule (Sutton and Barto [39]).
The gc-DQN has 2 additional sub-networks (below the standard DQN in Fig.
4.A):
the ﬁrst takes as inputs the n low-frequency Fourier modes over the state space after discretisation into 8 4 bins and has n actions = 2 outputs to allow representation for each action (left or right) given
A
B
C
D
E
F
G
H 2 2 2 2 2 2 2 2
G
G 2 2 2 2 2 2 2 2
S
G
G 2 2 2 2 2 2 2 2
S
G
S
G
Bounded grid
S
G
Bounded grid[-0 
-]
[Figure]
var=0.04var=0.08var=0.
[Figure]
|| x t x t || the input Fourier modes; the second takes as inputs the state variables and action and has 2 n outputs which serve as action-dependent multipliers to the connection weights from the input layer of the ﬁrst network.
The second sub-network receives state as well as action inputs to mitigate the absence of translation-invariance.
The outputs of the ﬁrst network and of the standard DQN are fully connected to an output layer to learn the updated Q values.
We also evaluated a model-based version of gc-DQN based on deep Dyna agents (Peng et al.
Preliminary results in Fig.
4(B) show that the gc-DQN and deep gc-Dyna-Q greatly accelerates learning comparing to the baseline agents, with relatively minor increase in the model complexity and computational costs.
The results support our hypothesis that Fourier eigenvectors weighted by action-speciﬁc values can aid prediction (in this case, of future value).
See Appendix.
E for details.
The focus of this paper is proposing a theoretical framework for state representation, prediction, planning and path integration via grid-like eigenvectors and action-dependent eigenvalues.
The proposed gc-DQN is only a preliminary attempt towards a neural network implementation of the proposed approach (see also Mahadevan and Maggioni [29]), more rigorous studies in this direction is left for future work.
A UNIFYING FRAMEWORK FOR MODELS OF GRID CELL FIRING
Our focus so far has been on proposing a ﬂexible and efﬁcient extension of the prediction models of grid cells (Baram et al.
[2]; Dordek et al.
[15]; Stachenfeld et al.
[38]) to arbitrary directed transitions.
However, many other computational models of grid cells emphasise path integration rather than inputs from place cells, such as continuous attractor network (CAN) models, in which grid-like patterns emerge in recurrently connected networks performing path integration (Fuhs and Touretzky [18]; Burak and Fiete [6]; Corneil and Gerstner [13]), and oscillatory interference (OI) models in which grid-like patterns reﬂect coincidence detection of velocity-dependent oscillatory inputs during path integration (Hasselmo [22]; Burgess [8]; Welday et al.
Here we build upon the previous work of unifying the prediction and path integration models of grid cells ﬁring by Sorscher et al.
[37], to show the equivalence of our generalised prediction framework with the CAN models of grid cells; we additionally show the equivalence between the proposed model and oscilla- gc-DQN
A
B
I npu t : s t a t e
D Q N _ h
i dd e n _
D Q N _ h
i dd e n _
O u t pu t :
D Q N
I npu t : ( s t a t e , a c t i o n )
S A _ h i dd e n _
O u t pu t : S t a t ea c t i o n m u t li p li e r s
I npu t : n
F o u r
i e r m o d e s
F o u r
i e r _
h
i dd e n _
O u t pu t : a c t i o n w e
i g h t s
O u t pu t : a c t i o n v a l u e s tory interference models in terms of their interpretations of path integration and theta phase precession.
4. R ELATION TO CONTINUOUS ATTRACTOR NETWORK MODELS OF GRID CELLS
One of the most prominent unifying analyses for different grid cell models (Sorscher et al.
[37]) proves the equivalence between maximising a spatial representation objective function under the prediction models and the pattern formation dynamics of CAN models of path integration.
Their analysis, however, does not include non-zero velocity inputs, corresponding to the asymmetric velocity-dependent connectivities in CANs which perform path integration (Fuhs and Touretzky [18]; Burak and Fiete [6]).
We can explicitly address this equivalence using our framework.
Assuming that grid cell ﬁring rate, g , reﬂects linear combinations of selected Fourier modes (e.g., 6 Fourier modes at π 3 radians increments with the same frequencies):
[Equation]
where f j ’s are the selected Fourier modes with weights w j .
Note that our proposed model is on the theory-level, rather than the implementation level of CAN models, hence we do not assume any speciﬁc neural network structure here.
Utilising the grid cell ﬁring described by Eq.
9, followed by similar analysis as in Sorscher et al.
[37], we show that, under non-zero velocities, the differential equations governing the dynamics of the CAN model updates are equivalent to the derivative of the Lagragian equation underlying the optimsation problem of the prediction models (up to scaling factors).
Hence we show that under non-stationary transition dynamics, the CAN models and the prediction models should yield identical updates to grid cell ﬁring (up to scaling factors).
The complete proof can be found in Appendix B.
4. R ELATION TO OSCILLATORY INTERFERENCE MODELS OF GRID CELLS
Another major class of computational models of grid cells is the oscillatory interference model, here we show the equivalence between the generalised prediction model and the OI models of grid cells by showing that they perform path integration via similar phase coding.
In OI models of grid cells, path integration is achieved via the phases of the “velocity controlled oscillators" (VCOs), which encode movement speed and direction by variations in burst ﬁring frequency.
The VCOs generate grid-like ﬁring patterns via coincidence detection by the grid cells
(Burgess [8]; Hasselmo [22]; Welday et al.
The variation of frequency with velocity produces a phase code for displacement, enabling the modelled grid cells to perform path integration.
Namely,
VCOs change their frequencies relative to the baseline according to the projection of current velocity, v ( t ) , onto the VCO’s "preferred direction", d :
[Equation]
where β is a positive constant, and f b ( t ) is the baseline frequency (the 4 − 11 Hz EEG theta rhythm).
It follows that VCOs perform linear path integration since the relative phase between VCO and baseline at time t , φ ab ( t ) =
φ a ( t ) − φ b ( t ) , is proportional to the displacement in the preferred direction:
[Equation]
(11) where x ( t ) is the agent’s location at time t .
The interference of VCOs whose preferred directions differ by multiples of π/ 3 generates grid-like patterns, provides an explanation of “theta phase precession" in grid cells (in which ﬁring phase relative to the theta rhythm encodes distance travelled through the ﬁring ﬁeld;
Hafting et al.
[21]; Burgess [8]) and complements the attractor dynamics given by symmetrical connections between grid cells (Bush and Burgess [9]).
We note that the main experimental results held against OI models (that bats and humans do not have reliable theta frequency oscillations) has recently been resolved: the required phase coding (theta phase precession) can occur relative to a variable baseline frequency (Bush and Burgess [10]) and has now been found in both bats and humans (Eliav et al.
[17]; Qasim et al.
We simulated the ﬁring of a grid cell with 6 Fourier mode inputs (Fig.
5 A-B), each ﬁring a spike at its complex phase in the current state, as a leaky integrate and ﬁre
neuron performing coincidence detection, using a real trajectory of a rat exploring a 50 cm × 50 cm box.
At each time step (corresponding to one theta cycle), the phase of each Fourier mode is updated according to Eq.
7 given the current velocity, and ﬁres a spike at this phase if it is within the interval [ − π/ 4 , π/ 4] (modelling modulation by the baseline theta oscillation).
The grid cell ﬁres a spike at the current location if the integrated inputs reach a threshold.
Note that we could also simulate a set of grid cells, with different offsets (depending on the initial phases of the Fourier modes) and different scales and orientations (depending on the choice of Fourier modes), such that the grid cells, like the Fourier modes, comprise a basis for the state space and do so on the basis of path integration (for which environmental inputs are only required to prevent error accumulation [8; 7]).
Grid cells show "precession" in their ﬁring phase relative to theta as the animal moves through the ﬁring ﬁeld (signaling distance travelled within the ﬁeld; Hafting et al.
[21]; Jeewajee et al.
[23]; Climer et al.
Our model captures this, like an OI model.
The Fourier modes whose wavevectors that are aligned with the current direction of translation advance in phase as the movement progresses.
Phase precession results from assuming that the Fourier modes aligned with movement direction are the dominant inﬂuence on grid cell ﬁring (c.f. those aligned to the reverse direction).
The baseline “theta frequency" corresponds to the mean rate of change of phase of all Fourier components and so could vary (e.g., for noise reduction, see Burgess and Burgess [7]; Burgess [8]), without precluding phase coding (Eliav et al.
[17]; Bush and Burgess [10]).
By simulating straight runs, we can see clear late-to-early phase precession (Fig.
5 C), as observed in grid cells.
Thus, the OI model and our model perform path integration or prediction in the same way: the phase of each VCO changes corresponding to the component of translation along the VCO’s preferred direction, which is exactly analogous to the complex phases of the Fourier modes being updated to reﬂect transitions along their wavevectors (multiplication by corresponding eigenvalues, Eq. 7).
D ISCUSSION
Understanding how different actions affect the agent’s state across environments is essential for generalisation.
Existing models are capable of such prediction under a single ﬁxed transition matrix, e.g. corresponding to symmetrical diffusion, by using eigenvectors of this transition matrix as a basis for representing state occupancy (Baram et al.
[2]; Corneil and Gerstner [13]; Stachenfeld et al.
Here we generalised these models to provide a mathematical framework for predicting the effects of speciﬁc actions, so long as their effects (and corresponding transition matrices) are translation invariant.
This uses a common set of eigenvectors of all such matrices (Fourier modes of the state space) to represent state occupancy, so that the effects of actions correspond to multiplication by action-speciﬁc eigenvectors.
This model explains how grid cells (as superpositions of Fourier modes) could support prediction of the effects of actions across environments that share the same underlying transition structure (see also Whittington et al.
[43]), and could, for example, perform sense-of-direction planning in new environments (i.e., ﬁnding combinations of actions that most likely lead to the target state by ignoring local obstacles).
We assume that other (e.g., fronto-parietal) brain areas are responsible for detecting and avoiding obstacles following the overall direction provided by the grid cells [16; 28].
However, topology-dependent modiﬁcations to grid ﬁring patterns could be used to accommodate local deviations from translation-invariance (e.g. obstacles), utilising the Woodbury inversion formulae, see
Fig.
3E, Appendix A3 and Piray and Daw [34].
We also show show our framework corresponds to other computational models of grid cells based on path integration, and provide a functional explanation for theta phase precession.
A number of questions and predictions are raised by the proposed model.
If a basis of neurons with
Fourier-mode-like ﬁring patterns act as inputs to cells in entorhinal cortex, then grid cell ﬁring patterns are only a small proportion of the set of ﬁring patterns that could be synthesised.
This is consistent with the existence of periodic non-grid cells in entorhinal cortex that resemble combinations of small numbers of Fourier modes (Krupic et al.
We predict the use of the same set of grid cells (superpositions of Fourier modes inputs) for indicating “sense of direction" to goal locations across different environments after a single visit, i.e. showing generalisation across Euclidean environments.
Finally, the proposed model predicts future state occupancy from the transition matrix, future work could also consider the reversed direction: inferring the translation between two locations given the phase codes for each (see Bush et al.
[11]), as linked discriminative and generative models.
The current model applies to translation-invariant transition structures, and use of the Fourier shift theorem to calculate eigenvalues also assumes a Euclidean state space.
We demonstrated ways of generalising planning to bounded or locally non-translation invariant transition structures in Section 3.
We note that machine learning methods based on a similar premise (creating a single representation to support planning via multiple different actions) might work even when the transitions are not strictly translation invariant (e.g., family trees, see Whittington et al.
Here we showed that, by giving standard DQN agents the ability to represent the current state as eigenvectors of the state space weighted by state- and action-speciﬁc values, signiﬁcantly improves learning of the CartPole task (Fig. 4), which is not strictly spatial or translation invariant.
Hence our approach offers potential generalisation to non-spatial tasks, such as transitive inference (Von Fersen et al.
[41], Appendix D).
Future work will involve more rigorous study of the neural network implementation of the proposed grid cell model and its application to reinforcement learning.
A future direction in generalising the current model to non-spatial tasks will be to consider Fourier analysis on groups of operators utilising group-theoretic knowledge (Kondor and Trivedi [25]; Gao et al.
A CKNOWLEDGEMENTS
C.Y. thanks a DeepMind PhD studentship offered by the UCL Centre for Artiﬁcial Intelligence.
T.E.J.B and N.B. thank the Wellcome Trust for support.
The authors would like to thank Daniel Bush,
Talfan Evans, Kimberly Stachenfeld, Will de Cothi and Maneesh Sahani for helpful comments and discussions.
The authors declare no competing ﬁnancial interests
A S OME PROOFS IN S ECTIONS
Proposition A.1 Given our assumption of periodic boundary condition, the transition matrix, T ∈ C N × N
(Eq. 3), is indeed a circulant matrix.
Proof
It is easy to see that the proposition holds trivially for transition matrices with only one-step translations but without Gaussian spread.
Hence here we only show the proof for the case where the transition structure includes both Gaussian spread and one-step translations.
Consider for an arbitrary transition matrix T for a 2D rectangular environment with length L and width W , and the underlying transition velocity is v = ( v x , v y ) , remembering that T is the LW × LW 2D matrix formed from concatenating rows from what would be the 4D matrix of transitions between all pairs of states in a L × W 2D state space.
An arbitrary entry on the k th lower subdiagonal is
T i,i − k
= P ( x ( t + 1)
=
i − k | x ( t ) = i ) for any suitable state i given k (i.e., i ≥ k ).
If the Gaussian spread is radially symmetric with constant variance across states, the value of T i,i − k only depends on the distance between state i − k and the state i v , where i v is the translated state of state i given the effect of the velocity v .
The states i − k
and i v are equivalent to the states (( i − k ) //L
, ( i − k ) mod L ) and ( i//L + v x , i mod L + v y ) in the two-dimensional spatial domain respectively (where a//b denotes the integer part of a/b ).
Note that we need to have the velocity v ∈ [ ± L/ 2 , ± W/ 2] so that the translation leaves the actual distance d unchanged.
The distance between the state i − k and the expected state i v in the 2D state space is then
[Equation]
For any arbitrary i (cid:48)
[Equation]
=
i such that i (cid:48)
=
i + m , we could compute similarly the distance between state i (cid:48) − k and its corresponding expected state (Gaussian center) i (cid:48)
+ v .
After some algebra, we have that the distance between states i − k
and i + v equals the distance between states i (cid:48)
− k and i (cid:48)
+ v .
[Equation]
Now if we look at the two square terms within the square root separately, we have
[Equation]
[Equation]
The second equality holds due to the fact that ( i − k ) mod L + δ mod L ≤ L since this is simply the x -position of state i + δ − k , which is never larger than L , hence (( i − k ) mod L + δ mod L ) mod L
=
( i − k ) mod L + δ mod L .
Similarly ( i mod L + δ mod L ) mod L = i mod L + δ mod L .
Hence we have
[Equation]
Hence all entries on the k th lower subdiagonal are identical, i.e. T i,i − k
= T i (cid:48) ,i (cid:48)
− k for all 1 ≤ k ≤
LW − 1 .
And by similar arguments, we could show that all entries on the k th upper subdiagonal are identical (for 1 ≤ k ≤ LW − 1 ), and equals to the corresponding entries on the LW − k th lower subdiagonals.
And the fact that all the main diagonal entries are identical is immediate from the problem setting.
Hence our target transition matrix is indeed a circulant matrix.
(Note that in simulations the transition matrix will only be approximately circulant due to normalisation and numerical issues.)
Now we consider the corresponding ( LW − k )
th upper subdiagonal (to the k th lower subdiagonal), by similar arguments, we have that for any T
i (cid:48)(cid:48) ,i (cid:48)(cid:48)
+ k = P ( x ( t + 1)
=
i (cid:48)(cid:48)
+ k | x ( t )
=
i (cid:48)(cid:48) ) for suitable i (cid:48)(cid:48)
(i.e. i (cid:48)(cid:48)
+ k ≤ L ), the distance between the state i (cid:48)(cid:48)
+ k and expected next state
i (cid:48)(cid:48)
+ v t are the same as d , which is equivalent to T i (cid:48)(cid:48) ,i (cid:48)(cid:48)
+ k = T i,k − i .
Hence all entries on the ( LW − k )
th upper subdiagonal are identical and equal to the entries on the k th lower subdiagonal.
This holds for arbitrary 1 ≤ k ≤ LW − 1 .
[Equation]
Proposition A.2 For any circulant matrix T ∈ C N × N as shown in Eq.
3, its k th eigenvector takes the form:
[Equation]
where ω k = exp (cid:0) 2 πkiN (cid:1) is the k th N th root of unity, and the set of eigenvalues equals to the set of
DFTs of an arbitrary row/column of T .
Proof
Firstly, note that the product between the circulant matrix T and an arbitrary vector v is equivalent to a convolution.
[Equation]
And we immediately have that
[Equation]
This is true due to the periodicity of the entries given by the circulant structure.
Then if we take the dot product of T and an arbitrary vector v m of the form shown in Eq.
17, the l th entry of the output vector has the following form.
[Equation]
where the equality holds since ω
[Equation]
Note that the last sum in Eq.
20 is independent of the choice of l since both T j and ω j are periodic hence any change in l is simply rearranging the terms in the summation.
Also we have that ω
ml
=
ω lm is the l th entry of the eigenvector v m .
Hence we have
[Equation]
21) where
[Equation]
for m = 0 , . . .
, N − 1 .
Hence for an arbitrary N × N circulant matrix T , the eigenvalues take the form as shown in Eq.
and the corresponding eigenvectors take the form as shown in Eq.
17, and the eigenvalues are equivalent to the DFT of the ﬁrst row of the circulant matrix immediately follows from Eq.
22 and the deﬁnition of DFT (Bracewell [4]).
[Equation]
The predicted phase change in the eigenvalues over the eigenvalues of the baseline symmetric transition matrix computed with Fourier modes computed via Fourier shift theorem (Eq. 7) under our formulation perfectly captures the actual phase changes caused by the one-step translations in the eigenvalues between the symmetric and and asymmetric transition matrices, as shown in Fig.
6A.
However, when the transition dynamics is a combination of diffusion and one-step translations, the predicted phase changes in eigenvalues will no longer perfectly match the actual phase changes observed as shown in Fig.
6B, and the oscillation is caused by the diffusion process.
Namely, although the expected translation is indicated by the velocity, the actual translation spans a range of states depending on the width of the diffusion ﬁeld.
[Figure]
Proposition A.3 The updated SR given the insertion of a barrier is
[Equation]
where S 0 and S are the initial and updated SR, R = S
[ J, :] and C = S
[: , J ] are the J -th rows and columns of S 0 respectively, where J is the index set of states adjacent to the inserted barrier.
Proof
This derivation is inspired by Piray and Daw [34].
Given the deﬁnition of the SR, we have
[Equation] [Equation]
where N is the number of states.
Given the insertion of a barrier, S and S 0 only differ in their j -th rows for j ∈ J where J is the index set of states adjacent to the barrier.
Hence we could write:
[Equation]
Then if we have E ∈ R | J |× N with zeros everywhere but ones on the j -th rows for j ∈ J , then by setting W = I − T and W
= I − T 0 , we could write:
[Equation]
[Equation] The Woodbury inversion formula is usually use in cases whn we are trying to compute the inverse of a matrix given a low-dimensional perturbution (Riedel [36]).
[Equation]
Hence by applying the Woodbury inversion formula, we have:
[Equation]
where C = ES 0 are the j -th columns of S 0 for j ∈ J .
[Equation]
Proposition A.4 The "sense of direction", θ ∗ , is given by the form shown in Eq. 8.
Proof
Essentially, we wish to ﬁnd value of θ such that under the drift velocity v θ = ( v cos( θ ) , v sin( θ )) , given the start and target states, s 0 and s G , the future discounted occupancy of s G starting from s 0 (or W [ s 0 , s G ] , where W is the SR matrix), is maximised.
Under our formulation,
W can be calculated as follows:
[Equation]
where F is the DFT matrix (Eq. 4), and Λ v θ is the set of eigenvalues of the transition matrix given velocity v θ .
From our analysis based on Fourier shift theorem (Eq. 7), for each λ v θ i ∈ Λ v θ , we have that:
[Equation]
where D i is the i th eigenvalue of the symmetric (baseline) diffusion transition matrix, and k i is the wavevector for the
i th Fourier mode.
Then using linear algebra, we immediately arrive at the expression in Eq. 8.
[Equation]
B S OME PROOFS IN S ECTION
Proposition B.1 The equations governing the dynamics of the prediction model and the CAN model of path integration are equivalent.
Proof
We show the proof under the single-cell formulation, which can be immediately generalised to the situation with multiple cells.
We ﬁrstly note that the prediction model can be mathematically categorised as minimising the following reconstruction objective function.
[Equation]
where g ∈ R n G represents the grid cells ﬁring rates, and w ∈ R 1 × N represents the linear readout weights.
Following Sorscher et al.
[37], we replace w by its optimal value given a ﬁxed g , i.e., ˆ w = ( g T g )
− 1 g T T .
Note that any scaling of g can be absorbed into a corresponding reversed scaling into ˆ w , hence g is assumed to be of unit modulus (or the matrix G can be taken to be orthonormal in the multi-cell case).
Additionally, following the non-negativity constraint proposed in Dordek et al.
[15], the overall optimisation problem becomes.
[Equation]
Hence we can immediately write down the Lagrangian as follows.
[Equation]
where γ and µ are the multiplicative constant for the additive penalty terms corresponding to the constraints in Eq.
The derivate of the Lagrangian with respect to g then takes the following form.
[Equation]
[Equation]
where σ ( · ) is the rectiﬁed linear function.
Inserting the grid cell ﬁring representation as a linear summation of the Fourier modes into Eq.
34, we obtain the following.
[Equation]
where λ j are the corresponding eigenvalue of f j with respect to the (symmetric) transition matrix,
T 0 .
The dynamics of the grid cells under the CAN models can be written as following.
[Equation]
36) where W is the recurrent connectivity matrix, b ( v ) is the velocity-dependent feedforward input to the grid cell under the CAN model which involves a constant baseline term and a velocity dependent term (Burak and Fiete [6]).
Now suppose the agent is moving under non-zero velocity, v .
Given the grid cell ﬁring represented by the linear summation of Fourier modes (Eq. 9), Eq. 35 can be written as following,
[Equation]
where e j is the unit-norm wavevector of the Fourier mode f j for all j .
Now check with Eq.
36 by setting
[Equation] [Equation]
By checking that when g > {+0 , σ ( 2 πiN γ T 0 g + σ ( 2 πiN γ (cid:80) Gj =
+}
λ j w j f j ( (cid:104) v, ˆ e j (cid:105)− 1)+ µ ))
= 2 πiN γ T 0 g + σ ( 2 πiN γ (cid:80) Gj =1 λ j w j f j ( (cid:104) v, ˆ e j (cid:105)− 1)+ µ ) , we see that under non-zero velocity inputs, by appropriately adjusting the additive velocity input term, b ( v ) , the equations governing the dynamics for the normative and mechanistic models are equivalent.
[Equation]
C 2D F OURIER MODES
We know that the Fourier basis vectors from Eq.
4 form plane waves as shown in Fig.
From standard Fourier analysis in 2D space, the 2D Fourier modes form an orthonormal basis, and takes the following form.
[Equation]
where the 2D Fourier basis vectors are encoded by the position vectors
u = ( u 1 /L, u 2 /W ) ∈ [0 , 1] ×
[0 , 1] (position vectors of each location in the L × W environment projected onto [0 , 1] × [0 , 1] ).
The direction of the encoder position vector u represents the direction of the plane wave and the frequency of the plane wave is the unnormalised direction vector
|| u (cid:48)
|| (where u (cid:48) =
u × ( L, W ) ), note that u (cid:48) is also the wavevector for the plane wave.
This is a slightly different formulation comparing to the formulation given in Eq. 4, which consider the state space as a 1-dimensional ﬂattened vector of the 2-dimensional environment, hence the Fourier basis vectors are the corresponding 1-dimensional
Fourier modes.
Though both formulation give us the same set of Fourier basis vectors, under the deﬁnition in Eq.
39, we could easily track the frequency and direction of the plane wave formed from the 2D Fourier modes.
And the phase shift via the Fourier shift theorem 7 equivalently applies for this 2D Fourier formulation.
The Fourier modes comprises a basis for representing any distribution over the task state space, so we could use a linearly weighted combination of Fourier modes to reconstruct any ﬁring patterns, such as those observed in place cells (Welday et al.
[42], Fig.
However note that the coincidence detection of small number of oscillators with different frequencies will generate periodic patterns, e.g., grid cells, and more oscillators will be needed for those with more local ﬁring ﬁelds such as place cells.
Note that the total number of Fourier modes equals the number of states in the environment (e.g., LW for the L × W rectangular environment on a square grid), and it could be infeasible and inefﬁcient to compute and store a large number of such Fourier modes (or neurons with VCO-like ﬁring patterns) in the brain.
Hence here we only use the principal modes (taking the top n Fourier basis vectors in terms of the corresponding eigenvalues (frequencies)), within contain the majority of the information is contained, with the number of principal modes depending on the desired reconstruction resolution.
We utilised the top 100 principal Fourier modes for most of the simulations in the main text (see
Fig.
7 for a typical ﬁxed set of Fourier modes).
Fig.
8 demonstrates that the small number of Fourier modes are able to reconstruct grid cells ﬁring ﬁelds with various spacings and orientations, and place cells ﬁring ﬁelds.
D T RANSITIVE I NFERENCE
In the main paper we argued that the same set of eigenvectors can be used to predict future occupancy distribution given the transition matrix for symmetrical relations like diffusion between adjacent states and directed transitions (e.g. moving N S E W).
Here we brieﬂy discuss the generalisation of our model to non-spatial tasks.
We could apply our method to the one-dimensional transitive inference tasks of this type.
e.g., given
A > B , B > C , C > D , then infer if A > D (Von Fersen et al.
This would be like having a 1D track (and Fourier eigenvectors for 1-step transitions in both directions) corresponding to actions "greater" or "smaller", and using "intuitive planning" to see if using eigenvalues for "greater" will take you from A to B in the discounted future more likely than eigenvalues for "smaller".
In order to deal with the non-periodicity of the task, we simulate transitive inference in a small subset of the state space of the torus.
As shown in Fig.
9, we see that our framework correctly predicts the transitive relationship between the chosen state x 129 and states close to x 129 .
Despite the simplicity of 1D transitive inference (Fig. 9), our model is still an advance on the original intuitive planning method in being able to predict the effects of both "greater" and "smaller" transitions with the same set of eigenvectors, rather than being restricted to prediction with one or the other alone.
E S IMULATIONS E. F URTHER DETAILS OF THE GC -DQN AGENT
The overall architecture can be found in the graphical illustration in ﬁg. 4.
At each timestep, the state values and a speciﬁc action value are fed into a neural network for all possible values of actions (Blue box in the bottom left of ﬁg. 4), which outputs n actions output, where n is the number of
Fourier modes inputs to the second network.
The output can considered as the speciﬁc updates to each Fourier modes corresponding to the action in the current location, like the phase shift in the
Fourier shift theorem.
The inputs to the grid cell network are the ﬁrst n Fourier modes, whose dimensions ( D ) are determined by the size of the state space.
When the state variables are continuous, we compute an approximate size of the state space by discretising each state variable.
The number of principal Fourier modes ( n ) is chosen arbitrarily as long as the majority of the information can be reconstructed from the chosen set of Fourier modes.
Higher values of n leads to ﬁner details of the prediction, but also induces higher computational costs.
[Figure] [Figure] At each timestep, the n Fourier modes is fed as the input to the grid cell network (shown in the middle row of ﬁg. 4(A)).
Each action multiplier (outputs from the state-action network) is multiplied with the corresponding column of the weight matrix between the input layer and the hidden layer of the grid cell network.
The outputs of the hidden layer is then transposed, and forward propagate to the output layer of the second network.
The computations of the grid cell network is considered to be equivalent to using the Fourier modes to construct a weight value for choosing each action at a given state that aids navigation/planning.
The outputs from the grid cell network and the standard DQN agent is then combined to output a vector, that acts as the values for each action that guides action choice in the current timestep.
E. S IMULATION DETAILS
All simulations were implemented in Python.
The simulation details for each task is as follows:
• Fig.
1: The state space is assumed to be a 1 D ring with 20 states, with the transition probabilities P ( s t
+1 =
i + 1 | s t = i ) =
P ( s t +
=
i − 1 | s t = i )
= 0 .
5 , and discounting factor γ = 0 .
9 for generating the resolvent (Eq. 2).
• Fig.
2: Variance of each (Gaussian) ﬁring ﬁeld (representing the strength of diffusion) is 3 ; B : (0 , 5) drift velocity with increasing diffusion (variance increase by 3 per step);
C : (3 , 3) drift velocity with 0 diffusion; E : The successor representation is computed using the Fourier modes and corresponding eigenvalues, with the discounting factor γ = 0 .
• Fig.
3: A : The wind effect causes (0, 2) (2 units southward) displacement at each timestep; B, D : The successor representation is computed given a transition matrix that assumes the variance at each (Gaussian) ﬁring ﬁeld is 1 . 5 , followed by directed actions under the wind effects (with 0 diffusion), the discounting factor is γ = 0 .
9 ; F, G : The optimal following the ascending values of the successor representation, without any wind effect.
The SR is computed given a transition matrix that assumes the variance at each (Gaussian) ﬁring ﬁeld is 1 . 5 , followed by directed actions, the discounting factor is γ = 0 .
9 ; All computations are done by working directly with the Fourier modes instead of the transition matrices.
• Fig.
The environment is the CartPole task (Barto et al.
[3]), and is simulated using the OpenAI gym environment (Brockman et al.
The state value consists of 4 variables: (Cart position, Cart velocity, Pole angle, Pole angular velocity), the action value is an integer takes value from { 0 , 1 } , where 0 represents moving left, and 1 represents moving right.
For constructing the Fourier modes, we discretised each state variable into 8 bins, hence resulting in 8 4 number of states, and we chose the top 50 low-frequency Fourier modes as the inputs to the grid cell network.
The standard DQN agent consists of two fully connected hidden layers with standard ReLU activations, with 48 and 24 units, respectively.
The target network is updated every 500 timesteps.
The deep Dyna-Q agent is a simpliﬁed version of the model proposed in Peng et al.
[33], with an additional 2-layer neural network learning the environmental transition dynamics, with 64 and 32 units in each hidden layer followed by ReLU activations.
At each timestep, the learnt environment model is called to generate K imaginary trajectories that are used for model-based updates to the DQN agent.
The number of model-based updates, K , is taken to be 2 .
The state-action network in the gc-DQN has one hidden layer, with 32 units followed by ReLu activation.
The grid cell network has one hidden layers, with hidden size ( n, A ) followed by ReLU activation, where n represents the number of input Fourier modes, and A represents the number of possible actions.
The deep gc-Dyna-Q has similar architecture as the gc-DQN agent, but with an additional environment network that learns the transition dynamics of the environment that is used for model-based updates (with same architecture as the standard deep Dyna-Q agent).
All models are learnt using the mean squared error loss function and Adam optimiser (Kingma and Ba [24]) with learning rate 0 .
001 and no learning rate decay.
The exploration strength, (cid:15) , is set to be 0 .at
the start of each independent run, decreases by 0 .
at each episode,
and is bounded below by 0 .
A total of 5 independent runs of 100 episodes are performed for each agent.
Note that 100 episodes were simulated for each independent run due to the limited time and computational resources, but the results show that it is sufﬁcient for demonstrating the increase in performance of the gc-DQN agent comparing to the baseline agents.
We will, upon acceptance, show simulations with more episodes (up to the points where convergence of the baseline agents are observed) in the camera-ready version.
All implementations are performed in the TensorFlow framework (Abadi et al. • Fig.
{+5: A: wavevectors of chosen input Fourier modes: k 
+}
=
, k 2 = (1 / 50 , 4 / 50) , k 3 = (3 / 50 , − 3 / 50) , k 4 = ( − 4 / 50 , − 1 / 50) , k 5 = ( − 1 / 50 , − 4 / 50) , k 6 = ( − 3 / 50 , 3 / 50) ; B: real rat trajectory projected onto 50 × 50 2D spatial domain, ﬁring phase interval of the input Fourier modes: [ − 2 .
5 π/ 12 , 2 . 5 π/ 12] , integration time interval: 8 , exponential decay rate: 0 . 2 , grid cell ﬁring threshold: 2 . 95 , directional bias: within ± π/ 2 of the head direction (the range of the relative difference between the direction of the wavevector and the head direction, within which the Fourier modes are allowed to ﬁre); C: running direction: arctan (1 / 3) .
• Fig.
The discounting factor: γ = 0 .
3 , the number of states: 26 , number of effective transitive inference states: 10 .
The Python-based implementations can be found at https://github.com/ucabcy/
Prediction_and_Generalisation_over_Directed_Actions_by_Grid_Cells .