Tuning the Weights: The Impact of Initial Matrix Configurations on Successor Features Learning Efficacy

In this study, we assume that an RL agent interacts with the environment through Markov decision processes (MDP, [12, 13]). An MDP is a tuple $M:=(\mathcal{S},\mathcal{A},R,\gamma)$ comprising of the following elements. Sets $\mathcal{S}$ and $\mathcal{A}$ are the state (e.g., spatial locations) and action spaces. The function $R(s)$ specifies the immediate reward received in state $s$, which can be expressed as $R:\mathcal{S}\rightarrow\mathbb{R}$. Here, the discount factor $\gamma\in[0,1)$ is a weight that reduces the reward in the distant future.

In RL, the agent’s objective is to discover a policy function $\pi:\mathcal{S}\mapsto\mathcal{A}$ that maximizes the cumulative discounted reward, often referred to as the return $G_{t}=\sum^{\infty}_{i=t}{\gamma^{i-t}R_{i+1}}$, where $R_{t}=R(S_{t})$. The return essentially represents the sum of all future rewards that an agent can expect to accumulate, discounted by the factor $\gamma$. In order to solve this optimization problem, a common approach is to employ dynamic programming, which defines and computes the value function of a policy $\pi$ as follows:

where $\mathbb{E}^{\pi}[\cdot]$ denotes the expected value when the agent follows policy $\pi$. After determining $V^{\pi}(s)$, also known as policy evaluation, the policy $\pi$ can be improved in a greedy manner. This process, referred to as greedy policy improvement, is defined as follows: $\pi^{\prime}(s)\in\text{argmax}_{a}Q^{\pi}(s,a)$ where $Q^{\pi}(s,a):=\mathbb{E}[R_{t+1}+\gamma V^{\pi}(S_{t+1})|S_{t}=s,A_{t}=a]$ [13]. Here, $Q^{\pi}(s,a)$ represents the expected return from taking action $a$ in state $s$ and following policy $\pi$ thereafter.

As proposed in literature [14], the central premise of SR learning lies in the decomposition of the value function (Eq. 1). It suggests that the value function can be decomposed into an expected visiting occupancy and reward of the successor states $s^{\prime}$ as follows:

where $\mathbb{I}(S_{i}=s^{\prime})$ yields a value of 1 when an agent visits the successor state $s^{\prime}$ at time $t$; otherwise, it returns 0. Consequently, $M(s,s^{\prime})$ represents the discounted expectation of visitation to the successor state $s^{\prime}$ from the state $s$. $M(s,s^{\prime})$ can be perceived as a comprehensive representation that integrates not only the immediate transition probabilities from state $s$ to state $s^{\prime}$, but also the cumulative impact of the agent’s policies and the array of potential future trajectories. This interpretation underscores the dynamism and predictive capacity of the SR, as it encapsulates the influence of the agent’s decisions and environmental dynamics on future state visitations [14].

The SR matrix $M$ can be incrementally learned by the agent through the use of the temporal difference (TD) learning algorithm. This approach allows the agent to continually update its understanding of the environment based on the difference between predicted and actual visitation. The specific TD learning equation for the SR matrix $M$ is derived as follows [4, 14]:

The classical form of SR learning is constrained to tabular environments, limiting its applicability to more complex, high-dimensional settings [15]. An effective means of circumventing this limitation is the application of a set of feature functions, denoted as ${\psi(s)}$, which allows for the generalization of SR learning [11].

By assuming the expected reward of state $s$ can be represented as a product of the feature vector and its corresponding reward weights, denoted as $R(s)=\phi(s)^{T}\mathbb{w}_{rew}$, we can reframe the value function (Eq. 1) in a way that accommodates these feature functions. The revised value function is given as follows:

where $\phi_{t}$ denote $\phi(S_{t})$. By incorporating a one-hot vector in $\mathbb{R}^{|\mathcal{S}|}$ for a tabular environment, $\psi(s)$ essentially mirrors the $M(s,:)$ vector of SR learning. This is because it represents the discounted sum of occurrences of $\phi(s^{\prime})$ when a transition unfolds under policy $\pi$. For clarity, we will henceforth refer to $\psi^{\pi}(s)$ as the successor features (SF) associated with state $s$ under policy ${\pi}$.

The introduction of SF marks a significant broadening of the SR learning framework, facilitating its application across a wider spectrum of MDP environments, such as partially observable MDPs and those characterized by continuous states [16].

In our approach, the SFs are approximated utilizing a linear function represented as follows:

This estimation leans on the presumption that $\phi(s)$ operates as a population vector of neurons that responds to the state $s$ observed by an agent. The utilization of a linear function aligns with neurobiological models of hippocampal place cells and finds support in the literature, reinforcing its relevance and applicability in our research [5, 15, 4].

To estimate $\psi(s)$, we apply the TD learning to update the weight matrix $W$. This procedure parallels the matrix $M$ updating method observed in successor representation (SR) learning, thereby offering a streamlined approach to SF estimation in reinforcement learning contexts.

It’s worth highlighting that Eq. (6) corresponds to Eq. (3) when $\phi(s)$ is presented as a one-hot vector. Alongside this, the expectation weight vector associated with rewards, denoted as $\mathbb{w}_{rew}$, is updated using a simple delta rule as follows:

With the established update rules, we are now ready to investigate the SFs’ learning with different initialization methods of the weight matrix $W$. Notably, the initial values of weight matrix $W$ are assumed to play a critical role in the learning performance and efficiency.

To explore the influence of different weight initialization methods on the learning dynamics of SFs, we initialized the weight matrix $W$ using three different methods: identity, zero, and small random matrices.

To investigate the learning process of each RL agent, we used a simple one-dimensional (1D) grid world of size $N\in\mathbb{N}$ spanning from 3 to 100 cells. In this environment, the agent navigates the grid world using left and right actions (Figure 1). In every episode, the agent starts at the leftmost position in the grid world. The ultimate goal is to reach the rightmost end of the world (also known as the terminal state). Upon reaching this terminal state, the agent receives a reward of 1 point. In contrast, all other states receive a score of 0, which means no reward. In our investigation, the discount factor $\gamma$ was set to 0.95.

To maximize the overall discounted reward, the agent selects actions predicated upon the estimated Q value utilizing an $\epsilon$-greedy policy. This policy prescribes a uniform random action selection with probability $\epsilon$, while at other times, with a probability of $1-\epsilon$, the agent chooses the action associated with the highest Q-value estimate. To foster adequate exploration and promote learning stabilization, the probability $\epsilon$ undergoes a decay according to the rule: $\epsilon_{k}=0.9\cdot 0.95^{k}+0.1$, where $k$ signifies the episode index [19].

The learning rates assigned to each learner—the matrix $M$ for the SR agent and the matrix $W$ for the SF agents—–were uniformly set at $\alpha_{M}=\alpha_{W}=0.1$. The learning rate allocated to the reward position vector, for both the SR agent and the SF agents, was fixed at $\alpha_{r}=0.1$. To equitably compare SR and SF agents, maze environment observations were utilized as state indices for the SR agent and one-hot coding vectors for the SF agents.

We evaluated the learning performance of SF agents with different weight initialization methods based on several metrics, including learning speed, final performance, and stability of learning. In addition to these performance metrics, we also analyzed the changes in the SR matrix and the weight matrix throughout the learning process. These analyses allowed us to better understand the dynamics and mechanisms underlying the influence of weight initialization on SF learning.

In this evaluation, each agent simulation test was run 10 times, and the mean and standard deviation of the results are presented in the experimental result section.

To elucidate the impact of disparate initial matrix forms on SF agents, we analyzed the transformational learning history of the weight variables of SF agents, juxtaposing the findings with those of the SR agent.

Drawing on Equations (2) and (4), a direct correlation can be established between the variance in the SR matrix learning and the RL agent’s performance. Herein, we analyze the anticipated value and step length to highlight the performance differences in the learning process.

Initiating the SF weight matrix as an identity matrix provides the agent with a unique starting position in its learning journey about the environment. As learning progresses, each element in the identity matrix corresponds to a particular state, thereby facilitating the updating of knowledge regarding state transitions. However, this initialization method could restrict the agent’s versatility in exploring and learning diverse environmental patterns, potentially resulting in slower learning as observed in previous research. This limitation could be particularly consequential for an agent’s adaptability in increasingly complex or dynamic environments.

Alternatively, initializing the SF weight matrix with zero establishes a ”tabula rasa” situation for the agent. Devoid of any prior knowledge, these agents are heavily influenced by their environmental interactions and the inherent learning algorithm. Although this approach broadens the exploration scope, it may decelerate learning due to the absence of initial guidance. This downside was apparent in studies where agents initiated with a zero matrix took a longer time to converge compared to their randomly initialized counterparts.

Contrastingly, random matrix initialization strikes a balance between exploration and exploitation. Incorporating random elements into the SF weight matrix equips the agent with a degree of ”innate knowledge” guiding its initial steps while preserving a vast spectrum for exploration and learning. Consequently, this initialization method may enhance the learning efficiency, offering a promising avenue for improving SF learning algorithms.

Effective initialization methods, contingent upon the activation function, have been well-documented in the study of Artificial Neural Networks (ANN) utilizing backpropagation algorithms. The normalized Xavier initialization method [17] is typically employed with sigmoid and tanh functions, while the He initialization method [18] sees frequent usage with ReLU functions.

In this study, the absolute values derived from the Xavier and He methods were employed to initialize the random agent. It’s worth noting that inclusion of negative numbers in the weight matrix can result in a negative SR value corresponding to future occupancy, as we make use of a single-layer function approximator devoid of an activation function. To address this issue, we can utilize multilayer ANNs as a function approximator. When a deep neural network is employed as an SF approximator, it begs the question as to which activation function in the hidden layer is optimal, and consequently, the most effective weight initialization method.

Though this paper focused on exploring MDP-based agent learning of environmental characteristics via the SF algorithm, numerous other algorithms are available that describe animal-environment interactions and learning mechanisms. For instance, Particle Swarm Optimization (PSO) that models avian foraging behavior [22, 23]. Among the latest advancements to the PSO algorithm, multi-swarm PSO has been successfully implemented in feature learning for sentiment analysis of Massive Online Open Course lecture reviews [24].

In the context of RL, the feature vector of the input layer offers a snapshot of the agent’s current position. Subsequently, this information is transformed by the SF weight matrix into a population vector, effectively encoding the anticipated future occupancy given the policy at hand. This sequence of operations bears striking resemblance to the neurobiological mechanisms believed to underpin spatial learning.

A collection of studies [7, 5, 4] suggest that hippocampal CA1 place cells encode SR through population codes. Viewed through this neurobiological lens, the SF weight matrix may be interpreted as a close analog to the synaptic weights connecting CA1 place cells to preceding layers of neurons in the neural hierarchy, such as those in the CA3 and entorhinal cortex. This parallel between the functioning of RL algorithms and the neuronal processes that facilitate spatial learning lends support to the use of such algorithms in the investigation of cognition and its underlying biological substrates.

While the exact mechanisms by which the brain might implement the synaptic update rule used in our study remain elusive, a body of research has found substantial evidence that TD learning parallels the activity of dopaminergic neurons in response to reward prediction errors [25, 26]. This aligns with the hypothesis that the neural instantiation of TD learning might be facilitated through neuroplasticity rules, such as spike timing-dependent plasticity and heterosynaptic plasticity [15, 27]. This conjecture, if further corroborated, could add an extra layer of understanding to our exploration of the intersections between artificial intelligence and neurobiology.

From a biological standpoint, it seems reasonable to posit that place coding and reward prediction coding might be processed in tandem within the brain, which subsequently synthesizes these elements into anticipated values for a given state. This line of thought supports the perception of the brain as a device engaging in parallel distributed processing, as suggested by [28]. Underpinning this proposition, the backpropagation algorithm has exhibited exceptional capability in tasks such as image recognition [29, 30]. Moreover, a convolutional neural network (CNN) trained with this algorithm has exhibited activation patterns that bear resemblance to those observed in the visual cortex and the inferior temporal cortex of the brain [31]. Notably, when the activation pattern of a trained CNN was used to manipulate an image, it was found to predict neuronal responses in the V4 visual cortex of macaque monkeys [32]. Nevertheless, it is still a matter of ongoing debate and remains unconfirmed whether the backpropagation algorithm is genuinely operative within the brain [33, 34].

While the biological embodiment of SR learning remains elusive, particularly regarding the brain’s processing location and method for the inner product of the feature vector and reward vector, there is a notable correlation between the outcomes of SR learning and the behavior of hippocampal place cells [35, 36, 4]. However, it warrants further exploration to fully understand how the brain learns and signifies the sequences of state transitions, rewards, and state values. Experimental findings have associated the representation of the reward signal with the orbitofrontal cortex (OFC) [37, 38], suggesting the anterior cingulate gyrus as a probable area for the integration of the OFC’s reward signal and the HPC’s SF signal [39, 40]. In contrast, a study by [41] postulated that the HPC directly encodes the position of the reward.

Transitioning our focus to the question of ’how’, we are confronted with the challenge of extracting a scalar value from the successor feature vector and reward vector [42]. Although this issue extends beyond the scope of the current study, we can glean some neurobiological insights. Specifically, if the synaptic weights in a neural network at the developmental stage are randomly initialized, they demonstrate faster convergence to an optimal state.

While our study offers valuable insights into the impact of SF weight matrix initialization on learning efficiency and convergence, it’s important to note that these findings are based on a one-dimensional grid world. The learning patterns and efficiency we observed may vary with different types of environments. For instance, we noticed more distinct learning trajectories in larger grid worlds, while smaller grid worlds exhibited similar patterns. Therefore, the grid world size is a potential limitation of our study, and our conclusions might be more applicable to larger grid worlds. Future studies could broaden the scope by investigating a wider range of MDP environments, such as two-dimensional grid worlds, to enhance the applicability of our findings.

Our study relied on certain evaluation metrics, including MSE of value error, step length, and PCA of SR place matrix, to analyze learning efficiency and convergence. While these metrics provided significant insights, they might not encapsulate all aspects of an agent’s learning trajectory. For instance, MSE and step length predominantly focus on the speed of learning, potentially overlooking other critical dimensions such as stability and adaptability of learning. Additionally, PCA, while effective in dimensionality reduction and visualizing high-dimensional data, may oversimplify complex learning patterns.

This study embarked on an exploratory journey into the role of matrix initialization in SF learning within the framework of RL. We discovered notable differences in the learning trajectories of agents with different matrix initialization forms - identity, zero, and random (Xavier, He, uniform distribution). Our findings suggest that random matrix initialization, particularly using Xavier and He methods, led to more efficient learning and faster convergence to the optimal state, as evidenced by a quicker decrease in value error and step length. The PCA further revealed distinct patterns of SR place matrix evolution among different agents, reinforcing the importance of matrix initialization in shaping learning dynamics.

The study highlights the significance of weight initialization in the learning process. Our observations demonstrate that the choice of initialization method significantly influences the learning trajectory and efficiency of the agents. Specifically, agents initialized with random matrices demonstrated accelerated learning and quicker convergence to the optimal state. These findings underline the value of exploring diverse initialization techniques to enhance the effectiveness of SF learning.

The implications of this research extend beyond SF learning and RL, contributing to our broader understanding of intelligence from both a neuroscientific and artificial intelligence perspective. By drawing parallels between SF learning and the functioning of place cells in the brain, the study offers intriguing insights into the neurobiological processes underlying learning. An intriguing direction for future research is to delve deeper into the parallels and disparities between biological learning and AI learning algorithms. The results of this study shed light on the learning efficiency of agents, mirroring the learning process of place cells in the brain. Continued efforts to bridge this gap could lead to the development of more biologically-inspired AI models, possibly leading to breakthroughs in our understanding of both artificial and natural intelligence.