The goal of reinforcement learning is to maximize the value function:

Here, r_t is the reward given at time t, s_t is the state at time t, and γ is a parameter called discount factor. 𝔼_π[·] means expected values when the agent acts according to a policy π. The value function represents the expected cumulative discounted reward.

The value function can be approximated with some basis functions as follows:

Here, V is the estimated value function, w is a weight vector, and x(s) is a feature vector representing state s. The weight w can be learned by standard TD learning adapted for linear function approximation:

where α^w is a learning rate and δ is the TD error. The TD error is defined as

The successor representation (SR) is a state representation based on the prediction of state transitions. The SR matrix M is defined as

The rows of the SR matrix can be used as x(s) in Equation 2 (Russek et al., 2017):

where ws′ is a component of vector w corresponding to a state s′. Then, according to Equation 3, w can be learned by

for all states s′.

We use one of the methods of learning SR using neural networks proposed in a previous study (Fang et al., 2023).

A recurrent neural network (RNN) is used to compute the SR. It is assumed that the transition probability matrix T is encoded in the synaptic weights of the RNN. Then, the steady-state activity of the network in response to one-hot input ϕ retrieves a row of the SR matrix, M^⊤ϕ.

The dynamics of the RNN is defined with the following equation:

Here, x is the activity of RNN neurons, J is the weight matrix of RNN, f is an activation function, ϕ is the input, and γ is a scaling factor of recurrent activation. This dynamics leads to the steady-state activity

when f is the identity function.

The weight J is updated as follows:

where η is a learning rate. For each synapse,

The first term is a temporally asymmetric potentiation term which is similar to spike-timing-dependent plasticity (STDP). The second term is a form of synaptic depotentiation, and similar inhibitory effects are known to be elements of hippocampal learning (Kullmann and Lamsa, 2007; Lamsa et al., 2007).

Although it works with a static learning rate, to accelerate learning, the authors introduced an adaptive learning rate calculated by n(t)=∑t′<tx(t′),η=min(1nj(t),1) for synapses from neuron j. Modulating synaptic learning rates as a function of neural activity is consistent with experimental observations of metaplasticity (Abraham and Bear, 1996; Abraham, 2008; Hulme et al., 2014).

The authors assumed that the timescale of the plasticity is longer than the timescale of the RNN dynamics and that x can be regarded as converging to the steady state in the update of the weight. Under this assumption, the plasticity (Equation 10) leads to

where T is the transition probability matrix. T gives the probability that the agent transitions from a state s to a state s′ in one time step: Ts,s′=P(st+1=s′|st=s). From Equations 9, 12 and

which is derived from Equation 5, we obtain

When f is a hyperbolic tangent, the steady state approximates the rows of the SR matrix, and the model becomes stable for larger γ values compared to when f is the identity function (Fang et al., 2023). Therefore, we use tanh as f. As in the previous study, we use x after repeating the update (Equation 8) for t_max steps such that γmaxt<10-4 as the steady-state activity x_ss.

We adopt an actor-critic method with a policy gradient method (Sutton and Barto, 2018).

The value function is approximated by

and the weight w is learned by

as described in Section 2.1. The inner product in Equation 15 is calculated by synapses whose weight represents w and activity of presynaptic cells represents x(s). The weight update Equation 16 can be interpreted as synaptic plasticity dependent on presynaptic cell activity.

A policy π is defined with an exponential soft-max distribution:

where π(a|s, θ) represents the probability of choosing action a in state s parametrized by θ, h(s, a, θ) represents preference of action a in state s parametrized by θ, and β is a parameter scaling the preference. For the tasks we use, which are maze tasks with 4 actions, preference h is defined as

where x(s) is the feature vector representing state s and

The parameter θ is learned by a learning rule of policy gradient methods (Sutton and Barto, 2018):

when the agent chose action a at state s, where α_θ is a learning rate and γ is a discount factor. For policies expressed by Equation 17, this becomes

when the agent chose action a_i at state s. Here, δ_ik is Kronecker delta. We omit γ^t for learning efficiency. The inner product in Equation 18 is calculated by synapses whose weight represents θ and activity of presynaptic cells represents x(s). The learning rule Equation 21 can be interpreted as synaptic plasticity dependent on presynaptic and postsynaptic cell activity.

We use the row corresponding to state s of the SR matrix or the one-hot vector corresponding to state s as the feature vector x(s) in Equations 15, 18. The row corresponding to state s of the SR matrix is obtained as the steady-state activity x_ss of the RNN in Equation 14 by making ϕ the one-hot vector corresponding to state s.

We refer to the use of x in Equation 15 as “using x for the Critic” and the use of x in Equation 18 as “using x for the Actor.” Different feature vectors can be used for the Critic and the Actor. The structures of our model with possible combinations of state representations are shown in Figures 1A–D.

We adopt Q-learning and SARSA as representative methods for learning the state-action value function. The state-action value function is defined by

The state-action value function can be approximated with some basis functions as follows:

Here, Q is the estimated state-action value function, w is a weight vector, and x(sa) is a feature vector representing state-action pair sa. The update of w in Q-learning adapted for linear function approximation is

where

The update of w in SARSA adapted for linear function approximation is

where

A state-action version of the successor representation matrix is defined as

We use the row corresponding to state-action pair sa of the SR matrix or the one-hot vector corresponding to state-action pair sa as the feature vector x(sa) in Equation 23. Using a version of the RNN in Section 2.2 where each neuron corresponds to a state-action pair instead of a state, the row corresponding to state-action pair sa of the SR matrix is obtained as the steady-state activity x_ss of the RNN by making ϕ the one-hot vector corresponding to state-action pair sa. The inner product in Equation 23 is calculated by synapses whose weight represents w and activity of presynaptic cells represents x(sa).

We use policies similar to Equation 17:

We calculate the state value function by

when we visualize it.

Unless otherwise noted, the following values were used for each parameter: γ = 0.8, α_w = 0.3, α_θ = 0.3, β = 1. The initial values of x, J, w, θ are a zero vector or a zero matrix.

The tasks we used are similar to the latent learning task and the policy revaluation task in a previous study (Russek et al., 2017).

We generate barriers on a 7-by-7 grid world. Simply generating barriers randomly can result in the start and the reward location being separated by barriers, making it impossible to reach the reward. Instead of regenerating the maze if the space is divided, we adopt a generation method that avoids division. To clearly demonstrate differences in results of different learning methods and state representations, mazes are constructed by generating narrow, short paths with dead ends. Two paths of each of the lengths 2, 3, and 4 are generated in random order. Path generation is performed by extending a path randomly from a starting point and making an entrance of the path at its endpoint (the starting point becomes a dead end). When generating paths, states adjacent (including diagonally) to existing paths are excluded. First, paths whose starting point of generation is the top-left, top-right, bottom-left, and bottom-right corners are generated in order. After that, starting points of generation are randomly determined. This generation method ensures the space remains undivided.

Figures 2A–D show typical examples of the learned value function and optimal actions. The learned value function is depicted in a color map, and the direction of the action with the highest preference in each state is indicated by an arrow. We can see that learning of the value function proceeds faster when the SR is used for both the Critic and the Actor (Figure 2A) than when one-hot encoding is used for both (Figure 2D). Appropriate policies were learned with all four combinations of state representations.

Figure 2E shows the mean and standard error of the number of steps for each trial while training sessions of 50 trials were performed 500 times. The decreasing trends in the number of steps mean successful learning. The plots in Figure 2E indicate that using the SR for the Actor enhances learning efficiency, while the effect of using the SR for the Critic is small in comparison.

Figures 3A, B show typical examples of the learned value function and optimal actions for Q-learning and SARSA using the SR, respectively. We can see that learning of the value function proceeds faster for Q-learning than SARSA. This result is natural considering that under the same conditions, an update of Q(s, a) in Q-learning is larger than or equal to that in SARSA because maxaQ(st+1,a),≥Q(st+1,at+1) in Equations 25, 27.

Figure 3C shows the mean and standard error of the number of steps for each trial while training sessions of 50 trials were performed 500 times, for different learning methods using the SR. As for the actor-critic method (hereafter referred to as AC), the SR is used both for the Critic and the Actor. Note that learning of the SR occasionally failed for Q-learning and SARSA, and such sessions are excluded (also for figures below). This phenomenon has been mentioned in the previous study (Fang et al., 2023) that proposed the model for learning SR. The frequency of this failure was less than 10 times out of 500 times. We can see that learning by AC was significantly faster than Q-learning and SARSA, while Q-learning was slightly faster than SARSA. Q-learning and SARSA showed stagnation in learning. Given that both Q-learning and SARSA correctly learned the optimal action at each state, the large number of steps to reach the reward location can be attributed to the small difference between the probability of selecting the optimal action and the probability of selecting other actions. Indeed, in locations distant from the reward location, Q-values are small, and therefore the difference between the Q-value of the optimal action and that of other actions is also expected to be small. In contrast, with AC, the difference between the preference (denoted by h) of the optimal action and that of other actions can become sufficiently large as learning progresses, even in locations distant from the reward location. The fact that Q-learning was slightly faster than SARSA is thought to be a consequence of Q-learning learning the value function faster, as mentioned earlier. The slightly faster learning speed of Q-learning compared to SARSA can be interpreted as a consequence of Q-learning learning the value function faster than SARSA.

Figures 3D, E show typical examples of the learned value function and optimal actions for Q-learning and SARSA using one-hot encoding, respectively. In both cases, there is no significant difference compared to when the SR is used.

Figure 3F shows the mean and standard error of the number of steps for each trial while training sessions of 50 trials were performed 500 times, for Q-learning with the SR and Q-learning with one-hot encoding. Learning by Q-learning with the SR was slightly faster than learning by Q-learning with one-hot encoding. Figure 3G shows the mean and standard error of the number of steps for each trial while training sessions of 50 trials were performed 500 times, for SARSA with the SR and SARSA with one-hot encoding. There is no significant difference in learning speed.

Figures 4A–D show typical examples of the learned value function and optimal actions. There is no significant difference in the learned value function across the four combinations of state representations. When the SR is used for the Actor (Figures 4A, B), the learned optimal actions are not appropriate at some locations in short paths with dead ends. The preference for actions in one state is influenced by the learning of actions in other states when the SR is used for the Actor, and this can result in such inappropriate policies.

Figure 4E shows the mean and standard error of the number of steps for each trial while training sessions of 50 trials were performed 500 times. The decreasing trends in the number of steps mean successful learning. The plots in Figure 4E indicate that both using the SR for the Critic and using the SR for the Actor enhance learning efficiency, with using the SR for the Actor having the greater effect.

As these results show, when the SR is used for the Critic and one-hot encoding is used for the Actor, appropriate policies can be learned and learning is faster than when one-hot encoding is used for both, making it a promising combination of state representations.

Figures 5A, B show typical examples of the learned value function and optimal actions for Q-learning and SARSA using the SR, respectively. As in the water maze task, we can see that learning of the value function proceeds faster for Q-learning than SARSA. The learned optimal actions are appropriate even in short paths with dead ends.

Figure 5C shows the mean and standard error of the number of steps for each trial while training sessions of 50 trials were performed 500 times, for different learning methods using the SR. As for AC, the SR is used both for the Critic and the Actor. As in the water maze task, we can see that learning by AC was significantly faster than Q-learning and SARSA, while Q-learning was slightly faster than SARSA.

Figures 5D, E show typical examples of the learned value function and optimal actions for Q-learning and SARSA using one-hot encoding, respectively. In both cases, there is no significant difference compared to when the SR is used.

Figure 5F shows the mean and standard error of the number of steps for each trial while training sessions of 50 trials were performed 500 times, for Q-learning with the SR and Q-learning with one-hot encoding. Learning by Q-learning with the SR was slightly faster than learning by Q-learning with one-hot encoding. Figure 5G shows the mean and standard error of the number of steps for each trial while training sessions of 50 trials were performed 500 times, for SARSA with the SR and SARSA with one-hot encoding. There is no significant difference in learning speed.

As previously described, in Figure 6A, the agent learns value and action in each placement in order. Then, Figure 6B shows the mean and standard error of the number of steps for each rewarded trial while training sessions of 40 rewarded trials were performed 500 times. Because we set an upper limit on the total number of steps per session during simulation, some runs end before reaching the reward location 20 times after reversing the reward location and the punishment location. Such runs are excluded from the mean and standard error. The number of such runs was 5 for “Critic: SR, Actor: SR,” 148 for “Critic: one-hot, Actor: SR,” 0 for “Critic: SR, Actor: one-hot,” 1 for “Critic: one-hot, Actor: one-hot.” Therefore, the actual difference between “Critic: one-hot, Actor: SR” and other combinations is larger than the plot indicates. The number of steps required to reach the reward location increases after the reversal because the new reward location is the original punishment location. As learning progresses, the number of steps required to reach it decreases. When the SR is used for the Critic and one-hot encoding is used for the Actor, the increase in steps after the reversal is smallest. The plots in Figure 6B indicate that using the SR for the critic suppresses the increase in steps after the reversal, while using the SR for the actor results in a larger increase in steps after the reversal.

We analyze the learning process when using the SR for the Critic and one-hot encoding for the Actor, which showed the smallest increase in steps after the reversal, from perspectives other than the number of steps.

Figure 6C shows a typical example of the history of rewards earned during learning when the SR is used for the Critic and one-hot encoding for the Actor. In this example, the reward and punishment positions are reversed around the 30th trial. Immediately after the reversal, the probability of reaching the punishment position is higher than the probability of reaching the reward position because the policy at that time is to move toward the punishment position (= the original reward position). We can see that the probability of reaching the punishment position is decreasing again after that.

Figures 6D, E show typical examples of the learned value function and optimal actions when the SR is used for the Critic and one-hot encoding for the Actor. It can be seen that after the reversal, the new value function is being learned. Unlike in the barrier maze task, the learned optimal actions lead to dead ends in some short paths. Since this occurs on paths relatively close to the punishment location, it is reasonable to consider that the cause is that the learning effect to move away from the punishment location outweighs the learning effect to move toward the reward location in some short paths.