Deep Reinforcement Learning Assisted Federated Learning Algorithm for Data Management of IIoT

The continuous expanded scale of the industrial Internet of Things (IIoT) leads to IIoT equipments generating massive amounts of user data every moment. According to the different requirement of end users, these data usually have high heterogeneity and privacy, while most of users are reluctant to expose them to the public view. How to manage these time series data in an efficient and safe way in the field of IIoT is still an open issue, such that it has attracted extensive attention from academia and industry. As a new machine learning (ML) paradigm, federated learning (FL) has great advantages in training heterogeneous and private data. This paper studies the FL technology applications to manage IIoT equipment data in wireless network environments. In order to increase the model aggregation rate and reduce communication costs, we apply deep reinforcement learning (DRL) to IIoT equipment selection process, specifically to select those IIoT equipment nodes with accurate models. Therefore, we propose a FL algorithm assisted by DRL, which can take into account the privacy and efficiency of data training of IIoT equipment. By analyzing the data characteristics of IIoT equipments, we use MNIST, fashion MNIST and CIFAR-10 data sets to represent the data generated by IIoT. During the experiment, we employ the deep neural network (DNN) model to train the data, and experimental results show that the accuracy can reach more than 97\%, which corroborates the effectiveness of the proposed algorithm.


I. INTRODUCTION
Human society is rapidly moving towards the era of industry 4.0 [1]. The global distribution of user equipments (UEs) is widespread and decentralized due to the influence of geographic location. Limited by the costs of transmission media (cables, optical fibers) and communication delays, traditional network infrastructures are not suitable for the development of Industry 4.0 [2]. On the other hand, wireless networks are widely used in driverless cars, smart grid and smart medical [3]. Fig. 1 shows the IIoT scenario under the background of rapid development of social science and technology. A large number of IIoT equipments access to IIoT frequently, which can produce a huge amount of data in a short period of time [4], [5]. Radio network resource management faces severe challenges, including storage, spectrum, computing resource allocation, and joint allocation of multiple resources [6], [7]. With the rapid development of communication networks, the integrated space-ground network has also become a key research object [8]. How to efficiently manage, store and use these time series data has become an important research topic.
The data generated by IIoT equipments often involves users' private information [9], [10]. Service providers and users do not want this information to be exposed to the third party, but this kind of data is usually vulnerable to attacks from heterogeneous networks, heterogeneous equipments or malware [11]. Therefore, there needs an approach to support IIoT, that cannot only maintain the heterogeneity and privacy of data, but also reduce the communication cost and model training deviation [12], [13]. As a new type of machine learning (ML) paradigm, federated learning (FL) has attracted great attention in the industry [14]. FL prevents the leakage of users' personal private information to a certain extent by separating the central server's direct access to the original data from the model training [15], [16]. In addition, FL can effectively maintain the heterogeneity of data and reduce the deviation of model training. Reference [17] has proved that FL can be effectively applied to multiple learning tasks. Considering the privacy, heterogeneity and wide distribution of the device data in IIoT, FL in IIoT may have a positive effect.
With the rapid popularization of IIoT technology, the geographical distribution of IIoT equipments is becoming more and more extensive and the equipment forms are quite different from each other [18]. This leads to a series of problems such as uneven data quality of equipment and high communication cost. In order to increase the model training rate of IIoT equipments and reduce the communication cost of model aggregation, driven by the current artificial intelligence (AI) technologies, we propose a deep reinforcement learning (DRL) assisted FL framework [19], [20]. DRL algorithm is suitable for solving decision-making problems in high dimensions [21], [22], so it can use the decision-making effect of DRL to select some high-quality local IIoT equipment models for aggregation. Considering that a large amount of data generated by IIoT equipments will lead to excessive training data, increase the burden on wireless network channels and may cause privacy leaks, we employ a distributed method to train  DRL agents, i.e., apply the DRL assisted FL method on each IIoT equipment to realize the effective application of FL in IIoT. The main work of this paper is as follows: 1) Aiming at the problem of difficult management of IIoT devices due to the generation of large amounts of heterogeneous and private data, this paper adopts a FL method to train and manage these data. 2) In order to improve the performance and efficiency of FL, this paper proposes a DRL-assisted FL framework. DRL algorithm based on deep deterministic strategy gradient (DDPG) is mainly used to select IIoT device nodes with high data quality, so as to increase the model aggregation rate and reduce communication costs. 3) We analyze the characteristics of IIoT device data, and then use MNIST, Fashion MNIST and CIFAR10 (IID and non-IID) data sets for experimental evaluation. Finally, we verify the effectiveness of FL technology based on DRL to process IIoT data.
The remaining part of this paper is organized as follows: Section II reviews the work related to the application of AI technology in IIoT. Section III introduces in detail the related issues of the application of FL in IIoT. Section IV introduces the realization of FL algorithm assisted by DRL and the node selection process of applying DRL algorithm. Section V describes the experimental setup of FL algorithm assisted by DRL and shows the experimental results. The last section concludes the paper.

A. Related Work of Industrial Internet of Things Based on Deep Reinforcement Learning
IIoT has been widely concerned by the industry. Especially in the field of IIoT, there have been many solutions to solve the practical problems of IIoT by using AI technology. Due to the explosive growth of user equipments and data streams, IIoT has seen a shortage of spectrum resources. Shi et al. [23] proposed a spectrum resource management scheme for IIoT networks. Specifically, in order to consider the differentiated communication needs of different users, the authors proposed a modified deep Q-learning network (MDQN) and designed a new reward function to drive the learning process. Afterwards, the authors established a simple medium access control model, which used the base station as the sole agent to manage spectrum resources. In the end, the solution promoted the sharing of spectrum between different types of users. Chen et al. [24] studied the joint power control and dynamic resource management of multi-access edge computing (MEC) in IIoT. The authors transformed this problem into Markov decision process (MDP) and used the dynamic resource management algorithm based on DRL to solve this process. The algorithm fully considered the dynamic and continuity of task generation, and finally used the DDPG to optimize the long-term average delay of dynamic resource management. The experimental results showed that the method was effective. Based on blockchain technology, Liu et al. [25] proposed an emerging data collection and sharing scheme to deal with the two main problems faced by IIoT. The first was to achieve efficient data collection within the limited energy and sensing range of mobile terminals. The second was to ensure the security of data sharing between mobile terminals. The authors combined blockchain and DRL algorithms, used DRL to maximize data collection, and then used blockchain technology to ensure data security. This scheme embodied the superior performance of the two technologies.

B. Related Work of Industrial Internet of Things Based on Federal Learning
The latest development of FL in IIoT was an efficient deep anomaly detection framework based on equipment joint learning proposed by Liu et al. [26], which was used to detect time series data in IIoT. The authors first used a FL framework to enable edge equipments to co-train an anomaly model. After that, they used the attention mechanism-based convolutional neural network (CNN) long short-term memory (LSTM) model to capture important fine-grained features. In the end, the authors used a gradient compression mechanism based on Top-k selection to improve communication efficiency. A large number of experiments on real data sets showed that the algorithm can effectively reduce communication overhead. FL was originally proposed by Brendan et al. [27]. They first proposed a decentralized method for private data training-FL, and then conducted extensive experimental evaluations using five different model structures and four data sets to verify the robustness of the FL method. This research opened a precedent for the exploration of FL technology. In addition, there were some typical representatives who applied FL to solve different actual network problems [28]- [32], which had shown good results in solving the computing, caching and communication, and IIoT problems of the intelligent mobile edge.
We summarize the application classification of FL and DRL in TABLE I. Similar to the above works, we have jointly paid attention to the data privacy issues in the IIoT scene, and proposed feasible solutions from the perspective of design framework and algorithm implementation. However, the biggest difference between our work and the above works is that we use DRL algorithm to select high-quality IIoT equipment nodes for FL, so as to improve the rate of model aggregation and reduce the communication cost. This is not reflected in the above work.

III. FEDERATED LEARNING APPLICATION PROBLEM DESCRIPTION IN INDUSTRIAL INTERNET OF THINGS
The core goal of FL is to carry out efficient ML among multi-UEs or multiple computing nodes under the premise of ensuring the security and privacy of data communication [33]. IIoT equipments need to upload the local model to a central server after using local data for model training, and the central server will optimize the global model. The general steps are: 1) IIoT equipments use local data for local calculation, while minimizing the predefined empirical risk function, and then update the calculated weight to the wireless network access point. 2) The wireless network access point collects the weight of IIoT equipment update and accesses the FL unit to generate the global model.
3) FL redistributes the output of model training to IIoT equipments, which use global models for a new round of local training. 4) Repeat the above steps 2 and 3 until the loss function converges or reaches the maximum number of iterations. Suppose that there are N IIoT equipment nodes in an IIoT scenario, the local data set composed of local data of each IIoT-DN is {X 1 , X 2 , ..., X N }. The IIoT equipment node downloads the global model θ from the central server and trains them by local data set alignment. IIoT equipment node uploads new weights or gradients to a central server to update the global model. Therefore, the data sample size from N IIoT equipment node is N n=1 x n = X. Then the loss function of IIoT equipment node with data set X N is where f j (θ) is the loss function of data sample j. Optimize the global loss function by minimizing the weighted average of the local loss function F n (θ) of each IIoT equipment node training sample: We summarize the loss functions of several common ML models that can be used for FL in TABLE II [34]. DRL agent training real data from IIoT equipments has more advantages than agent data provided by the data center. The data generated by IIoT equipments is usually highly secretive and large in scale. DRL agents also need to record them on a central server when they are trained based on local models. In order to better verify the effectiveness of FL assisted by DRL, we extract the data characteristics generated by IIoT equipments as follows: (1) high privacy, (2) unbalanced amount of data (a lot of user output is used, the user output is low), (3) the distribution scale is large, and (4) the user equipment is limited by the communication quality. Corresponding to the actual situation of the data generated by IIoT equipments, our experiments are carried out on MNIST, Fashion MNIST and CIFAR-10 (independent and identically distributed (IID) and non-independent and identically distributed (non-IID)). The reason for using the above data sets is that they can reflect the characteristics of data heterogeneity, scale difference and dispersion to a certain extent.
In the case of non-IID, the data between users can be divided equally or unequal. We assume that there is a fixed number of IIoT devices in each round of communication between the server and devices. When the communication process starts, the devices are randomly divided into several groups, and the server will send the current global model parameters to each device. In order to improve the communication quality, we use DRL algorithm based on DDPG to select some device nodes to participate in the training. After that, each selected device will calculate according to the global state and local data set, and then send the update to the server. The server uses this repeated method to update the global model parameters. In where vn is the weights connecting the neurons, N is the number of neurons and φ() is the activation function.
this way, the non-IID characteristics of IIoT device data can be considered. The experimental setup part will be introduced in detail later.

IV. IMPLEMENTATION OF FEDERATED LEARNING ALGORITHM ASSISTED BY DRL
In this part, we will give a FL framework assisted by DRL in section IV-A, which is the core technology to realize the data training of IIoT equipment. Based on this framework, we will describe in detail the implementation steps of the FL algorithm assisted by DRL in section IV-B. Finally, we will introduce the process of applying DRL algorithm to achieve IIoT equipment selection.

A. Framework
The purpose of integrating DRL into FL is to intelligently use the cooperation between IIoT equipments and nodes to exchange learning parameters, so as to better train the local model [35]. Due to the limited cache and performance calculation of the edge IIoT equipments and nodes, and the direct transmission of a large amount of data may cause network channel congestion or even data leakage, and so the data is trained in the local equipment. Therefore, we make full use of the efficient model training performance of DRL and conduct data confidentiality training based on FL. The final purpose is to verify the training effectiveness of equipment data in IIoT. We deploy DRL agents in every edge IIoT equipment. Considering the inherent characteristics of wireless network, if the amount of data to be trained is large, it will increase the burden on the wireless network channel. Moreover, the data needs to be transformed for privacy reasons, so the correlation of the central server-side agent data may not be as good as the data on the edge equipments. In addition, deploying DRL agents on edge equipments alone may cause additional energy consumption.
To deal with the above problems, we propose a FL framework assisted by DRL, as shown in Fig. 2. Considering a large number of data factors, we adopt the distributed deep learning (DL) method. As shown in TABLE I, distributed DL has good performance of fast training and edge efficiency. Therefore, the fast training of the data on edge equipments can avoid network congestion to certain extents. Considering that the correlation of agent data on the central server is not as good as that of the edge equipment, in the IIoT scenario, a large number of IIoT equipments are equipped with perceptron, which can obtain rich and personalized data to update the global model. Based on these data, the central server-side agent data can be effectively updated to match the data of the edge equipment side. The FL framework based on DRL mainly completes three tasks: node selection, local training and global aggregation. Node selection is achieved by using the DRL algorithm based on DDPG, in order to select IIoT nodes with high data quality to participate in the model training process. The purpose of local training is to derive local model parameters suitable for local nodes. The purpose of global aggregation is to generate a global model suitable for data training, which is achieved by uploading local model parameters to each local server.

B. Implementation Steps
The FL framework assisted by DRL has three main implementation stages: 1) Initialization stage: The central server evaluates the connection request of IIoT equipments. After that, the central server randomly selects a subset of user equipment from the connected IIoT equipments to participate in this round of training [36], [37]. After training, a global model θ t is sent to each selected IIoT equipment. 2) Training stage: The selected IIoT equipment uses local data to train the global model, and the training process is θ t → θ n t , get the global model θ n+1 t after each iteration. For the n-th IIoT equipment, the optimization objective of loss function is where F (θ) is the n-th local loss function. for the next iteration. After that, the central server issues the new global model to the newly selected IIoT equipment collection. Repeat the above process until the loss function converges or reaches the maximum number of iterations.
In the FL algorithm assisted by DRL that we designed, we select IIoT equipments with a proportion of C in each round and then calculate the gradient loss of these IIoT equipment data. For IIoT equipment n, calculate the gradient g n = ∇F n (θ t ) under the current model θ t and then aggregate these gradients by the central server. Use the following formula to update: The operation of selecting specific IIoT devices is called a scheduling strategy. Since the reading and execution of data in the central server and the reading and execution of data in various IIoT devices are not in the same order of magnitude, uploading the training and updated models of all IIoT devices to the global controller will cause a lot of computing time and communication overhead. Therefore, it is an effective way to select a specific part of equipment upload parameter model through the scheduling strategy. Scheduling strategies play a crucial role in allocating wireless channels with limited resources to appropriate IIoT devices. Use S = C N to represent the ratio of the number of IIoT devices to the number of subchannels. There are mainly three commonly used scheduling strategies: 1) Random scheduling strategy: In each round of communication, the central server randomly selects N related IIoT devices for parameter update. At the same time, a dedicated sub-channel is built between the central server and the selected user equipment to transmit training parameters.
2) Cyclic scheduling strategy: The central server divides all IIoT devices into G groups. Each time a group is selected to build a channel for communication, and the model parameters are updated cyclically during each communication. 3) Proportional fairness strategy: In each round of communication, select N from C related IIoT devices according to the following calculation method: where i = {i 1 , i 2 , ..., i N } is a vector of length N . i * = {i * i , i * 2 , ..., i * N } represents the index number of IIoT device. ρ ii,t and ρ ii,t represent the instantaneous and time average signal-to-noise ratios (SNRs) of IIoT device i n during the t-th communication, respectively.
The SNR received at the central server is expressed as follows: where β represents the link loss index, h k represents the smallscale fading, σ 2 is the variance of Gaussian additive noise, and Ψ k u represents the position of IIoT device k.
An important purpose of model training is parameter update. At the number of communication rounds t, when the central server establishes a transmission channel with a certain IIoT equipment and successfully decodes the sent data, the parameter δv t k can be updated from IIoT equipment to the central server. We use the probability of successful parameter update to indicate the transmission performance of the wireless channel, which is expressed as follows: where γ k,t represents the instantaneous and time average signal-to-noise ratio at the central server. S b k,t ∈ {0, 1} represents the selection index, where b represents different scheduling strategies. When S b k,t = 1, it means that the communication time has started between the central server and IIoT equipment k, otherwise there is no communication.

C. IIoT Equipment Node Selection Based on DRL
Due to the different uses and geographical distribution of IIoT equipments, the data they generate will have the characteristics of heterogeneity, and the quality of the data will also vary. The problems of low data training efficiency, long model aggregation time and high wireless communication cost caused by data characteristics are major challenges in the application of FL in the field of IIoT. Therefore, in order to improve the quality of model aggregation and reduce communication costs, inspired by reference [20], we use DRL algorithm to select IIoT equipment nodes with accurate learning models. The main idea is to model the local training cost of IIoT equipments and describe the problem as a MDP. We use DDPG to find the optimal solution of MDP. DRL agent continuously interacts with the environment to accumulate maximum reward and the essential purpose is to minimize the cost of FL.
We first give the cost index for selecting IIoT equipment nodes. The total cost of node selection in FL algorithm is composed of training time cost and training quality cost. The training time cost of local IIoT equipment is where In addition, we use the indicator of learning quality to characterize the loss of training accuracy: where σ i = j Loss(y j −ŷ j ) is used to quantify the quality of the network model, w t represents the training model aggregated in time slot t, and d i represents the training data of IIoT equipment i. Therefore, the total cost of DRL assisted FL algorithm in time slot t is The node selection of IIoT equipment is a combinatorial optimization problem. We model it as a MDP, denoted by M = {S, A, P A , C A }. Among them, S represents the state space of IIoT equipment node, A represents the action space, P A represents the state change probability of the action taken, and C A represents the cost of the new state produced by the action. Therefore, the selection problem of IIoT equipment nodes is expressed as: where δ t represents the selection status of the IIoT equipment node. When the IIoT equipment node i is selected, δ t i = 1, otherwise δ t i = 0. The DRL agent trains the local model by interacting with the environment, and then we use the DDPG to select the optimal solution for IIoT equipment node. In the DRL algorithm, we use the reward function to evaluate the effect of taking action a t . The evaluation method is as follows: where T m represents the number of CPU cycles required to train the model m on the data d i , µ i represents the computing resources available to IIoT equipment i, and τ represents the transmission rate of the wireless network channel. The DDPG uses a value function to determine the strategy, which mainly includes an actor deep neural network (DNN) π(s t |θ pi ) and a critic DNN Q(s t , a t |θ Q ). The DDPG uses historical experience stored in local replay memory to perform information conversion. The converted information includes the current state s t , the action taken a t , the next state s t+1 and the reward r(s t , a t ) obtained by taking the action a t . The target network generates target values to train the critic DNN model.
The specific pseudo code is given in Algorithm 1, where the client label is indexed by n, B represents the local mini batch size, E represents the number of local epoch, α represents the learning rate.

V. EXPERIMENTAL SETUP AND RESULT ANALYSIS
In this part, we evaluate the performance of FL algorithm assisted by DRL based on MNIST, Fashion MNIST and CIFAR-10 (IID and non-IID). Since the above data sets reflect the characteristics of IIoT equipment data, we mainly test the accuracy of the algorithm in training the above data sets. This shows that the FL algorithm assisted by DRL has advantages in training IIoT equipment data.

A. Experimental Setup
Considering the host performance, we mainly use several small agent data sets to test the algorithm performance. MNIST data set and Fashion MNIST data set are mainly used for number recognition tasks. The former uses a simple multilayer perception (MLP) with two hidden layers, in which each layer has 200 units and is activated by the ReLU function [38]. The latter is trained by CNN, which has a fully connected layer (512 units and ReLU function), a softmax output layer and two 5 × 5 convolution layers. To fully study FL, we use two data partitioning methods for IIoT equipments. The first is the IID of data shuffling and then divided into 100 equipment Algorithm 1 Federated learning algorithm assisted by deep reinforcement learning Input: θ π , θ Q ; 1: Set θ tar π = θ π and θ tar Q = θ Q ; 2: Initialize actor DN N and critic DN N parameters; 3: K clients, B batch size, α learning rate; 4: Server executs: 5: initialize θ 0 ; 6: for each round do 7: max(C · K, 1) → num; 8: (random clients(num)) → S t ; 9: select client n ∈ S t ; 10: ClientU pdate(n, θ t ) → θ n t+1 ;

25:
Sample a mini batch of experiences 26: f rom local replay memory; 27: update π(s|θ π ) and Q(s, a|θ Q ); 28: update the target network parameters 29: to local replay memory; 30: store the new experiences; 31: end for 32: end for 33: return θ to server; receiving 600 data samples respectively. The second is the non-IID which classifies the data according to the digital tags, which is divided into 200 pieces of 300 pieces and each client is allocated 2 pieces. This setting method can explore the breaking degree of the algorithm on highly non-IID data. CIFAR-10 data consists of 10 classes of 32 × 32 images with three GRB channels. We used 50,000 training cases and 10,000 test cases and distributed them equally to 100 equipment.

B. Performance Evaluation
We first explore the impact of the number of IIoT equipment, i.e., the number of clients, on the MNIST data set and fashion MNIST data set models. The client ratio C controls the number of multi client parallelism. Because we use 100 IIoT device as clients, when C = 0, it represents one client, followed by 10, 20, 50 and 100 clients in turn.  clients account for to reach the required precision for training the two data sets. When B = ∞ is used for data training, all 600 data samples are treated as a single batch processing in each round. At this time, increasing the proportion of clients can only show a small advantage. When B = 10 is used for batch processing (especially when C ≥ 0.1), the advantage of communication round number change is obvious. In order to balance the computational efficiency and the convergence rate, we fix C = 0.1. Figs. 3 and 4 show the gradual improvement of the accuracy of MNIST data set and Fashion MNIST data set after different training rounds. Each data set is trained based on IID data and non-IID data, respectively. It can be seen that after about 1,000 rounds of communication, the accuracy rate of MNIST data set can reach 99% and then it tends to be stable. After about 2,000 rounds of communication, the training accuracy of the Fashion MNIST data set reached 97% and stabilized.  Since the CIFAR-10 data set has no natural user partition, we consider the balancing and IID settings. The architecture of the training model is derived from Tensorflow, which mainly includes two convolution layers, two full connection layers and one linear conversion layer. In the process of training, we use stochastic gradient descent (SGD) method for small batch training with size of 100. Fig. 6 shows the relationship between the test accuracy of CIFAR-10 data set and the number of communication rounds. Through the above experiments, we can see the advantages of the FL algorithm assisted by DRL in training the data of Our algorithm 20 10 17 173 IIoT equipment. Considering the outstanding characteristics of heterogeneity and privacy of data generated by IIoT equipments, we use MNIST, Fashion MNIST and CIFAR-10 data sets to represent them. Based on MLP and CNN to train the above data sets, the training accuracy is rising in the process of continuous communication between the client and the central server. Among them, CIFAR-10 data set can achieve 85% test accuracy, while MNIST and Fashion MNIST data sets can achieve more than 98% training accuracy. Therefore, it can be shown that the FL algorithm assisted by DRL has excellent performance and has the ability to effectively manage and train the equipment data of IIoT.
In addition, we compare the algorithm proposed in this paper with a baseline algorithm called FedSGD [39]. In each round of communication, we select IIoT devices with a proportion of C, and then calculate the loss gradient of these devices. We fix the value of C to 0.1, and then discuss the number of communication rounds required for the algorithm to achieve the target accuracy when B or E changes. The results are shown in TABLE IV.
It can be seen from the table that the change of E value and B value has obvious influence on the number of communication rounds. We set the target accuracy to 98%. For IID data, increasing the value of E and reducing the value of B can significantly reduce the number of communication rounds to reach this accuracy. But for non-IID data, although the number of communication rounds will be reduced, the overall effect is not obvious. In general, the larger E is and the smaller B is, the optimization effect is more obvious. So overall, the algorithm proposed in this paper is effective in improving the data training rate. Each selected IIoT device calculates the average gradient of local data based on the current model, and then the central server aggregates these gradients and updates them using equation (4). Each IIoT device uses its local data to perform one-step gradient descent based on the current model, and then the server performs a weighted average of the generated model. The final algorithm achieved good results.

VI. CONCLUSION
IIoT is the key technology to realize industry 4.0 and it is also the objective embodiment of the development degree of industry 4.0. What cannot be ignored is that the current IIoT is facing the management and training problems brought by the explosive growth of user data. The emergence of FL provides a new solution paradigm for heterogeneous data and private data training, and it can support the development of IIoT in a way that reduces model deviation. This paper mainly aims at the problem of how to manage and train a large amount of data produced by IIoT, and proposes a FL algorithm assisted by DRL in wireless network environments. DRL based on DDPG is mainly used for the selection of IIoT equipment nodes. We fully analyze the heterogeneity and privacy of the data generated by IIoT equipments. In the experimental phase, we use MNIST, Fashion MNIST and CIFAR-10 data sets to represent IIoT equipment data. The final results show that the FL algorithm assisted by DRL can effectively train the above data sets and achieve a high accuracy rate, which shows the effectiveness of the FL algorithm assisted by DRL in the management of IIoT equipment data.