Spatial-aware decision-making with ring attractors in Reinforcement Learning systems

Ring attractors commonly consist of a configuration of excitatory and inhibitory neurons arranged in a circular pattern. We can model the dynamics of the ring using the Touretzky ring attractor network (Touretzky, 2005). In this model, each excitatory neuron establishes connections with all other excitatory neurons, and an inhibitory neuron is placed in the middle of the ring with equal weighted connections to all excitatory neurons. This creates a network that facilitates complex information processing.

Excitatory neurons’ input signal. Let $x^{i}_{n}\in\mathbb{R}^{s}$ denote the input signal from source $i$ to the excitatory neuron $n=1,\ldots,N$. The total input to neuron $n$ is defined as the sum of all input signals $I$ for that particular neuron: $x_{n}=\sum_{i=1}^{I}x^{i}_{n}$, where $x_{n}\in\mathbb{R}$. To model input signals $x^{i}_{n}$ of varying strengths, these signals are commonly viewed as Gaussian functions $x^{i}_{n}:\mathbb{R}^{s}\rightarrow\mathbb{R}^{s}$. These functions allow us to represent the input to each neuron as a sum of weighted Gaussian distributions. The key parameters of these Gaussian functions are: $K_{i}$, the magnitude variable for the input signal $i$, which determines the overall strength of the signal; $\mu_{i}$, which defines the mean position of the the Gaussian curve in the ring for the input signal $i$, representing the central focus of the signal; $\sigma_{i}$, the standard deviation of the Gaussian function, which determines the spread or reliability of the signal; and $\alpha_{n}$, which represents the preference for the orientation of the neuron $n$ in space. These parameters combine to the following:

$$x_{n}(K_{i})=\sum_{i=1}^{I}x^{i}_{n}(\alpha_{n})=\sum_{i=1}^{I}{\frac{K_{i}}{\sqrt{2\pi\sigma_{i}}}\exp\left(-\frac{1}{2}\frac{(\alpha_{n}-\mu_{i})^{2}}{\sigma_{i}^{2}}\right)},\quad n=1,2,\ldots,N$$

(1)

Neuron activation function. We employ the rectified linear unit (ReLU) function $f(x)=\max(0,x+h)$, where $h\in\mathbb{R}^{+}$ as the activation function for each neuron, where $h$ is a threshold that introduces the non-linear behaviour in the ring.

Excitatory neuron dynamics. The dynamics of excitatory neurons in the ring is described as:

$$\frac{dv_{n}}{dt}\approx\frac{\Delta v_{n}}{\Delta t}=\frac{v_{n+\Delta t}-v_{n}}{\Delta t}=\frac{f(x_{n}+e_{n}+i_{n})-v_{n}}{\tau}.$$

(2)

In Eq. 2 $v_{n}\in\mathbb{R}$ represents the activation of the excitatory neuron $n$, $x_{n}$ is previously defined in Eq. 1 is the external input to the neuron $n$ of Eq. 1, $e_{n}\in\mathbb{R}$ represents the weighted influence of the other excitatory neurons activation, which is defined mathematically in Eq. 4. $i_{n}\in\mathbb{R}$ is the influence from the weighted inhibitory neuron activation to the target excitatory neuron $n$, and $\tau=\Delta t$ is the time integration constant. This equation captures the evolution of neuronal activation over time, considering both excitatory and inhibitory activations.

Inhibitory neuron dynamics. The activation of the inhibitory neuron, which regulates network dynamics, is described by:

$$\frac{du}{dt}\approx\frac{\Delta u}{\Delta t}=\frac{u_{+\Delta t}-u}{\Delta t}=\frac{f(e_{n}+i_{n})-u}{\tau}$$

(3)

Here, $u\in\mathbb{R}$ represents the inhibitory neuron’s activation output, $e_{n}$ is the weighted sum of excitatory activations where in this case $n$ is the inhibitory neuron, and $i_{n}\in\mathbb{R}$ is the weighted self-inhibition activation term. This equation models how the inhibitory neuron integrates inputs from the excitatory population and its own state.

Synaptic weighted connections: The influence between neurons decreases with distance, as modeled by the weighted connections. This weighted connection applies to both excitatory and inhibitory neurons. For excitatory neurons: $w^{(E_{m}\rightarrow E_{n})}=e^{-d^{2}_{(m,n)}}$, where $d_{(m,n)}=|m-n|$ is the distance between neurons $m$ and $n$. For the inhibitory neuron: $w^{(I\rightarrow E_{n})}=e^{-d^{2}_{(m,n)}}$. Note that our model contains a single inhibitory neuron placed in the middle of the ring, with a distance of 1 unit to all excitatory neurons.The excitatory ($e_{n}$) and inhibitory ($i_{n}$) weighted connections are also known in the literature as neuron-proximal excitatory and inhibitory voltage or potential.These are then defined as follows:

$\displaystyle e_{n}=\sum_{m=1}^{N}w^{(E_{m}\rightarrow E_{n})}v_{m}\ \ \ \ \ i_{n}=w^{(I_{m}\rightarrow E_{n})}u$

(4)

Full excitatory neuron dynamics. Combining all influences, the complete dynamics of excitatory neurons are described by:

$$\frac{dv_{n}}{dt}=\frac{1}{\tau}\left(f\left(\sum_{m=1}^{N}w^{(E_{m}\rightarrow E_{n})}_{m,n}v_{m}+x_{n}+w^{(I\rightarrow E_{n})}u\right)-v_{n}\right)$$

(5)

This equation updates the activation for all neurons based on recurrent excitation, external input from the summation of the input signals for neuron $n$, $x_{n}$; and inhibitory influence.

Full inhibitory neuron dynamics. The complete dynamics of the inhibitory neuron are given by:

$$\frac{du}{dt}=\frac{1}{\tau}\left(f\left(w^{(I\rightarrow I)}u+\sum_{m=1}^{N}w^{(E_{m}\rightarrow I)}_{m}v_{m}\right)-u\right)$$

(6)

This equation models how the inhibitory neuron integrates self-inhibition and excitatory inputs from the entire network. These equations collectively describe the complex dynamics of the ring attractor network, capturing the interplay between excitatory and inhibitory neurons, external inputs, and synaptic connections.

To integrate the ring attractor model with RL, we need to establish a connection between the estimated value of state-action pairs and the input to the ring attractor network. This integration allows the ring attractor to serve as a behavior policy, guiding action selection based on the values learned. We begin by reformulating the input function for a target excitatory neuron $n$. The key modification is setting the scale factor $K_{i}$ to the Q-value $Q(s,a)$ of the state-action pair $(s,a)$, that is $K_{i}=Q(s,a)$.

This formulation ensures that actions with higher estimated values are given more weight in the ring attractor dynamics, naturally biasing the network towards more valuable actions. The orientation of the signal within the ring attractor is determined by the direction of movement in the action space. We represent this as $\mu_{i}=\alpha(a)$, where $\alpha(a)$ is the angle corresponding to the action $a$ in the circular action space. We define our circular action space $\mathcal{A}$ as a subset of $\mathbb{R}^{2}$, where each action $a\in\mathcal{A}$ is represented by a point on the unit circle. The function $\alpha:\mathcal{A}\rightarrow[0,2\pi)$ maps each action to its corresponding angle on this circle. To account for uncertainty in our value estimates, we incorporate the variance of the estimated value for each action into our model: $\sigma_{i}=\sigma_{a}$.

This allows the network to represent not just the expected value of actions, but also our confidence in those estimates. Combining these elements, we arrive at the following equation for the action signal $x(a)$:

$$x_{n}(Q)=\sum_{a=1}^{A}{\frac{Q(s,a)}{\sqrt{2\pi\sigma_{a}}}\exp\left(-\frac{1}{2}\frac{(\alpha_{n}-\alpha_{a})^{2}}{\sigma_{a}^{2}}\right)}$$

(7)

This equation represents the input to each neuron as a sum of Gaussian functions, where each function is centered on an action’s direction and scaled by its Q-value. The dynamics of the excitatory neurons in the ring attractor, now incorporating the Q-value inputs, are described by:

$$\begin{gathered}\frac{dv_{n}}{dt}=\frac{1}{\tau}\left(max\left(0,\left(\sum_{m=1}^{m=N}w^{(E\xrightarrow{}E)}_{m,n}v_{m}+x_{n}(Q)+w^{I\xrightarrow{}E_{u}}u\right)\right)\right)-v_{n}\end{gathered}$$

(8)

This equation captures how the activation of each neuron evolves over time, influenced by the action-value functions $x_{n}(Q)$, and both excitatory and inhibitory feedback. This equation captures how the activation of each neuron evolves over time, influenced by the action-value input Gaussian functions $x_{n}(Q)$, the excitatory feedback $e_{n}=\sum_{m=1}^{N}w^{(E_{m}\rightarrow E_{n})}v_{m}$, and the inhibitory feedback $i_{n}(=w^{(I\rightarrow E_{n})}u$.

To translate the ring attractor’s output into an action in the 2D space, we use the following equation:

$$\text{action}=\operatorname*{arg\,max}_{n}\{\mathbf{V}\}\cdot\frac{N^{(A)}}{N^{(E)}}$$

(9)

where $n\in\{1,...,N^{(E)}\}$, $N^{(E)}$ is the number of excitatory neurons in the ring attractor, $N^{(A)}$ is the number of discrete actions in the action space $\mathcal{A}$, $\mathbf{V}=[v_{1},v_{2},...,v_{N^{(E)}}]$.

This equation assumes that both the neurons in the ring attractor and the actions in the action space are uniformly distributed. This approach allows for nuanced action selection that takes into account both the spatial relations between actions and their estimated values. A visualization of the ring is presented in Fig. 1.

In the field of DRL, for any state-action pair $(s,a)$, the Q-value $Q(s,a)$ can be expressed as a function approximation algorithm $\Phi_{\theta}(s)$ taking as input the current state $s$. This function approximation algorithm ($\Phi_{\theta}(s)$) can be expressed as the weight matrix of our function approximation algorithm transposed $\theta^{T}$ times the feature vector extracted from the input state $x(s)$: $Q(s,a)=\Phi_{\theta}(s)=\theta^{T}x(s)$. As stated in Section 3.1.2, the variance of the Gaussian functions, input to the ring attractor, will be given by the variance of the estimate value for that particular action $\sigma_{i}=\sigma_{a}$.

Among the diverse methods to compute the uncertainty of the action ($\sigma_{a}$) we have chosen to compute a posterior distribution with Bayesian linear regression (BLR). BLR acts as output layer for our neural network (NN) of choice. We choose a linear regression model because it does not compromise the efficiency of the NN, while at the same time it provides a distribution to compute the variance for the state-action pairs. The implementation is based on Azizzadenesheli et al. (2018), where a Bayesian value-based DQN model was instantiated to output an uncertainty-aware prediction for the state-action pairs. With this approach, the new $Q$ function is defined as:

$$Q(s,a)=\Phi_{\theta}(s)^{T}w_{a},$$

(10)

where $w_{a}$ are the weights from the posterior distribution of the BLR model. Eq. 10 represents the parameters of the final Bayesian linear layer.

When provided with a state transition tuple $(s,a,r,s^{\prime})$, where $s$ is the current state, $a$ is the action taken, $r$ is the reward received, and $s^{\prime}$ is the next state. This tuple represents a single step of interaction between the agent and the environment in the RL framework. The model learns to adjust the weights $w_{a}$ of the BLR and the function approximation algorithm, i.e. neural networks (NNs), ($\Phi_{\theta}$) to align the Q values with the optimal action $a=argmax(\gamma\Phi_{\theta}(s)^{T}w_{a})$, equation 11.

$$Q(s,a)=\Phi_{\theta}(s)^{T}w_{a}\xrightarrow{}y:=r+\gamma\Phi_{\theta_{\text{target}}}(x^{\prime})^{T}w_{\text{target}}\hat{a}$$

(11)

where $\gamma$ is the discount factor, $\theta_{target}$ refers to the parameters of the target function approximation algorithm used for learning, and $\hat{a}$ is the predicted optimal action in the next state $s^{\prime}$. The construction of the Gaussian BLR prior distribution and the weights $i$ sample $w_{a,i}$ collected from the posterior distribution are performed through Thompson Sampling. This process allows us to incorporate uncertainty into our action-value estimates. For details on the construction of the Gaussian prior distribution and the specifics of the sampling process, we refer readers to Azizzadenesheli et al. (2018). Both the mean $\mu_{a}$ and the variance $\sigma_{a}$ from Eq. 12 are calculated from a finite number of samples $I$.

	$\displaystyle\mu_{a}$	$\displaystyle=\frac{\sum^{i=I}_{i=0}w_{a,i}^{T}\Phi_{\theta}(s_{t})}{I}=\frac{\sum^{i=I}_{i=0}{Q(s,a)_{i}}}{I}$		(12)
	$\displaystyle\sigma^{2}_{a}$	$\displaystyle=\frac{\sum^{i=I}_{i=0}{\left(w_{a,i}^{T}\Phi_{\theta}(s_{t})-\mu_{a}\right)}^{2}}{I-1}$

To shape circular connectivity within a RNN, the weighted connections in the input signal $V(s)$ and the hidden state or attractor state $U(v)$ are computed as follows:

	$\displaystyle V(s)_{m,n}$	$\displaystyle=\frac{1}{\tau}\Phi_{\theta}(s)^{T}w^{I\xrightarrow{}H}_{m,n}=\frac{1}{\tau}\Phi_{\theta}(s)^{T}e^{\frac{d(m,n)}{\lambda}}$		(13)
	$\displaystyle d(m,n)$	$\displaystyle=\min{(abs(m-n\frac{M}{N}),N-abs(m-n\frac{M}{N}))}$
	$\displaystyle U(v)_{m,n}$	$\displaystyle=h(v)^{T}w^{H\xrightarrow{}H}_{m,n}=h(\Phi_{\theta})^{T}e^{\frac{d(m,n)}{\lambda}}$
	$\displaystyle d(m,n)$	$\displaystyle=\min{(abs(m-n),N-abs(m-m))}$

This circular structure mimics the arrangement of excitatory neurons in the ring attractor. Eq. 13 shows the input signal to the recurrent layer $V_{m,n}$ from neuron $m$ from the previous layer in the DL agent to neuron $n$ in the RNN. The hidden state, $U_{m,n}$ mimics an attractor state, representing the recurrent connections in the RNN.. The weighted RNN connections input-to-hidden and hidden-to-hidden $w^{I\xrightarrow{}H}_{m,n}$, $w^{H\xrightarrow{}H}_{m,n}$ depend on a parameter $\lambda$ that drives the decay of the potential over distance and distance between neurons $d(m,n)$ where $N$ is the total number of neurons for the RNN and $M$ is the count of neurons in the previous layer of the NN architecture. The function $\phi_{\theta}(s):\mathbb{R}^{S}\rightarrow\mathbb{R}^{M}$ maps the input state $s$ of the DL agent to a representation of characteristics that will be the input of the recurrent layer. Likewise, $\theta$ represents the parameters of this function (i.e., the weights and biases of the NN layers preceding the RNN layer, which extract relevant features from the input). The function $h(v):\mathbb{R}^{N}\rightarrow\mathbb{R}^{N}$ is a parameterized by learnable weights transformation that maps the information from previous forward passes into the current hidden state. The learnable parameter $\tau$ is the positive time constant responsible for the integration of signals in the ring. It defines the contribution of input states $\phi_{\theta}(s)$ to the current hidden state, imitating the attractor state, applied to neural networks.

Finally, the action-value function $Q(s,a)$ is derived from the RNN layer’s output by applying the neurons activation function to the combined input $V(s)$ and hidden state information $U(s)$. The activation function of choice is a hyperbolic tangent $\tanh$, this function is symmetric around zero, leading to faster convergence and stability. However, the output range of $\tanh$ (-1 to 1) is not fully compatible with value-based methods, where the DL agents needs to output action-value pairs in the range of the environment’s reward function. To address this issue and prevent saturation of the $\tanh$ activation function, we scale the action-value pairs by multiplying them with a learnable scalar $\beta$, as $Q(s,a)=\beta h_{t}=\beta\tanh((V(s)+U(v))=\beta\tanh((\frac{1}{\tau}\Phi_{\theta}(s_{t})^{T}w^{I\rightarrow H}+h_{t-1}(v)^{T}w^{H\rightarrow H}))$.