2.1. Successor representations

The successor representation and the more general successor features describe future expectations of a quantity, conditional on the current state of the world. They are most easily defined in the following setting: Assume the environment of an agent/animal consists of a set of states . The states of the world are changing according to a time homogeneous Markov chain, denoted , with transition probabilities encoded in the matrix P such that , the probability transitioning from state s to state in one timestep. Then for any feature/observation function

(1)

one can define an expectation of weighted, cumulative future values of that function, given the current state:

(2)

The weighting factor puts relatively more importance on proximal times. In the case that is an injective function, that is for every state of the world there is an unique observation value, it makes sense to define the ’successor representation’ (SR)

(3)

We use this terminology here in a little more generality than is usual, but it can be seen that one important special case leads to what is usually called SR: Let , where e_s is the unit vector with a 1 at the entry corresponding to state s. That is, assigns to each state its indicator vector. Then one obtains

(4)

(5)

where the last equality is the well known identity for the Neumann series. The matrix is widely referred to as the SR, so our definition encompasses this special case. One can see from these equations that the SR of indicator vectors thus gives a weighted sum of expected future visitation probabilities of states, and it is this expression that has originally been used to model the predictive representations that hippocampus ought to encode [8,9].

2.2. Learning successor representations with local learning rules

We construct a simple model of two neuron population activities p_t which have dynamics of the form

(6)

(7)

Here, W^r is a recurrent connectivity matrix that feeds the activity of the first population back into itself, while W^f is feedforward matrix, which encodes how activity of the first population is fed into the second. This architecture of one population of highly recurrently connected neurons feeding into a second one with little recurrent connectivity is reminiscent of hippocampal subfields CA3 and CA1 respectively [26]. The two populations obtain additional external inputs , which might represent the input to CA3 via mossy fibers, or directly through the perforant path, and the input to CA1 from EC through the latter, respectively. In the simplest case, these inputs are just indicator-functions for particular states, that is A more realistic shape, which we employ in our experiments in Sect 2.5, might be given by Gaussian inputs of the form The are scalar-valued global gain factors which control the relative strengths of inputs to the populations. Note that we intentionally used the same symbol as for the timescale factors in the successor representation, as those will turn out to be equivalent. For analysis, we assume that the activation function is the identity (i.e. the system is linear), and that the population vectors take the equilibrium values of the above dynamics, that is

(8)

(9)

We then define a learning rule for the synaptic weights, using these equilibrium values. Hence, we implicitly assume that neural dynamics happen on a timescale which is way quicker than that of learning, . Taking the equilibrium values is the limit case which simplifies analysis, but in practice one can also take and simply update activities and weights concurrently. We also assume that our weight matrices are initialized in such a way that these are stable equilibria of the dynamics, which for example will be the case if all weights are initialized to sufficiently small non-negative values.

The learning rule we use is a slight modification of the learning rule used in [25]. In particular, we use the same general learning rule for both recurrent and feedforward weights, only varying certain parameters. Let p^post,i,p^pre,i be the activity of the i-th post/pre-synaptic neuron respectively. We then update the weight from the j–th presynaptic neuron to the i-th postsynaptic neuron via

(10)

The update rule contains terms of the form p^post,ip^pre,j, which are simply Hebbian terms. The other summands perform a subtractive normalization: they subtract the total overall input to a post-synaptic neuron, such that only activity exceeding this input will actually be considered positive. This term has been interpreted as a decorrelative term by [25] - in total, the learning rule can be understand as a predictive coding rule approximating a conditional expectation operation, as we explain in Appendix A in S1 Appendices. In matrix-vector notation the update rule reads

(11)

The parameters control the sensitivity of this update to the order of the activity, put differently, the temporal symmetry of the rule: If we write our synaptic weight change as then for parameters we obtain a learning rule that is invariant under a reversal of time, that is

(12)

while for the original learning rule with no such relation holds - see also Fig 1. For we would obtain a rule that is antisymmetric in this sense - it turns out however that this would yield unstable dynamics.

Download:

• PPT
  PowerPoint slide
• PNG
  larger image
• TIFF
  original image

Fig 1. Model and learning rule.

(A) Cartoon depiction of the model we are using in the main text. A recurrently connected population of neurons p₁, putatively CA3, receives external input , putatively from dentate gyrus or entorhinal cortex(EC). It projects to another population p₂, which receives input . The latter could be CA1 receiving input again from EC. Note that there are no recurrent connections in the second layer and no backwards connections. (B) Quantities relevant for the update of synapse W_ij: pre- and postsynaptic activities, as well as the sum of the total input to the postsynaptic neuron through the synapses W. (C), (D) Invariance of learning rules with respect to temporal order. We plot synaptic weight change of a single synapse in a setup with a single pre- and postsynaptic neuron, respectively. The right column has the same pre- and postsynaptic activities as the left column, only in reverse order. In (C), the learning rule with parameters is used, while in (D) Only in the latter the synaptic weight changes are preserved (in reverse order), while in (C), postsynaptic activity before presynaptic activity leads to a net weight decrease. Note that in this illustrative example W is fixed, in reality, network dynamics and weights would influence each other and lead to more complex changes.

https://doi.org/10.1371/journal.pcbi.1013056.g001

2.3. Network learns successor representation and successor features

Before exhibiting the representations that are learned under the modified learning rule, it might be helpful understanding which representations a two layer-network as the above learns with the simplest choice of parameters (). Let us assume the network has been exposed extensively to features under the same random walk with transition probabilities P, such that the synaptic weights could converge. In practice this means simulating discretized dynamics of (Eq 6) together with the learning rule of (Eq 10), which then results in updates of the form ():

(13)

(14)

Then, as we show in Appendix A and Appendix D in S1 Appendices, the synaptic weights converge in such a way that the network computes successor features. Indeed, the equilibrium is best explained by stating what the population activities encode once the weights have converged. We find that the activities in the network, after extensive exposure to features , when presented with new inputs , compute successor representations/features: They take the equilibrium values

(15)

(16)

In words, this means that the first, recurrent layer computes the weighted cumulative sum of the predicted values of the feature it was trained on, but given the possibly new feature - one could see this a form of pattern completion, but with a predictive component. Similarly, the second layer computes predictions of , only that these predictions themselves depend on those predictive representations passed to it from the first layer. In particular, when the inputs are the same as the model was trained on, and the maps are injective (i.e., sufficiently rich features exist for the environment), then these equations simplify to

(17)

(18)

That is, in this case the two layers simply learn to encode the successor representations of their respective inputs.

2.4. Influence on the representations by choice of parameters

We now proceed to ask the question “which representations would be learned depending on the choice of the parameters ?”. It turns out that the resulting representations are still successor representations, albeit corresponding to transition probabilities that are not necessarily faithful to those of the environmental dynamics anymore. To be precise, we define a weighted sum of transition matrices

(19)

Here, P^forward contains the transition probabilities of the actually observed process (that is, ’forward’ in time), while P^backward contains the transition probabilities of the reverse process, i.e. - see also (Eq 39) We then show in Appendix D in S1 Appendices that under suitable conditions, with the learning rule defined above, the model is able to learn the successor representations under . This entails that the neuronal population activities at convergence yield

(20)

with the respective parameters used in the learning rule. In terms of the successor representation matrix, this simply means

(21)

That is, in particular, under the regime , this corresponds to the ’true’ successor representation. For one obtains the ’predecessor representation’ [16]. Under a symmetric regime, the transition probabilities forward and backward in time are averaged over, that is one obtains

(22)

In fact, in this case the transition probabilities are reversible in time - this is not surprising, as the learning rule was defined to be invariant under a time reversal and hence should only extract aspects of the dynamics which are reversible. We note here that such reversible dynamics have been typically assumed in the theory of SR when construing grid cells as efficient representations of the geometry of an environment through the eigenvectors of the SR - see Appendix C in S1 Appendices.

2.4.1. Activities in the model converge to theoretically obtained limits.

Having theoreticallly obtained the limits of the weights and the corresponding activities, we next verified these limits empirically in simulations. In these simulations, we consider an environment with a discrete number of states and inputs which are functions of these states. In the simplest case, , we obtain the classical successor representation. In particular, for a symmetric learning rule we obtain a symmetrized successor representation - which shows less dependence on the policy. For example, on a circular track, the representation becomes indifferent to whether the agent is performing a clock-wise or a anti-clockwise walk (Fig 2). One might argue that a reversible representation encodes more of the geometry of the underlying state space and less of the actual dynamics (although there is still an indirect influence of the dynamics through the stationary distribution). Indeed, we show in Appendix G in S1 Appendices that the symmetrized transition probabilities are always closer to a uniform policy than the unsymmetrized ones.

Download:

• PPT
  PowerPoint slide
• PNG
  larger image
• TIFF
  original image

Fig 2. Successor Representations learned in circular random walks.

We construct a circular state space with possible actions stay, move clockwise and move anti-clockwise. We simulate three random walks, one where the actions are selected uniformly (first row), one where clockwise actions are preferably selected (second row) and one where anti-clockwise actions are preferably selected (third row). The first column shows an example trajectory of the respective walk. The second and third column show the successor representations learned by the first and second layer of our model, using a symmetric ( and an asymmetric () learning rule, respectively. Note how the successor representation learned with a symmetric rule does not distinguish between the policies. Here, the inputs to the cells are one-hot vectors encoding the respective states and the plotted successor representations are obtained by taking the average population activity in the respective states.

https://doi.org/10.1371/journal.pcbi.1013056.g002

We also verified our theoretical results in more complex scenarios: the convergence prevails also when the features are random inputs instead of one-hot vectors, and also when the random walk is arbitrary instead of circular - see Fig 3. Additionally, we investigated the stability under the choice of parameters . Here we found that it seems that the model is only stable when the positive weight is bigger (in absolute value) than the negative weight, with no convergence at the boundary case of . The latter is in line with the theoretical results - note that (Eq 19) is undefined in this case.

Download:

• PPT
  PowerPoint slide
• PNG
  larger image
• TIFF
  original image

Fig 3. Successor representations are learned for a variety of inputs, dynamics, and parameters.

Top: Convergence of recurrent (red) and feedforward (blue) matrices to their theoretical limit with random features in circular (left) and arbitrary (right) random walks. Bottom: Convergence of recurrent weight (left) and feedforward weight (right) for different parameters . The other set of parameters is fixed to (1,0) and in these experiments, respectively. In graphs, we measure convergence by the loss term as explained in Methods section. In the bottom row, we compute the fraction of the loss at the final step over the initial loss and display the result in a logscale. Thus, negative values indicate converging towards the target. Note that the values on the antidiagonal are approximately 0.

https://doi.org/10.1371/journal.pcbi.1013056.g003

2.5. Place fields under symmetrized rule show less shift

Although both areas CA3 and CA1 show place cells, these cells exhibit different properties and dynamics [33,34]. It is well known, that on a linear or circular track place fields shift backwards opposite the direction of travel in both regions [35,36]. However, when directly comparing cell recordings from both, it has been observed for example in [37] that the shift in CA3 is in general less pronounced, that is, the center of mass of these cells is more stable than in their counterparts in CA1. We hypothesized that a difference in learning rules could explain this effect. Indeed, it is not hard to see theoretically why this should be the case: Through learning, the features get replaced by their successor features . If there is a preferred direction of travel, then states preceding those where the feature puts a lot of mass will also have more mass, since they are predictive of the former states. If there is an asymmetry in the policy, the same will not hold true for succeeding states, hence one observes a shift towards the predecessors. If now however one has symmetric transition probabilities, then there is no directionality, hence this shift would not occur. Indeed, we provide a simple proof of this in Sect 4.5.

We confirmed this intuition, running our two-layer model in a simple linear track where the agent repeatedly moves from the left side to the right. Indeed, we find a tendency to shift in the CA1 cells, which isn’t as pronounced in the CA3 population (Fig 4). Qualitatively, our results match those obtained in [37]. Importantly, these results only hold when using the symmetrized version of the learning rule for CA3, while the asymmetric variant yields almost no distinction.

Download:

• PPT
  PowerPoint slide
• PNG
  larger image
• TIFF
  original image

Fig 4. Symmetric learning rule leads to more stable place fields in linear track.

We simulated an experiment with a rat repeatedly running on a linear track, similar to [37]. A two-layer SR network was used where the recurrent weights had a symmetric (A, middle row) or asymmetric (A, top row) learning rule. In the the symmetric case, there is less shift of the centre of mass of place fields in the modelled CA3 population (red) than in the CA1 population (blue), which is not the case in the asymmetric version. Histograms show distribution of shifts comparing last five laps versus first five laps, while the rightmost plot shows shift relative to the 12-th lap. The results in the symmetric case are qualitatively similar to data (A, bottom row) from Ca^2 + recordings of hippocampal neurons in a similar experiment - figure adapted from Dong, C., Madar, A. D., & Sheffield, M. E. (2021)([37]). In B, we show firing rates of an exemplar cell from CA3 and CA1 respectively, where the symmetric learning rule is used for CA3. The firing rates in each position are averaged over the first and last five laps, and plotted relative to the centre of mass in the first laps. With experience, only the place field in CA1, not the place field in CA3 shifts backwards (arrow indicates direction of travel).

https://doi.org/10.1371/journal.pcbi.1013056.g004

2.6. Generalization and learning with (a)-symmetric rules

Having derived the different kinds of successor representations that symmetric and asymmetric learning rules encode, we next sought to understand what the functional benefits of these representations might be. In particular, we hypothesized that a symmetric learning rule for successor representations might be a relatively simple inductive bias that would favour learning such representations that are invariant under time reversal. This could be useful in such environments where there is a symmetry in transition structure. A particularly simple example of this setting - but still likely for biological agents to encounter - is when the transition structure of the environment is deterministic and transitions in both directions between states are possible. In this case, the state space becomes a metric space and the metric a particular invariant under the symmetry - that is This then hints at a possible benefit of using a learning rule biased towards such symmetry: In a metric space, an optimal policy for navigating towards any target depends only on the metric - in fact one may give a closed form expression as we show in Appendix H in S1 Appendices. Hence, the more the representations that are learned in one particular tasks encode something akin to the distance on the underlying space, the more useful these representations should be for generalizing to new such tasks. In other words, one would want to bias the representations to encode more of the geometry of the space and less of the dynamics of the particular tasks.

In the light of this hypothesis, we trained a reinforcement learning agent, equipped with a temporal difference (TD) learning rule that for a fixed policy would converge to the same weighted representation under , and investigated performance in simple navigation tasks. Note that since we now want to understand the benefits on the computational level irrespective of the biological implementation, we use a classical RL model and not the neural network model from the preceding sections - however, all experiments could also be conducted using such a model. The agent we use encounters transitions of states and updates an internal matrix M which serves as an estimate of the successor representation. When transitioning from state , the update equations are given by

(23)

(24)

where is a fixed learning rate. Note that for this yields the standard TD learning rule for the successor representation [14]. Furthermore, the agent learns a reward vector R, and computes the value function of states via . Together with a local transition model (which we assume as given), the agent can then define Q values of state action pairs Q(s,a) and take the next action based on these Q–values [11,38].

The task for the agent is split in two parts: In the first part, in each epoch, the agent is initialized in a random location and has to navigate to a fixed goal s_target, where a unit reward is received -i.e., This goal does not change over epochs. After a fixed number of epochs, the goal is changed to a new location , randomly drawn from all other locations. Importantly, the agent is then only allowed to relearn the reward vector, not the successor representation matrix.

We find that both the classical SR, corresponding to parameters , as well as the symmetrized version are able to learn the navigation tasks with similar mean learning curves. Although it might be counterintuitive at first that symmetrization yields a policy which still navigates to the correct target, we actually prove in Appendix H in S1 Appendices that indeed an optimal policy is stable under such symmetrization - which shows at least that once such a policy is learned, it can be maintained. However, this result only shows stability once an optimal policy is reached, and indeed one may observe in Fig 5 that the symmetric agent generally shows a higher variation. Furthermore, we find that the symmetric agent seems to be more sensitive to the choice of hyperparameters: Indeed, we find that for smaller learning rates, there is a steeper degradation of performance caused by high variability in the results for the symmetric agent. However, for higher learning rates, both show a similar performance (Fig 6).

Download:

• PPT
  PowerPoint slide
• PNG
  larger image
• TIFF
  original image

Fig 5. Symmetric successor representation agent affords better generalization in simple navigation tasks.

Left: Agents started in random locations in the environments and had to learn to navigate to fixed targets. After 400 episodes, reward location was switched to a new random location, where agents could only relearn the reward prediction vector but not the SR. (Generalization) performance is visualized by total number of steps taken per episode, for an agent using the classical rule (blue) and an agent using the symmetric rule (red). Dashed line indicates change of target location. We show the average performance over different environments as performance is qualitatively similar, see S3 Fig for plots in individual environments. Right: Similar to left plot, but instead of switching target after a fixed number of episodes, the target was switched when the previous target was found with a fixed accuracy. Violin plots show distribution of suboptimality (steps - optimal number of steps) over all environments, for individual environments see S4 Fig. For an outline of the environments see S2 Fig.

https://doi.org/10.1371/journal.pcbi.1013056.g005

Download:

• PPT
  PowerPoint slide
• PNG
  larger image
• TIFF
  original image

Fig 6. Comparison of policy entropy and sensitivity to learning rate among agents with symmetric and asymmetric learning rule.

Left: Policy entropy of the two agents show different trajectories in the generalization experiment. We calculated the entropy of the agent’s policy, averaged over all states, at the end of each episode. This reveals that during learning to navigate to the first target, the symmetric agent has more entropy, which is then reversed when the new target has to be reached. Right: Symmetric agent shows more sensitivity to learning rate parameter for lower learning rates. We trained the agents repeatedly until a fixed accuracy in navigation to the target was met. We then recorded the number of episodes it took until that criterion was reached. Curves show median and interquartile range of this number for the two agents.

https://doi.org/10.1371/journal.pcbi.1013056.g006

Importantly, we then find empirically that on the new targets, the symmetric learning rule outperforms the classical one on average, while both show higher variation in these tests (Fig 5). Thus, one may argue that the successor representation in a symmetric learning regime affords better generalization - at least in a navigational setting. This is not merely an effect of the trajectories that are sampled with the different learning rules - that is, in particular it cannot be attributed to the higher variation during training on the first target: We repeated the above experiment while learning the successor representations based on the transitions obtained from the classical agent alone - the results remain unchanged, suggesting that the symmetric rule yields representations more apt to generalization without the need of a different sampling regime (S5 Fig). Similarly, the results also hold when controlling the norm of the updates, such that both the asymmetric and symmetric update make an equally big update step at each point. In other words, the agent can concentrate on solving the current task and still gets afforded a map of the environment which is less influenced by the current policy. Nevertheless, when learning to navigate to the first target, the symmetric agent still learns a policy that on average has more entropy than the one learned by the asymmetric agent. This increase in entropy then yields a generalization advantage: as soon as the new target is introduced, the effect is reversed and the asymmetric agent has the more entropic policy (Fig 6).

Indeed, one can show that the transition probabilities encoded in will be closer (than the observed transition probabilities) to those that correspond to a uniform policy, choosing every transition with equal probability - see Appendix G in S1 Appendices - in particular meaning they have higher entropy. The successor representation of the uniform policy in turn is closely related to the shortest-path distance [39]. It can thus be used to generalize to any navigational target, while the successor representation under other policies will not neccesarily have this property. Together with our previous considerations on the symmetry of the state space, this led us to hypothesize that the generalization effect of the symmetric learning rule should vanish as soon as there is no such symmetry in the state space any more. We thus repeated the above experiment on a state space that corresponds to a directed graph, where the number of transitions needed to go from s to is not necessarily equal to those needed to travel from to s. Indeed, we find that in such a setting the effect is reversed: there, the classical learning rule leads to better generalization (Fig 7).

Download:

• PPT
  PowerPoint slide
• PNG
  larger image
• TIFF
  original image

Fig 7. Symmetric learning rule provides no advantage in generalization experiment on a directed graph.

We conducted the same kind of experiment as in Fig 5 on a directed graph. Left: The state space is tree-like, with the addition that from the leaf nodes at the last level one travels back to the central node (orange dashed line). Right: In this scenario an SR agent with the classical learning rule (blue) performs better in generalization than one with the symmetric learning rule (red).

https://doi.org/10.1371/journal.pcbi.1013056.g007

2.7. Variation in symmetry

So far the temporal difference learning rules we used had either perfect symmetry or no symmetry at all. In biological agents, such perfect symmetry might rarely be given, rather, one might expect that the temporal sensitivity profile of learning rules could vary both in space and time, meaning that cells exhibit different learning rules at different moments and furthermore two cells might vary in how they learn. We wondered how robust our findings are to such variations. Thus, we investigated the generalization performance of our successor representation agents under variations that might correspond to imperfect symmetry. First, we investigated how well agents would generalize that have a learning rule in between the symmetric and the classical ones we have discussed so far. To do so, we chose parameters , where s is a parameter that we varied. Note that for s = 0 we get the classical rule, while for s = 1 we get the symmetric rule. We found that generalization performance increased with the parameter s, with the highest generalization at s = 1, again confirming that the symmetric rule might afford a representation that generalizes best.

To further assess the robustness of our results, we then investigated whether dropping the assumption of globally fixed values would qualitatively alter the results. Again, this is motivated by the assumption that in biological agents, learning rules will not be perfectly static but might vary.

To capture this in the RL setting, first,we introduced noise to the parameters at every time step. That is, our update rule for the successor representation would become

(25)

where is independent noise that is added at each timestep to the parameters. Secondly, we studied a condition where we introduced a state-dependent heterogeneity in the parameters. That is, for each possible transition in the state space, there is a separate set of parameters which are randomly initialized at the start of learning and then fixed. The update rule thus is

(26)

To determine the parameters, first some putative values are drawn from Gaussian distributions with identical variance and different means , and afterwards the absolute value of those is taken. This is to ensure that we don’t have negative parameters, which might result in qualitatively different and unstable learning rules. The means are then either set to (1,0) to mimic the asymmetric with inhomogeneous parameters, or to for the symmetric agent. In both variations we again find qualitatively similar results as before: even with noisy parameters, both agents are able to learn simple navigation tasks (Fig 8). The symmetric agent shows more variation in performance, but is in turn able to generalize better to new targets.

Download:

• PPT
  PowerPoint slide
• PNG
  larger image
• TIFF
  original image

Fig 8. Variations of symmetry in the learning rule

Experiment for all plots is the same as in Fig 5. Top: Generalization for parameters . Violin plots show distribution of differences (steps-optimal number of steps) when evaluated on new targets. Distributions broaden towards the optimal value of 0 with increased symmetry. Middle: Generalization with parameters , randomly initialized for each pair of states. Bottom: Generalization with noise added to , at each timestep.

https://doi.org/10.1371/journal.pcbi.1013056.g008

2.8. Maze tasks

The grid worlds we used to investigate the generalization capability so far were topologically very simple. In reality, biological agents will be exposed to more complicated navigation situations, with multiple paths and detours. We thus repeated our setup in a more contrived maze with multiple chambers and more convoluted paths. In particular, instead of keeping the environment static, when we tested the generalization abilities of the agents we introduced slight modifications, such that previously open paths were blocked. Again, in this setting the symmetric agent showed a better generalization, indicated by a higher probabilty of solving the new task with a number of steps close to the optimal one, as indicated by the distribution plots in Fig 9. However, in comparison to simpler grid worlds, both agents are worse at generalization in this setting, which might be caused by the additional complexity of the environment.

Download:

• PPT
  PowerPoint slide
• PNG
  larger image
• TIFF
  original image

Fig 9. Generalization performance in maze task with blocked paths

Top: Grid world mazes used for generalization task. Leftmost maze was used for training, the other three environments for testing the generalization. Bottom: Violin plots show distribution of suboptimality (steps - optimal number of steps) of the agents when using the successor representation trained on one target and tested on another one. Training and test targets are randomly drawn from all possible states in the respective environments. The distribution for the symmetric agent is broader around 0, which indicates optimal generalization, and less pronounced at 400, which was the maximum number of steps allowed.

https://doi.org/10.1371/journal.pcbi.1013056.g009

4.1. Methods

4.1.1. Neural network model.

We consider two populations of rate-based neurons , respectively. The population p¹ is recurrently connected via a matrix of synaptic weights , and feeds its activity forward to population p² via the matrix of synaptic weights . Both regions receive additional external inputs, , and decay to equilibrium in the absence input. The temporal evolution of the populations is then modelled by the following differential equations:

(27)(28)

Here, is an activation function, and the parameters control the relative weight of the different types of inputs. On the computational level, these parameters will correspond to the discounting factor of the successor representation, effectively scaling how far in the future predictions are made.

On the level of biological implementation, there are two speculative interpretations one may have for these parameters. First, the might be associated with acetylcholine. Indeed, this transmitter is believed to control the relative strength of external and internal input in hippocampus [77]. In particular, a low value of would correspond to an encoding mode, where the network is relatively stronger driven by external input and thus susceptible to encoding associations or transition structures between these inputs. For high , the network is in retrieval mode and its outputs will depend more on the recurrent weights and therefore previously learned transition structure. Indeed, this interpretation has been pursued in [25], who showed how a recurrent net could learn weights at low and then use those learned weights to retrieve successor representations at arbitrary high . Another interpretation of would be that it represents the respective proportions of inputs of different kinds a certain subpopulation receives due to anatomic difference. That is, it could for example encode the number of synaptic contacts from other CA3 axons for a certain subpopulation - then understood not in absolute terms, but compared to other cells of the same type. Indeed, there is evidence that both connections in as well as connections to hippocampus are not homogeneous but vary along the septotemporal and proximodistal axis [78,79]. Thus, it might be plausible to assume a varying degree of strength of the different input types. In this work, however, we do not explicitly model these biological details, and thus the precise role of remains undetermined.

Assuming the weights (and the external inputs) change on a slower timescale than the population dynamics, we can let the above differential equations go to equilibrium for analysis - in practice, we can also integrate the ODE above with a timescale orders of magnitude smaller than the timescale of learning. If is approximately linear, this results in

(29)(30)

We assume that the animal receives inputs depending on the current state of the environment, which we will denote by the process S_t. In the reinforcement learning setting, the agent influences the way in which the state-process is sampled by selecting actions A_t, but at this point this is not relevant since we are only building up a predictive representation of states, not actions. The inputs to the two neural populations thus take the form

(31)

where is a function mapping the state space to an activation pattern in neural space. If the state space is discrete, we can also write this as

(32)

with now being a matrix of appropriate dimension, and is the unit-vector corresponding to state S_t, which hence selects the corresponding activation pattern from s columns. In the special case , we have a so called one-hot encoding, where hence every cell corresponds to a distinct state and fires if and only if the agent is in that state.

4.1.2. Learning rule.

The synaptic weights - both feedforward and recurrent - in our model are subject to a learning rule which is controlled by two parameters :

(33)

Here, p^pre,p^post denote the vectors of pre- and postsynaptic activities. The parameters do not model a biological substrate, they control the qualitative behaviour of the learning rule: they serve as a (crude) discrete approximation for a continuous plasticity kernel as is typically observed in spike-time dependent plasticity protocols. Indeed, for the learning rule strengthens connections where post-synaptic activity at the later timestep was high (measured relative to overall input), and presynaptic activity at the earlier timestep was high as well. In contrast, for , the learning rule strengthens connections where postsynaptic activity at an earlier timestep was high, and so was presynaptic activity at the later timestep.

Note that a synaptic update of the form

(34)

may be seen as performing a sort of conditional expectation or regression objective: At equilibrium, the expected update in synaptic weights, given the current activity, should be zero

(35)

We have that this holds in general only if

(36)

That is, if W encodes an optimal prediction of the post-synaptic activity given the presynaptic activity. Hence, heuristically, emphasizes to learn weights that are optimally tuned to predict the next state of their output. In our setting, the cells are driven by external input identifying the state S_t of the world, so this amounts to optimally predicting the next state S_t + 1. In turn, emphasizes a backward prediction, where S_t should be predicted by S_t + 1.

4.2. Forward and backward process, reversibility

Recall that a time-homogeneous Markov process on a finite state space is determined by its one-step transition probabilities

(37)

Under standard assumptions, the process has a unique stationary distribution [80,81], that is, a probability distribution over states such that

(38)

That is, if we are in stationary distribution, the probability mass is left invariant when transitioning according to our process. In particular, if we assume we start in stationary distribution, then . We may now define the reverse process of our original process, by simply reversing the order of time. That is, define a process R_t with transition probabilities

(39)

These probabilities will generally depend on the timestep t, but they will not if our process S_t is in stationary distribution. In that case, we again have a single matrix that encodes all transition probabilities:

(40)

(41)

(42)

In matrix notation, we have

(43)

where is a diagonal matrix with the values of on the diagonal. It is easy to see from this definition that P^backward is again a valid transition matrix, and that it has the same stationary distribution as P = P^forward.

Forward and backward transition probabilities are equal if and only if the process fullfills the detailed balance condition

(44)

In that case, the process is called reversible. In (Eq 19) we have used a weighted sum of transition probabilities of the forward and backward process Note that the symmetrized process with transition probabilities

(45)

will always be reversible, as can either be seen by a simple algebraic calculation, or by the following argument: Let P^* be the adjoint of P with respect to the inner product induced by . This means that for any vectors x,y we have

(46)

We then have that

(47)

as may be seen by plugging in the definition of P^backward into (Eq 46). For the adjoint, we have that . Thus, in particular,

(48)

This also explains that it is warranted to talk about a ’symmetrization’ here: indeed, the adjoint with respect to the standard inner product is the transpose of a matrix A, in which case one would obtain the standard symmetric part of a matrix, that is, .

4.3. Reinforcement learning and successor features

In Reinforcement learning, one generally considers a Markov decision process (MDP), that is a tuple , where are the sets of possible states and actions, gives the transition probability from state s to state when choosing action a, R(s,a) is obtained reward when chosing action a in state s, and is a discount factor. The goal in RL usually is to find an optimal policy, that is a probability distribution over actions, given states, which maximizes the expected discounted cumulative reward

(49)

where the states and actions are sampled according to the policy and the transition probabilities T. In our experiments, we are only interested in the particularly simple case where the transitions are deterministic - that is, taking an action a in state s surely leads to a specified state . Furthermore, we only consider the situation where the reward is a function of the state only. However, all definitions in the paragraph below readily generalize to functions of states and actions. Assume a fixed policy and a process S_t following said policy. Now consider any mapping from the state space

(50)

which we can interpret as an observable generated by the state space. Then define a function on the state space

(51)

thus gives the expected (exponentially weighted) cumulative future sum of the observation or feature , given the current state s. This is hence called a ’successor feature’ in the literature. Define . We then have

(52)

(53)

(54)

(55)

In particular, if the transition probabilities of S_t are time-homogeneous, we see that SF itself does not depend on t, and we can write

(56)

Temporal Difference learning (TD-learning) uses this relationship to construct a target to update an estimate of online. Indeed, if is the current estimate the agent has for the successor features, then consider update rules

(57)

(58)

Then we see that if , we get . Thus, here is used as a bootstrapping estimate of the target , and is updated according to the error to that target - in total, is thus also called TD-error.

Successor features encompass important special case examples: For the choice of the reward function, SF_R becomes the value function: The value function under a policy , which is typically denoted as , hence encodes the weighted cumulative sum of expected rewards, given the current state. In particular, the current estimate of the reward function may be used to define a new policy as

(59)

That is, the policy deterministically selects the action a^* which leads to the next state with the highest value according to the current value function estimate. Iterating this process then successively learns a better value function - this process of updating an estimate of the value function and then choosing an optimal policy with respect to it is the basis of the classical TD-learning algorithm [38]. Indeed, for our navigation experiments we use this approach, only replacing the maximum in (Eq 59) by a softmax which smoothens the transition probabilities. Note that the above assumes a model of which actions lead to which next states - which posing as given should be a sensible assumption in navigation problems - but a completely model-free approach simply computes the value of a state-action pair instead.

Successor Representations. In the case that we have that is in fact an injective mapping, we can define a modified version of successor features as

(60)

where Here, the injectivity of the mapping is necessary to ensure SR defined as above still enjoys desirable properties like homogenity in time and the Bellman equation, but besides of that one could also define a similar quantity without using injectivity. Now, in the injective situation, we would like to call the above ’successor representation’. This is because when we take mapping , where e_S is a unit vector in , then we obtain

(61)

which corresponds to the classical definition of successor representation. In general, if we assume is the matrix of features (could also be an operator if we allow for continuous state space), then we have that

(62)

where is such that .

Generalization and Successor Features. The idea behind using successor features for some set of function , instead of simply directly computing the the value function is that of generalization/transfer: Assume the reward function R can be written as a linear combination of the features , that is

(63)

Then also for the respective predictive representations one has

(64)

Now assuming that the reward changes to a new function , which also can be expressed with the features, then the only thing that has to be relearned is the weights w_i, while the successor features can be reused. Thus, one may then generalize more easily to a new task, since the transition structure under a policy is essentially already learnt. In particular, in a discrete state space, a fixed reward function can R can of course be encoded by a reward vector R, that is

(65)

The value function is then simply given through the classical SR as

(66)

Thus, for a fixed policy, this separates the computation of value into learning a successor representation and learning a reward vector. Our navigation experiments probe the generalization capability of this approach, by introducing a new reward vector , after an optimal policy and a SR for a reward vector R have been learned. Importantly, only the reward vector is allowed to be relearned, while the SR has to be reused from the previous task.

Symmetric TD-learning. In our reinforcement learning experiments, we use a modified version of TD-learning to mimic the behaviour of the local learning rule in the SR-network. In practice, Successor features are typically parametrized by some parameter (e.g., the weights of a neural network), which is then updated to reduce the TD-error. That is, for each value of we obtain a map of states. We can then update the parameter via

(67)

(68)

(69)

where are parameters corresponding to the ones controlling the local learning rule. In particular, for one obtains the classical TD-learning rule for function approximation [38] and for one obtains the ’predecessor representation’ [16]. The case yields a symmetrized version of TD-learning. In our experiments, we use a particularly simple version of the above: for a discrete state space, one can simply parametrize by means of a matrix , where d is the number of features and k is the number of states. Then one has , and hence the update rule

(70)

4.4. Successor representations and hitting times

In the theory of Markov chains, it is a common exercise to study the fundamental matrix of the Markov process, which encodes properties about the process via the first hitting times of states [82]. The first hitting time of a state, , is defined as the first time a process hits the state s. The successor representation may similarly be interpreted as an operator that encodes certain expectations related to the first hitting times. Indeed, in Appendix H in S1 Appendices we prove the following formula:

(71)

Here, denotes the largest number of times which state may be hit with nonzero probability when starting from state s. It may only take values 0, if may not be hit from s at all, 1, if is a transient state, or , when the process may arbitrarily often return to a state. (Eq 71) provides a convenient reformulation of the value function. In particular, when we are in the setting of a navigation task, where the goal is navigating to target s^*, the reward is a unit reward at s^*. Hence, for an optimal policy which only chooses among shortest paths to the target, the successor representation becomes

(72)

where is the distance on the state space. The value function may then be read of as the column corresponding to s^*. We use this formula in Appendix H in S1 Appendices to prove a stability result of optimal policies in the navigational setting under symmetrization: If we are in the deterministic navigation setting, and our policy is optimal, then after taking the symmetrization of the transition probabilities and computing the value function with these probabilities, the policy is still optimal.

4.5. A predictive representation leads to backward shift only in the asymmetric case

Here we want to give an analytic explanation for the stronger backward shifts of the centre of mass observed in Fig 4 when using an asymmetric learning rule, while a symmetric learning rule does not result in a shift. Recall that we simulated an agent repeatedly running on a linear track, always in the same direction, which was then reset to the starting position. The effect is most easily analyzed if we ignore boundaries and go to a continuous situation. Let’s thus assume an agent is repeatedly running on the real line . In the beginning, before learning, the cells in our model are just driven by the external input, and thus can be taken equal to the features. That is, cell j will take the value when at position x. After learning, it will instead encode a successor feature. With continuous space it is also more convenient to assume continuous time. The successor representation may be easily transferred to this setting [23], where the definition then is

(73)

where is the probability of transitioning from state x to state y in time t. The centre of mass (COM) of a cells’ firing field is just the spatial mean of the cells’ activity when the latter is normalized to be a probability distribution. That is, for example before learning, when the cell j fires at position x according to the feature , the j-th centre of mass is

(74)

To understand the observed shifting effect in COM, it is now instructive to study the successor feature with deterministic dynamics: Assume we are moving with constant velocity v, that is,

(75)

then the successor feature takes the form

(76)

Now, taking the mean over all positions yields:

(77)

(78)

(79)

(80)

(81)

Similarly we have

(82)

(83)

(84)

Thus, in total we obtain

(85)

That is, the center of mass is shifted backward by a factor which is controlled by the prediction timescale and the velocity of the movement. We thereby also obtain the seemingly new prediction that faster running speed and smaller terminal place field size (as a proxy for ) should result in bigger shifts. Now if our forward process is a deterministic process with constant velocity v, then conceptually, our backward process is also deterministic with velocity –v (although stricly speaking, there is no stationary distribution). In particular, the symmetrized process is a process which in an infinitesimal amount of time travels either forward with velocity v or backward with velocity v with equal probability. The shifts thus cancel and the centre of mass stays the same as for the initial feature. The same relation of shift and velocity/timescale holds also when the dynamics come from a Brownian motion with constant drift, that is as we show in Appendix B in S1 Appendices. Now when we symmetrize the Brownian motion with drift, the drift term will disappear, that is, we have a standard Brownian motion. It is thus clear that under this motion, there should again be no shift of the centre of mass.

4.6. Experiments

4.7. Convergence experiments

For our initial convergence experiments (Fig 3, top row) we simulate our model in a circular random walk setting with 30 states and use n = 40 cells in each layer. The input at each state is drawn from an i.i.d. Gaussian distribution () (i.e., for each cell and each state we draw a value from a Gaussian distribution, this is then assumed to be the input for that cell whenever the specific state is visited). Throughout this and the following experiments, we use a moderately high value of - meaning a relatively large timescale/small discounting. This is an arbitrary choice, but relatively high values have been used in the past in SR-theories of hippocampus [83].

We observe 100000 transitions, and repeat the experiment 30 times. We then repeat the same experiment (top row, right of Fig 3), also drawing the transition matrix P randomly. To track the convergences, we consider the loss terms

(86)

(87)

where the terms on the right hand vanish at convergence and are defined in (Eq 45). The resulting convergence curves are shown in Fig 3. These experiments are conducted with a linear activation function, we also repeat this experiment with the activation functions as shown in S1 Fig.

To check the effect of different values , we conduct a parameter sweep in a circular random walk, and random Gaussian features as above. For each combination we run 30 random initializations for 1000 iterations, and check this combination once for the recurrent layer and once for the feedforward layer. Here, we track convergence similar as above, but now we take the fraction

(88)

with W₀ the initial weight, to judge if we are moving towards or away from equilibrium. This results in the matrices shown in Fig 3. To speed up the experiments, we use a smaller network for these experiments, with 20 neurons in each layer and 10 states. Again we use dt = 0.1.

4.8. Linear track experiments

To simulate the animal running on a linear track, we partitioned a track of length 300cm in 50 discrete states and let the agent perform a rightward biased walk that either took a step to the right (p = 0.9) or stayed in place (p = 0.1). We chose a time step of dt = 0.4, corresponding to a velocity of . The agent would run 25 laps on the linear track, with a short resting phase in between two laps. We set up a two-layer model, with n = 100 cells in each layer, and again . For the feedforward weights, we always used the parameters , while for the recurrent weights we tested both the symmetric case as well as the asymmetric case We then performed an analysis approach as in [37]. For each lap, we collected the mean activities of each layer per state and computed the center of mass (COM) by the formula

(89)

where x is position on the track and denotes the mean activity of cell j when at that location. To obtain a distribution of the overall shifts of the COMs over all laps, we calculated the average COMs for the first and the last five laps and subtracted them from each other, obtaining the histograms in Fig 4. To track the evolution of the shift, we subtracted the COMs from the 12-th lap, which yields a more gradual tracking of shifts. We plot the results of this approach in the scatterplots in Fig 4.

4.9. Reinforcement learning

We investigated the generalization capabilities of a symmetric over an asymmetric learning rule in TD-learning in different scenarios.

4.9.1. Navigation experiments.

Since the hippocampal formation is prominently involved in navigation tasks, we first focused on tasks of such nature. We thus studied different variations of the same general problem setup: Given a grid environment with a deterministic transition structure, the agent would receive a unit reward only when arriving at the designated target state s_T, upon which the episode is terminated, and the next episode starts in an initial state s₀, drawn uniformly at random. The agent is thus encouraged to navigate from all possible starting states to the target state to receive a reward. We used the Neuronav toolbox [84] to implement the grid environment and implemented our modified SR-agent in the framework of that toolbox. For completeness, let us here again give a concise summary of how the agent learns. The agent possesses a world-model and updates an internal estimate M of the successor representation as well as an estimate w of the reward vector. Having observed a transition from state s to state via action a and a reward r, the updates are

(90)(91)(92)

In summary, this is how the agent selects actions and updates its internal estimates:

• Given the current state s and current estimate M for the successor representation, compute
• Select an action according to
• Observe next state and reward r
• Update reward estimate w and successor representation estimate M.

To check for generalization capability, we would train agents which utilize the respective learning rules (symmetric,asymmetric) to navigate to a fixed s_T first. We trained the agent either for a fixed number of episodes (we used 400 episodes with a maximum number of steps of 400), or until a fixed accuracy criterion was met (the mean deviation from optimal performance for the preceding 8 episodes was lesser than 2). Then, we would randomly select a new target state from among all possible states. Importantly, after the modification of the target states, the agents were only allowed to modify the reward-prediction vector w, but not the successor representations M themselves. That is, they could only learn the reward structure of the new task, but had to rely on the transition structure that was encoded during learning the previous task, which also depends on the policy. We repeated this experiment for 200 times, with randomly drawn combinations () in each repetition. The results of these experiments are depicted in Fig 5. For the first experiments reported in the main text, we used a learning rate of 0.1 and set .We chose a relatively high value of , as this is typically done when modelling hippocampal place cells in standard navigation tasks [8,20], but we repeated the second kind of analysis for different values of and the learning rate, as shown in S6 Fig.

For the experiments in the maze, we used the same parameters for the maximum number of steps and the performance criterion, but higher values for the learning rate and the discount factor: . This is because for lower parameter values the agents needed too many repetitions to converge.

Since all the environments we tested in the above setting have a symmetric nature, that is, their underlying state space is an undirected graph, we repeated the same generalization experiment in a setting where the state space is a directed graph, and hence travel time between two nodes is not symmetric. We constructed a simple graph with 17 nodes, which is essentially a tree graph with the addition of a directed edge from the lowest level to the highest. On this graph, we performed the same kind of navigation experiment. We used the same values for the learning rate and (0.1,0.7), but reduced the number of episodes to 50 since the state space is smaller and hence can be learned faster.

4.9.2. Policy entropy.

In Fig 6 we computed the entropy of the current policy of the agents after each episode of learning. As explained previously, we computed the current policy by applying the softmax-function to the q-values

(93)

where w is the estimate of the reward vector and M is the current estimate of the successor representation. Then the policy is: , and we computed the entropy of the policy averaged over all states as

(94)