Dynamic Sparse Training for Deep Reinforcement Learning

1 Introduction

Deep reinforcement learning (DRL) has achieved remarkable success in many applications. The power of deep neural networks as function approximators allows RL agents to scale to environments with high-dimensional state and action spaces. This enables high-speed growth in the field and the rise of many methods that improve the performance and stability of DRL agents [Wang et al., 2020]. While the achieved performance is impressive, a long training time is required to obtain this performance. For instance, it took more than 44 days to train a Starcraft II agent using 32 third-generation tensor processing units (TPUs) [Vinyals et al., 2019]. The very long training time leads to high energy consumption and prohibitive memory and computation costs. In this paper, we ask the following question: Can we provide efficient DRL agents with less computation cost and energy consumption while maintaining superior performance?

Few recent works attempt to accelerate the inference time of DRL agents via pruning [Livne and Cohen, 2020] or training a compact network under the guidance of a larger network (knowledge distillation) [Zhang et al., 2019]. Despite the computational improvement achieved at inference, extensive computations throughout the training of dense networks are still consumed. Our goal is to accelerate the training process as well as the inference time of DRL agents.

The long training time of a DRL agent is due to two main factors: (1) the extensive computational cost of training deep neural networks caused by the very high number of network parameters [Jouppi et al., 2017] and (2) the learning nature of a DRL agent in which its policy is improving through many trial-and-error cycles while interacting with the environment and collecting a large amount of data. In this paper, we introduce dynamic sparse training (DST) [Mocanu et al., 2021; Hoefler et al., 2021] in the DRL paradigm for the first time to address these two factors 1 . Namely, we propose an efficient training approach that can be integrated with existing DRL methods. Our approach is based on training sparse neural networks from scratch with a fixed parameter count throughout training (1). During training, the sparse topology is optimized via adaptation cycles to quickly adapt to the online changing distribution of the data (2). Our training approach enables reducing memory and computation costs substantially. Moreover, the quick adaptation to the new samples from the improving policy during training leads to a faster learning speed.

In fact, the need for neural networks that can adapt, e.g., change their control policy dynamically as environmental conditions change, was broadly acknowledged by the RL community [Stanley, 2003]. Although prior works related to the automatic selection of function approximation based on neuroevolution exist [Heidrich-Meisner and Igel, 2009], perhaps the most connected in the spirit to our proposed method is a combination of NeuroEvolution of Augmenting Topologies (NEAT) [Stanley and Miikkulainen, 2002] and temporal difference (TD) learning (i.e., NEAT+Q [Whiteson and Stone, 2006]). Still, the challenge remains, and cutting-edge DRL algorithms do not account for the benefits of adaptive neural networks training yet.

Our contributions in this paper are as follows:

• The principles of dynamic sparse training are introduced

1 Code is available at: https://github.com/GhadaSokar/DynamicSparse-Training-for-Deep-Reinforcement-Learning. The extended version of this paper is available at https://arxiv.org/abs/2106.04217.

in the deep reinforcement learning field for the first time.

• Efficient improved versions of two state-of-the-art algorithms, TD3 [Fujimoto et al., 2018] and SAC [Haarnoja et al., 2018a], are obtained by integrating our proposed DST approach with the original algorithms.
• Experimental results show that our training approach reduces the memory and computation costs of training DRL agents by 50% while achieving superior performance. Moreover, it achieves a faster learning speed, reducing the required training steps.
• Analysis insights demonstrate the promise of dynamic sparse training in advancing the field and allowing for DRL agents to be trained and deployed on low-resource devices (e.g., mobile phones, tablets, and wireless sensor nodes) where the memory and computation power are strictly constrained.

2 RelatedWork

Sparsity in DRL. To the best of our knowledge, the current advance in deep reinforcement learning is achieved using dense neural networks. Few recent studies have introduced sparsity in DRL via pruning. PoPS [Livne and Cohen, 2020] first trains a dense teacher neural network to learn the policy. This dense teacher policy network guides the iterative pruning and retraining of a student policy network via knowledge distillation. In [Zhang et al., 2019], the authors aim to accelerate the behavior policy network and reduce the time for sampling. They use a smaller network for the behavior policy and learn it simultaneously with a large dense target network via knowledge distillation. GST [Lee et al., 2021] was proposed as an algorithm for weight compression in DRL training by simultaneously utilizing weight grouping and pruning. Some other works [Yu et al., 2019; Vischer et al., 2021] studied the existence of the lottery ticket hypothesis [Frankle and Carbin, 2018] in RL, which shows the presence of sparse subnetworks that can outperform dense networks when they are trained from scratch. Pruning dense networks increases the computational cost of the training process as it requires iterative cycles of pruning and retraining. This work introduces the first efficient training algorithm for DRL agents that trains sparse neural networks directly from scratch and adapts to the changing distribution.

Dynamic Sparse Training (DST). DST is the class of algorithms that train sparse neural networks from scratch and jointly optimize the weights and the sparse topology during training. This direction aims to reduce the computation and memory overhead of training dense neural networks by leveraging the redundancy in the parameters (i.e., being overparametrized) [Denil et al., 2013]. Efforts in this line of research are devoted to supervised and unsupervised learning. The first work in this direction was proposed by [Mocanu et al., 2018]. They proposed a Sparse Evolutionary Training algorithm (SET) that dynamically changes the sparse connectivity during training based on the values of the connections. SET achieves higher performance than dense models and static sparse networks trained from scratch. The success of the SET algorithm opens the path to many interesting DST methods that bring higher performance gain. These algorithms differ from each other in the way the sparse topology is adapted during training [Mostafa and Wang, 2019; Evci et al., 2020; Jayakumar et al., 2020; Liu et al., 2021; Sokar et al., 2021].

In this work, we adopt the topological adaptation from the SET method in our proposed approach. The motivation is multifold. First, SET is simple yet effective; it achieves the same or even higher accuracy than dense models with high sparsity levels across different architectures (e.g., multi-layer perceptrons, convolutional neural networks, restricted Boltzmann machines). Second, unlike other DST methods that use the values of non-existing (masked) weights in the adaptation process, SET uses only the values of existing sparse connections. This makes SET truly sparse and memory-efficient [Liu et al., 2020]. Finally, it does not need high computational resources for the adaptation process. It uses readily available information during the standard training. These factors are favorable for our goal to train efficient DRL agents suitable for real-world applications. We leave evaluating other topological adaptation strategies for future work.

3 ProposedMethod

In this section, we describe our proposed method, which introduces dynamic sparse training for DRL. Here, we focus on integrating our training approach with one of the state-of-theart DRL methods; Twin Delayed Deep Deterministic policy gradient (TD3) [Fujimoto et al., 2018]. We named our new approach Dynamic Sparse TD3 or "DS-TD3" for short. TD3 is an efficient DRL method that offers good performance in many tasks [Shi et al., 2020; Woo et al., 2020]. Yet, our approach can be merged into other DRL methods as well. The integration with soft actor-critic (SAC) [Haarnoja et al., 2018a] is presented in the extended version of the paper.

TD3 is an actor-critic method that addresses the overestimation bias in previous actor-critic approaches. In actor-critic methods, a policy π is known as the actor, and a state-value function Q is known as the critic. Target networks are used to maintain fixed objectives for the actor and critic networks over multiple updates. In short, TD3 limits the overestimation bias using a pair of critics. It takes the smallest value of the two critic networks to estimate the Q value to provide a more stable approximation. To increase the stability, TD3 proposed a delayed update of the actor and target networks. In addition, the weights of the target networks are slowly updated by the current networks by some proportion τ . In this work, we aim to dynamically train the critics and actor networks along with their corresponding target networks from scratch with sparse neural networks to provide efficient DRL agents. In the rest of this section, we will explain our proposed DST approach for TD3. The full details are also provided in Algorithm 1.

Our proposed DS-TD3 consists of four main phases: sparse topology initialization, adaptation schedule, topological adaptation, and maintaining sparsity levels.

Sparse Topology Initialization (Algorithm 1 L1-L4). TD3 uses two critic networks (Qθ1 ,Qθ2 ) and one actor network (πϕ) parameterized by θ1 = {θl1 }|Ll=1 , θ2 = {θl2 }|Ll=1 , and ϕ = {ϕl }|Ll=1 respectively; where L is the number of layers in a network. We initialize each of the actor and critic networks with a sparse topology. Sparse connections are allocated in each layer between the hidden neurons at layer l - 1 and layer l. We represent the locations of the sparse connections by a binary mask M = {Ml }|Ll=1 . Following [Mocanu et al., 2018], we use Erd˝os-R´enyi random graph [Erdos et al., 1960] to initialize a sparse topology in each layer l. Namely, the probability of a connection j in layer l is given by:

where λl is a hyperparameter to control the sparsity level in layer l, and nl-1 and nl are the neurons count in layers l - 1 and l, respectively. Mj ∈ {0, 1}; a value of 1 means the existence of a weight in location j. We omit the index l from the mask and weight matrices for readability. A sparse topology is created in each layer for the actor and critic networks:

where ⊙ is an element-wise multiplication operator and Mϕ, Mθ1 , and Mθ2 are binary masks to represent the sparse weights in the actor and two critic networks, respectively. The initial sparsity level is kept fixed during the training.

The target policy and target critic networks are parameterized by ϕ′ , θ′1 , and θ′2 , respectively. Initially, the target networks have the same sparse topology and the same weights as the current networks: ϕ′ ← ϕ, θ′1 ← θ1, θ′2 ← θ2.

After the topological and weight initialization, the agent collects enough data before training using a purely exploratory policy. During training, for each time step, TD3 updates the pair of critics towards the minimum target value of actions selected by the target policy πϕ′ :

where γ is the discounting factor, r is the current reward, s′ is the next state, and ϵ ∼ clip(N(0, ˜σ), -c, c) is the proposed clipped noise by TD3, defined by ˜σ, to increase the stability; where c is the clipped value. As discussed, TD3 proposed to delay the update of the policy network to first minimize the error in the value network before introducing a policy update. Therefore, the actor network is updated every d steps with respect to Qθ1 as shown in Algorithm 1 L17-L19.

During the weight optimization of the actor and critic networks, the values of the existing sparse connections are only updated (i.e., the sparsity level is kept fixed). The sparse topologies of the networks are also optimized during training according to our proposed adaptation schedule.

Adaptation Schedule. The typical practice in DST methods applied in the supervised setting is to perform the dynamic adaptation of the sparse topology after each training epoch. However, this would not fit the RL setting directly due to its dynamic learning nature. In particular, an RL agent faces instability during training due to the lack of a true target objective. The agent learns through trial and error cycles, collecting the data online while interacting with the environment. Adapting the topology very frequently in this learning paradigm would limit the exploration of effective topologies for the data distribution and give a biased estimate of the

Algorithm 1 DS-TD3 (λl , η, e, N, τ , d)

• 1: Initialize critic networks Qθ1 , Qθ2 and actor network πϕ with sparse parameters θ1, θ2, ϕ with a sparsity level defined by λl :
• 2: CreateMϕ, Mθ1 , andMθ2 with Erd˝os-R´enyi graph
• 3: θ1 ←θ1⊙Mθ1 , θ2 ←θ2⊙Mθ2 , ϕ←ϕ⊙Mϕ
• 4: Initialize target networks θ′1 ← θ1, θ′2 ← θ2, ϕ′ ←ϕ
• 5: Initialize replay buffer B
• 6: for t = 1 to T do
• 7: Select action with exploration noise a ∼ πϕ(s) + ϵ,
• 8: ϵ ∼N(0, σ) and observe reward r and new state s′
• 9: Store transition tuple (s, a, r, s′ ) in B
• 10: Sample mini-batch of N transitions fromB
• 11: ˜a ← πϕ′(s′ ) + ϵ, ϵ ∼ clip(N(0, ˜σ),-c, c)
• 12: y ← r +γmini=1,2 Qθ′ i (s′ , ˜a)
• 13: θi ← argminθi 1N ∑(y - Qθi (s, a))2
• 14: if t mod e then
• 15: θi ←TopologicalAdaptation(θi,Mθi , η) (Algo. 2)
• 16: end if
• 17: if t mod d then
• 18: Update ϕ by the deterministic policy gradient:
• 19: ∇ϕJ(ϕ) ← 1N ∑ ∇aQθ1 (s, a)|a=πϕ(s)∇ϕπϕ(s)
• 20: if t mod e then
• 21: ϕ ← TopologicalAdaptation(ϕ,Mϕ, η) (Algo. 2)
• 22: end if
• 23: Update target networks:
• 24: θ′ i ← τθi + (1 - τ )θ′ i
• 25: ϕ′ ← τϕ+ (1 - τ)ϕ′
• 26: θ′i ← MaintainSparsity(θ′ i , ∥θi∥0 ) (Algo. 3)
• 27: ϕ′ ←MaintainSparsity(ϕ′ , ∥ϕ∥0 ) (Algo. 3)
• 28: end if
• 29: end for

current one. To address this point, we propose to delay the adaptation process and perform it every e time steps, where e is a hyperparameter. This would allow the newly added connections from the previous adaptation process to grow. Hence, it would also give better estimates of the connections with the least influence in the performance and an opportunity to explore other effective ones. Analysis of the effect of the adaptation schedule in the success of applying dynamic sparse training in the RL setting is provided in Section 4.2.

Topological Adaptation (Algorithm 2). We adopt the adaptation strategy of the SET method [Mocanu et al., 2018] in our approach. The sparse topologies are optimized according to our adaptation schedule. Every e steps, we update the sparse topology of the actor and critic networks. Here, we explain the adaptation process on the actor network as an example. The same strategy is applied for the critic networks.

The adaptation process is performed through a "drop-andgrow" cycle which consists of two steps. The first step is to drop a fraction η of the least important connections from each layer. This fraction is a subset (cp) of the smallest positive weights and a subset (cn) of the largest negative weights. Thus, the removed weights are the ones closest to zero. Let ˜ϕp and ˜ϕn be the cp-th smallest positive and the cn-th largest

Algorithm 2 Topological Adaptation ( X, M, η)

1: c ← η ∥X∥0 2: cp ← c/2 ; cn ← c/2 3: ˜Xp ←get cp-th smallest positive(X) 4: ˜Xn ←get cn-th largest negative(X) 5: M j ←M j -1[(0 <Xj < ˜Xp ) ∨ (0 >Xj > ˜Xn )] 6: Generate c random integers {x}|c1 7: M j ←M j + 1[(j == x) ∧ (Xj == 0)] 8: X ←X⊙M

Algorithm 3 Maintain Sparsity (X, k)

• 1: ˜X←Sort Descending(|X|)
• 2: M j = 1[|Xj | - ˜Xk ≥ 0], ∀j ∈ {1, ... ∥X∥0 }
• 3: X =X⊙M

negative weights, respectively. The mask Mϕ is updated to represent the dropped connections as follows:

where Mjϕ is the element j in Mϕ, 1 is the indicator function, ∨ is the logical OR operator, and ∥.∥0 is the standard L0 norm. The second step is to grow the same fraction η of removed weights in random locations from the non-existing weights in each layer. Mϕ is updated as follows:

where x is a random integer generated from the discrete uniform distribution in the interval [1, n(l-1) × (nl )] and ∧ is the logical AND operator. The weights of the newly added connections are zero-initialized (ϕ = ϕ ⊙Mϕ).

Maintain Sparsity Level in Target Networks (Algorithm 3). TD3 delays the update of the target networks to be performed every d steps. In addition, the target networks are slowly updated by some proportion τ instead of making the target networks exactly match the current ones (Algorithm 1 L23-L25). These two points lead to a slow deviation of the sparse topologies of the target networks from current networks. Consequently, the slow update of the target networks by τ would slowly increase the number of non-zero connections in the target networks over time. To address this, after each update of the target networks, we prune the extra connections that make the total number of connections exceed the initial defined one. We prune the extra weights based on their smallest magnitude. Assume we have to retain k connections. The target masks of the actor (M′ϕ′ ) and critics (M′θ′1 , M ′θ′2 ) are calculated as follows:

where ˜ ϕ′ k and ˜ θ′ i k is the k-th largest magnitude in the actor and critics respectively, and |(.)j | is the magnitude of element j in the matrix. The target networks are updated as follows:

4 Experiments and Results

In this section, we assess the efficiency of our proposed dynamic sparse training approach for the DRL paradigm and compare it to state-of-the-art algorithms. Experimental settings are provided in the extended version of the paper.

Baselines. We compare our proposed DS-TD3 against the following baselines: (1) TD3 [Fujimoto et al., 2018], the original TD3 where dense networks are used for actor and critic models, (2) Static-TD3, a variant of TD3 where the actor and critic models are initialized with sparse neural networks which are kept fixed during training (i.e., there is no topological optimization), and (3) SAC [Haarnoja et al., 2018b], a popular off-policy algorithm in which the policy is trained to maximize a trade-off between expected return and entropy which results in policies that explore better.

Benchmarks. We performed our experiments on MuJoCo continuous control tasks, interfaced through OpenAI Gym. We evaluate our proposed approach on five challenging environments (HalfCheetah-v3, Hopper-v3, Walker2d-v3, Antv3, and Humanoid-v3).

Metrics. We use multiple metrics to assess the efficiency of the studied DRL methods: (1) Return which is the standard metric used in DRL to measure the performance of an agent, (2) Learning curve area (LCA) which estimates the learning speed of a model (i.e., how quickly a model learns) [Chaudhry et al., 2018] by measuring the area under the training curve of a method, (3) Network size (#params) which measures the memory cost via the number of network parameters, and (4) Floating-point operations (FLOPs) which estimate the computational cost required for training. It is the typical metric in the literature to compare a DST method against its dense counterpart. Details are in the extended version of the paper.

4.1 Results

Learning Behavior and Speed. Figure 1 shows the learning curve of studied methods. DS-TD3 has a much faster learning speed than the baselines, especially at the beginning of the training. After 40-50% of the steps, DS-TD3 can achieve the final performance of TD3. Static-TD3 does not have this favorable property which reveals the importance of optimizing the sparse topology during training to adapt to the incoming data. The learning behavior of DS-TD3 is also faster than SAC in all environments except one. Table 1 shows the learning curve area (LCA) of each method. DSTD3 has higher LCA than TD3 and static-TD3 in all environments. It is also higher than SAC in three environments out of five. This metric is important to differentiate between two agents with similar final performance but very different LCA.

Performance. Table 2 shows the average return (R) over the last 10 evaluations. DS-TD3 outperforms TD3 in all environments. Interestingly, it improves TD3 performance by 2.75%, 20.48%, 64.18%, 16.88%, and 37.51% on HalfCheetah-v3, Walker2d-v3, Hopper-v3, Ant-v3, and Humanoid-v3 respectively. Static-TD3 has a close performance to TD3 in most cases except for Humanoid-v3, where Static-TD3 outperforms TD3 by 30.98%. DS-TD3 has a better final performance than SAC in three environments.

4.2 Analysis

Memory and Computation Costs

We analyze the costs needed for the training process by calculating the FLOPS and #params for the actor and critics. We performed this analysis on Half-Cheetah-v3. #params for dense TD3 is 214784, which requires 1×(1.07e14) FLOPs to train. With our DS-TD3, we can find a much smaller topology that can effectively learn the policy and the function value, achieving higher performance than TD3 with a sparsity level of 51%. This consequently reduces the number of required FLOPs to 0.49×(1.07e14).

Adaptation Schedule

We analyze the effect of the adaptation schedule on the performance. In particular, we ask how frequently the sparse topology should be adapted? We performed this analysis on HalfCheetah-v3. Figure 2a shows the learning curves of DSTD3 using different adaptation schedules controlled by the hyperparameter e (Section 3). Adapting the topology very frequently (e ∈ {200, 500}) would not allow the connections to grow and learn in the dynamic changing nature of RL. The current adaptation process could remove some promising newly added connections from the past adaptation process. This would be caused by a biased estimate of a connection's importance as it becomes a factor of the length of its lifetime. Hence, the very frequent adaptation would increase the chance of replacing some promising topologies. With less frequent adaptation cycles, e = 1000 (the setting from the paper), DS-TD3 can learn faster and eventually achieves higher performance than other baselines. Giving the connections a chance to learn helps in having better estimates of the importance of the connections. Hence, it enables finding more effective topologies by replacing the least effective connections with ones that better fit the data. However, increasing the gap between every two consecutive adaptation processes to 2000 steps decreases the exploration speed of different topologies. As illustrated in the figure, DS-TD3 with e = 2000 has a close learning behavior and final performance to TD3. Yet, it still offers a substantial reduction in memory and computation costs. This analysis reveals the importance of the adaptation schedule in the success of introducing DST to the DRL field.

Sparsity Level

We analyze the performance of our proposed method using different sparsity levels. Figure 2b shows the learning curves of the dense TD3 and DS-TD3. By removing 25% of the connections and training the sparse topology dynamically using DS-TD3, we can achieve a faster learning speed and a performance increase of 2.11%. More interestingly, with a higher reduction in the size of the networks by 50%, we achieve a much faster learning speed. However, when the network has a very high sparsity level (i.e., 80%), it fails to learn effective representations for the reinforcement learning setting. Learning DRL agents using very high sparse networks is still an open-challenging task.

Learning Behavior and Speed

DRL agents learn through trial-and-error due to the lack of true labels. An agent starts training with samples generated from a purely exploratory policy, and new samples are drawn from the learning policy over time. Our results show that dynamic sparse agents have faster adaptability to the newly improved samples, thanks to the generalization ability of sparse neural networks [Hoefler et al., 2021]. This leads to higher learning speed, especially at the beginning of the training. We hypothesize that dense neural networks, being over-parameterized, are more prone to memorize and overfit the inaccurate samples. A longer time is required to adapt to the newly added samples by the improved policy and forget the old ones.

To validate this hypothesis, we analyze the behavior of a dense TD3 agent when it starts training with samples generated from a learned policy instead of a purely exploratory one. We use two learned policies trained for 5 × 105 and 7 × 105 steps to draw the initial samples (see the extended version of the paper). We performed this experiment on HalfCheetahv3. As illustrated in Figure 3, the learning speed of DS-TD3 and TD3 becomes close to each other at the beginning. Afterward, DS-TD3 performs better than TD3 since the new samples are generated from the current learning policies. With initial samples drawn from more improved policy (Figure 3b), dense TD3 learns faster. It achieves higher performance than the dense baseline that starts learning with samples drawn from the policy trained for 5 × 105 steps (Figure 3a). On the other hand, DS-TD3 is more robust to over-fitting, less affected by the initial samples, and quickly adapt to the improved ones over time.

5 Conclusion

Introducing dynamic sparse training principles to the deep reinforcement learning field provides an efficient training process for DRL agents. Our dynamic sparse agents achieve higher performance than the state-of-the-art methods while reducing the memory and computation costs by 50%. Optimizing the sparse topology during training to adapt to the incoming data increases the learning speed. Our findings show the potential of dynamic sparse training in advancing the DRL field. This would open the path to efficient DRL agents that could be trained and deployed on low-resource devices where memory and computation are strictly constrained.