Prediction and generalisation over directed actions by grid cells

1 INTRODUCTION . A `` cognitive map '' encodes relations between objects and supports flexible 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 . [ 20 ] ) . Each place cell fires when the animal is near a specific location , whereas each grid cell fires periodically when the animal enters a number of locations arranged in a triangular grid across the environment . Together , this system could support representation and flexible 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 . [ 43 ] ) . Recent work suggests that place cell firing provides a local representation of state occupancy , while grid cells comprise an eigenbasis of place cell firing covariance ( Dordek et al . [ 15 ] ; Stachenfeld et al . [ 38 ] ; Sorscher et al . [ 37 ] ; Kropff and Treves [ 26 ] ) . Accordingly , grid cell firing 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 . ∗Please send any enquiries to : changmin.yu.19 @ ucl.ac.uk and n.burgess @ ucl.ac.uk 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 . [ 2 ] ) . 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-specific eigenvalues . We define 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 firing . 2 “ INTUITIVE PLANNING '' WITH A SINGLE TRANSITION STRUCTURE . 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 . [ 2 ] ) . Specifically , consider a transition matrix , T ∈ RN×N , Tss′ = P ( st+1 = s′|st = s ) where st encodes the state at time t and N is the number of states . Then , Tn is the n-step transition matrix , and has the same set of eigenvectors as T . Specifically , the eigendecomposition of T and Tn are : T = QΛQ−1 , Tn = QΛnQ−1 ( 1 ) 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 : ∞∑ k=0 ( γT ) k = ( I − γT ) −1 = Qdiag ( w ) Q−1 , where w = { 1 1− γλ , for λ ∈ σP ( T ) } ( 2 ) The resolvent form ( Eq . 2 ) is an infinite 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 . [ 38 ] ) . 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 ) . 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 can not flexibly adjust its predictions corresponding to the effects of arbitrary directed actions ( i.e. , variable asymmetric transition structure ) , hence can not 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 . 3 FLEXIBLE 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 . [ 30 ] ) . 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 ∈ RLW×LW , with each row the vectorisation ( vec ( · ) ) of the matrix of transition probabilities starting from the specified location , i.e. , T [ j , : ] = vec [ P ( st+1|st = 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 : T = T0 TLW−1 · · · T2 T1 T1 T0 TLW−1 · · · T2 ... T1 T0 . . . ... TLW−2 · · · . . . . . . TLW−1 TLW−1 TLW−2 · · · T1 T0 ( 3 ) The normalised eigenvectors of the circulant matrix T ∈ RN×N ( N = LW ) are the vectors of powers of N th roots of unity ( the Fourier modes ) : qk = 1√ N [ 1 , ωk , ω 2 k , · · · , ω N−1 k ] T ( 4 ) where ωk = exp ( 2πiN k ) , for k = 0 , . . . , N − 1 , and i = √ −1 . Hence the matrix of eigenvectors ( as the columns ) , F = ( q0 , q1 , . . . , qN−1 ) , is just the ( inverse ) discrete Fourier transform matrix ( Bracewell [ 4 ] ) , where Fkj = ω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 specifies 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 first row ( or column ) of T , see Bracewell [ 4 ] ) : λm = N−1∑ j=0 Tjω m j , for m = 0 , . . . , N − 1 ( 5 ) where { T0 , . . . , TN−1 } are the N unique elements that fully specifies the circulant matrix T ( Eq . 3 ) . We can then utilise tools from Fourier analysis for efficient updating of the eigenvalues whilst leaving the universal eigenbasis unaffected . For a transition matrix Tv corresponding to an arbitrary action ( translation velocity ) v = ( vx , vy ) , each row of Tv is again a circulant , but shifted version of the corresponding row vector of the symmetric transition matrix corresponding to zero drift velocity , T0 . Specifically , the first rows of the two matrices are related as follows : Tv ( k ) = T0 ( k + vxL+ vy ) , for k = 0 , . . . , N − 1 ( 6 ) Given the eigenvalues for T0 , Λ0 = [ λ00 , λ 0 1 , . . . , λ 0 N−1 ] ∈ CN ( via the DFT of the first row of T0 , Eq . 5 ) , we can immediately derive the eigenvalues of Tv , Λv , via a one-step update based on the Fourier shift theorem ( Bracewell [ 4 ] ) without recomputing the eigendecomposition : Λv [ k ] = exp ( 2πi N ( vxL+ vy ) k ) Λ0 [ k ] , for k = 0 , . . . , N − 1 , for arbitrary v , i.e. , Λv = ΦδvΛ0 , Φδ ( v ) = [ 1 , ωδ ( v ) , ω 2 δ ( v ) , . . . , ω N−1 δ ( v ) ] , where δ ( v ) = vxL+ vy ( 7 ) 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 fixed eigenbasis ( Eq . 4 ) and the respective eigenvalues ( Eq . 7 ) for arbitrary transition structures , we can make efficient 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 efficiently generalised across environments . Sense of Direction . We define 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 . θ∗ = arg max θ ∑ j exp [ 2πi ( xG − x0 ) · kj ] 1− γDj exp [ 2πivθ · kj ] ( 8 ) where γ is the discounting factor , Dj , j = 1 , · · · , LW are the eigenvalues for the symmetric diffusion transition matrix , kj , j = 1 , . . . , LW are the wavevectors for the j-th Fourier components , x0 , xG 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 sG 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 finding 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 reflects 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 flexible 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 efficiently 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 , S0 , being embedded into a larger , periodically bounded pseudo state space Sp , at least twice as large in each dimension as S0 ( Fig . 3 F ) . We again follow the previous procedures , utilising the Fourier modes , this time computed on Sp , to perform predictions in S0 ( 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 S0 , 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 to 1000 time 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 .

Summary: The authors propose as extension of the successor-representation approach to Grid cells. The paper shows that this model can generate several experimentally observed properties of grid cells, and can be used in navigation of novel/mutable environments. Overall, the work should be of interest to any ICLR attendees who engage in research surrounding grid cells.