A Low Latency Adaptive Coding Spike Framework for Deep Reinforcement Learning

1 Introduction

Deep learning has had a significant impact on many areas of machine learning, including reinforcement learning (RL). The combination of deep learning algorithms and RL has led to the development of a new field called deep reinforcement learning (DRL) [Arulkumaran et al., 2017], which has achieved impressive results in a variety of RL tasks, including some that have reached or even surpassed human-level performance [Mnih et al., 2015; Hu et al., 2021; Silver et al., 2016]. However, DRL using deep neural networks (DNNs) requires a significant amount of resources, which may not always be available. For example, in the case of mobile robots, the processing system needs to be power-efficient and have low latency (inference time) [Niroui et al., 2019; Lahijanian et al., 2018]. As a result, researchers are actively exploring the use of alternative networks for DRL that are more energy-efficient and have lower latency.

Inspired by neurobiology, spiking neural networks (SNNs) use differential dynamics equations and spike information encoding methods to build computing node models in neural networks [Maass, 1997]. Compared with DNNs, SNNs are closer to biological neural networks and could greatly reduce energy consumption with neuromorphic hardware. Therefore, SNNs are well suited as a low-power alternative to DNNs in DRL. But there is still relatively little attempt at using SNNs in RL to achieve high-dimensional control with low latency (inference time steps) and energy consumption.

∗Corresponding author

The existing spiking reinforcement learning (SRL) methods are mainly divided into two categories [Neftci and Averbeck, 2019]. First, guided by revealing reward mechanisms in the brain, several works [Yuan et al., 2019; Bing et al., 2018; Fr´emaux et al., 2013] train spiking neurons with rewardbased local learning rules. These local learning methods can only be applied for simple control with shallow SNNs. Second, oriented toward better performance, replacing DNNs in DRL with SNNs can also implement SRL. Since SNNs usually use the firing rate as the equivalent activation value [Li et al., 2021], which is a discrete value between 0 and 1, it is hard to represent the value function of RL, which does not have a certain range in training. To avoid using discrete spikes to calculate the value function, this type of method is subdivided into three types. Tan [2021] and Patel [2019] use the SNN conversion method to convert the Deep Q-Networks (DQN) [Mnih et al., 2015] into SNNs. Zhang [2022] and Tang [2020] use the hybrid framework, which uses DNNs for value function estimation and SNNs for execution. Chen [2022] and Liu [2022] use surrogate gradients to train SNNs directly and fixed spike coders to estimate the value function.

However, these methods all have their shortcomings. Conversion methods have low energy efficiency and high latency, while the hybrid framework loses the potential biological plausibility of SNNs and limits the application of different types of RL algorithms. Although the directly trained SNNs have high energy efficiency and low latency, they are hard to estimate for continuous-valued functions using fixed coders. They have to find a trade-off between latency and accuracy [Li et al., 2022]. In addition, these SRL methods have a narrow range of applications because they are all limited to a specific online RL algorithm (DQN, TD3 [Fujimoto et al., 2018], and DDPG [Lillicrap et al., 2016]).

To overcome these shortcomings, we propose an adaptive coding method to compress or extend information in the temporal dimension through learnable matrix multiplication. This learnable code offers increased flexibility over fixed coding, allowing for better results with lower latency. Unlike sequence-to-sequence coding in [Shrestha and Orchard, 2018] and [Hagenaars et al., 2021], our proposed adaptive coder implements value-sequence coding. Moreover, our coders do not require a label signal and can automatically select the optimal form according to different tasks.

At the same time, in order to expand the scope of application, we design structures for online and offline SRL algorithms, respectively. Figure 1A shows the flow of the online algorithm, where the state and reward of the environment are directly interacted with the spiking agent after passing through the encoder. SNN calculates the value function and outputs it in the form of spikes, which are converted into real values by decoders. Finally, the action selection is made according to the value function. Offline algorithms are more complex (Figure 1B). It begins by collecting data through behavioral policies that interact with the environment. Then, the spiking agent interacts with the collected data and learns through the same process as online algorithms. We named the proposed SRL model adaptive coding spike framework (ACSF). The main contributions of this paper can be summarized as follows:

• We propose an adaptive coding spike framework for SRL that uses adaptive coders to represent states and estimate value functions. This framework reduces the latency of existing methods effectively while maintaining optimal performance.
• To our best knowledge, the ACSF is the first complete, directly trained SRL framework. It supports both online and offline RL algorithms, expanding the application range of SRL.
• For the RL scale, ACSF provides a low-power inference scheme for RL algorithms (up to 5X energy efficient). This provides potential assistance for mobile robot control using RL algorithms.
• From the experimental results, compared with conversion and other directly trained methods, ACSF has better versatility and lower latency (T = 4). At the same time, it achieves the best performance in all conditions.

2 RelatedWorks

2.1 Reward-based Local Learning

To explore the reward mechanism in the brain, applying reward-based training rules to the local learning of spiking neurons [Yuan et al., 2019; Bing et al., 2018; Fr´emaux et al., 2013; Legenstein et al., 2008] was the mainstream approach for early SRL. However, they have the common problem that they are limited to shallow SNNs and are hard to apply to complex tasks. E-prop [Bellec et al., 2020] successfully advances SRL to complex tasks with good performance. However, it needs to tune a lot of hyperparameters, making it difficult to apply to different RL tasks.

2.2 Convert DNNs to SNNs for RL

The conversion algorithms map the DNNs to the SNNs by matching the firing rate and activation values. It is a commonly used method for obtaining high-performance SNNs. Several works [Tan et al., 2021; Patel et al., 2019] obtain Deep Spiking Q-Networks (DSQN) using conversion methods. To ensure the accuracy of mapping, those methods often have high latency (about 500 time steps) and cannot outperform the source DQNs.

2.3 Hybrid Framework of SNNs and DNNs

In RL, some algorithms use the actor-critic architecture [Sutton et al., 1999], where value function estimation and action selection are implemented by critic and actor networks, respectively. Some works [Zhang et al., 2022; Tang et al., 2020] use deep critic networks and spiking actor networks to implement SRL. These methods can be applied to complex tasks with lower latency, but they only work for actor-critic structures, and they rely on DNNs, which makes them difficult to fully apply to neuromorphic chips.

2.4 Directly Trained SNNs for RL

Surrogate gradient methods introduce backpropagation (BP) into SNNs successfully and make it easy to train deep SNNs directly. Recently, a few papers [Chen et al., 2022; Liu et al., 2022] have used directly trained SNNs and fixed coders to construct DSQN. Although these methods have high energy efficiency, they do not form a complete framework and have a poor trade-off between accuracy and latency. In addition, their applications are also limited to discrete action-space environments.

3 Methods

In this section, we first introduce the DRL algorithms and the spiking neuron model. Then, we describe the adaptive coding method in detail. Finally, we derive the gradient of the direct training method. The overall architecture of ACSF is shown in Figure 2.

3.1 DRL Algorithms

In this section, we will introduce different baseline algorithms (DQN, DDPG, BCQ and behavioral cloning) used in ACSF. Due to space limitations, we only show the final loss functions. Please refer to the supplementary materials for more details. Assume the agent's current environment sampling result is state s, action a, reward r and next state s′ . Mark the reward discount factor as γ.

DQN

DQN is a typical RL algorithm, and it only needs to learn the value function (q-values). Its loss function [Girshick, 2015] is:

where Qθ1 and Qθ2 is the current and target q-networks, respectively.

DDPG

DDPG is based on AC architecture, including actor networks that select actions and critic networks that calculate value functions. The loss functions of the actor and critic networks are:

where Qθ1 and Qθ2 are current and target critic networks; Aϕ1 and Aϕ2 are current and target actor networks. N is the batch size.

BCQ

BCQ is an offline RL algorithm. In addition to training double critic networks Q1θ1 , Q2θ1 , target critic networks Q1θ2 , Q2θ2 , current perturbation network ξδ1 and target perturbation network ξδ2 , we also need to train a variational auto-encoder (VAE) [Kingma and Welling, 2014] Gω = {Eω1 ,Dω2 }. Their loss functions can be written as follows [Fujimoto et al., 2019]:

Behavioral Cloning

Behavioral cloning (BC) is an offline imitation learning method used to fit the distribution of buffered data D. We use a multilayer perceptron (MLP) Pκ in this work. The input of the MLP model is raw states (s), and the output is action (a). We use the mean squared error (MSE) as the loss function:

where m is the dimension of action.

3.2 Iterative LIFModel

In this work, we employ the iterative leaky integrate-andfire (LIF) neuron model to drive the proposed ACSF due to its biological plausibility and computational tractability. Its dynamics can be described by three processes: charge, discharge, and reset. The corresponding equations are:

where the superscript l indicates the l-th layer, the subscript t represents the t-th time step; V and Vˆ denote the membrane potential before and after reset respectively; β is the decay factor; O is the output spike; Θ(·) is the firing function; Vr is the reset voltage; ϑ is the firing threshold. The schematic diagram of neuron dynamics is shown in Figure 3.

3.3 Adaptive Coders

As mentioned earlier, existing directly trained SRL methods have limited application and poor trade-offs between performance and latency. The key to solving these problems is to design a common adaptive coder with learnable weights that allows the model to be adapted to different algorithms and environments. In addition, adaptive coding has better characterization capabilities than fixed coding and can achieve better performance with lower latency.

Spike Encoder

There are many fixed encoding methods in SRL, such as rate coding, population coding, and temporal coding. Although these methods have had some success, they are frequently rigid, with strict time step length requirements. To increase flexibility and reduce latency, we encode the raw state S with a learnable matrix WE (Figure 2). We use matrix multiplication to expand S in the temporal dimension by T times, yielding the temporal state Sτ . The temporal state Sτ is then fed into the spiking neurons, which produce spike trains Ot. The encoding strategy at the t-th time step can be stated as follows:

where S and Sτ denote the state and temporal state, respectively; Ot is the output spike at the t-th time step; ∗ represents the matrix multiplication; and WE is the learnable matrix that can be jointly optimized.

Decoders

The decoding method used in SRL is usually rate decoding. The accuracy of this decoding method is proportional to the length of the time step, so it is difficult to strike a good trade-off between performance and latency. To improve the performance of decoding at low latency, we use an operation similar to adaptive encoding. We multiply the spike trains with a learnable matrix WD to compress them T times in the time dimension: {︃Q = Qτ ∗WVD

where Qτ and Aτ are output spike trains of the value network and action network, respectively; WV D and WAD are decode matrices of the value decoder and action decoder; and Q/A are the value function and actions selection of RL. Specially, for RL algorithms with non-AC architectures such as DQN, we only decode the Q values.

3.4 Direct Training with Surrogate Gradients

Compared with SNNs obtained using conversion methods, directly trained SNNs tend to have lower latency and higher inference energy efficiency. However, the non-differentiable property of spikes makes directly training SNNs difficult. The surrogate gradient method solves this problem by using alternative activation functions and is widely used to train deep SNNs [Wu et al., 2018]. We adopt the arctangent function as the surrogate gradient function:

where α is a hyperparameter that controls the width of the surrogate function.

Suppose the loss function is L, according to Eq. 7 and the chain rule, for the decoding matrix, we have:

and for the final layer L:

where ∂LQ ∂Q and ∂LA ∂A are defined according to different algorithms; Qτt and Aτt are output spike trains, equal to OL t , so we only need to calculate ∂OL t ∂W L :

and for the hidden layers l ∈ {1, · · · ,L - 1}, according to Eq. 5:

where T is the total time steps.

The first factor in Eq. 12 could be derived as:

The second factor in Eq. 12 could be derived as:

The third factor in Eq. 12 could be derived as:

where:

Let:

Then:

Bring Eq. 13, Eq. 14 and Eq. 18 into Eq. 12, we can get the derivative of the loss functions in the hidden layers:

4 Experiments

In this section, we apply the ACSF to both online and offline methods and evaluate them in different environments. We compared the ACSF with the latest SRL methods to prove that the ACSF can achieve a better trade-off between performance and latency. We also conduct ablation studies to demonstrate the superiority of adaptive coders.

4.1 Experimental Settings

Environment

To evaluate the proposed ACSF, we apply it to sixteen Atari 2600 games from the arcade learning environment (ALE) [Bellemare et al., 2013] and four MuJoCo [Todorov et al., 2012] continuous robot control tasks from the OpenAI gym [Brockman et al., 2016]. We use the most stable version of the Atari game without frame skipping (NoFrameskipv4) and the latest version of the MuJoCo environment (-v3). The visualization of environments is shown in Figure 4. The training steps for Atari games and MuJoCo environments are 50 million (50M) and 1 million (1M), respectively, and the evaluation interval is 50 thousand (50K) and 5 thousand (5K).

Training and Evaluating

Baseline methods DQN, DDPG, BCQ, BC, and the proposed ACSF were all trained using the same structure as in [Fujimoto et al., 2019]. Our results are reported over 5 random seeds [Duan et al., 2016] of the behavioral policy, the OpenAI gym simulator, and network initialization. We save the model with the highest average reward for testing.

We evaluate the ACSF on Atari games by playing 10 rounds of each game with an ϵ-policy (ϵ = 0.05). Each episode is evaluated for up to 18K frames per round. To test the robustness of the agent in different situations, the agent starts with a random number (at most 30 times) of no-op actions. For MuJoCo environments, we played 10 rounds of each environment without exploration, and each test episode lasted for a maximum of 1K frames.

4.2 Results on Atari Games

To test the performance of ACSF on game controls (a discrete action-space environment), we apply it to the DQN algorithm and test it on 16 different Atari games. We also compare the ACSF with some other DSQN methods [Tan et al., 2021; Liu et al., 2022; Chen et al., 2022] to show it can achieve an excellent trade-off between latency and performance.

Figure 5: Learning curves for DQN and ACSF. During the training process, the performance of ACSF meets or exceeds that of the DQN algorithm. The learning curves have been smoothed for aesthetics.

The experimental results and comparisons are shown in Table 1 and Figure 5. The full learning curves are included in the supplementary material. Experimental results show that ACSF outperforms the original DNQ algorithm. Meanwhile, compared to other SNN-based methods, ACSF maintains similar or better performance and reduces latency by more than 50%. This also proves that adaptive coding performs better on RL tasks than fixed coding.

4.3 Results onMuJoCo Robot Control Tasks

The ACSF is a complete SRL framework that can be applied to both online and offline algorithms. Therefore, we apply ACSF to offline algorithms and conduct experiments in a discrete action-space environment to prove the versatility of ACSF. We compare the performance of ACSF with three baseline algorithms (DDPG, BC, BCQ). The results are shown in Table 2 and Figure 6. From the results, ACSF is also suitable for offline algorithms, and the adaptive coders improve the performance significantly. Even at ultra-low latency (T = 4), ACSF outperformed all the baseline algorithms.

The performance of ACSF will also vary under different time step settings because the latency (T) affects the dimension of the learnable matrix in adaptive coders, and thus affects the fitting ability of coders. According to experimental results, ACSF can achieve its best performance when T = 8 or T = 10. We can select the value of T flexibly according to the usage scenario.

4.4 Ablation Study

To verify that adaptive coders can improve performance effectively. We conducted ablation experiments to compare the performance of algorithms with different coding methods. We compared five settings:

• No coders+DNN: The original BCQ algorithm.
• Rate coders+SNN: Compute the firing rate of spikes as the value function.
• Accumulate coders+SNN: Modify the dynamics of the output neurons to make them stop firing spikes. Record the voltage at the last time step as the value function.
• Adaptive coders+DNN: On the basis of the original DNN, adaptive coders are added to eliminate the bias caused by the structural differences introduced by the adaptive coders.

Table 3: The performance of different coders.

• Adaptive coders+SNN: Combines adaptive coders and directly-trained SNNs (the proposed ACSF).

Table 3 displays the results. Rate encoding does not work well due to very short time steps. The accumulate coders lose their spike characteristics, so the performance is better and similar to that of pure DNN. The combination of adaptive coders and DNN is slightly superior to pure DNN, which also proves that adaptive coders can integrate state information and calculate the value function accurately. The combination of adaptive coders and SNN has optimal performance in different environments, demonstrating the effectiveness of ACSF in low-latency settings.

4.5 Power Estimation and Comparison

To calculate the power consumption of SNNs on hardware, we count the total synaptic operations (SynOps) [Wu et al., 2021]. It is defined as follows:

where fout denotes the number of outgoing connections to the subsequent layer, T denotes the total time steps, L is the total number of layers, N l means the number of neurons in layer l, and s denotes the spikes.

For DNNs:

where f l in represents the number of incoming connections to each neuron in layer l.

The total SynOps are shown in Table 4. The ACSF improves energy efficiency (up to 5X) on both algorithms, providing a more efficient application method for DRL.

5 Conclusion

In this paper, we propose a low latency adaptive coding spike framework (ACSF) for SRL. The ACSF solves the high latency problem in SRL by using adaptive coding. At the same time, ACSF has expanded the application range of SRL by designing structures for online and offline algorithms. In terms of versatility, ACSF is the first model to achieve compatibility with both online and offline algorithms and achieve optimal performance in both discrete and continuous environments. In addition, ACSF uses a directly trained SNN, which does not require DNNs, so it can be easily migrated to neuromorphic hardware. Extensive experiments show that ACSF can achieve the same or higher performance with ultralow latency compared to other SRL methods. Meanwhile, compared with traditional DRL methods, ACSF has a significant improvement on inference energy efficiency (up to 5X) while maintaining better performance.

Table 4: Total synaptic operations.

Acknowledgements

This work was supported by National Key Research and Development Program of China under Grant No. 2020AAA0105900, National Natural Science Foundation of China under Grant No. 62236007 and U2030204 and the Key Research Project of Zhejiang Lab under Grant No. 2021KC0AC01.