## 3 Problem Formulation
Therefore, selecting an optimal number of epochs for each device is a complex but crucial problem in FL on heterogeneous industrial IoT devices. An optimal number of epochs for each device should be selected to ensure that the model neither underfits nor overfits the data and also reduces the consumption of energy and computation resources. The objective is to minimize the expected loss over all devices while accounting for the communication and computation costs,  
\[ \begin{split}\min_{w,E_{i}}&=c_{i}\cdot\exp_{i} \Big{[}L(f_{w}^{(E_{i})},x)\Big{]}(1-acc(w,E_{i}))\\ &+\alpha\cdot TC(E_{i})+\beta\cdot R(E_{i})-\gamma\cdot D_{non- IID}(E_{i}),\end{split}  \tag{3}\]  
where \(E_{i}\) is the number of epochs selected for device \(i\), \(acc(w,E_{i})\) is the accuracy, \(L(f_{w}^{(E_{i})},x)\) is the loss function of the model on dataset \(x\) after \(E_{i}\) epochs, and \(TC(E_{i})\) is the total cost, which includes the sum of the computation and communication cost, \(R(E_{i})\) quantifies the reduction in communication rounds, and \(D_{non-IID}(E_{i})\) represents the variability due to the distribution of non-IID data between devices. \(\alpha\), \(\beta\), and \(\gamma\) are introduced to control the trade-offs between model quality, total cost, communication rounds, and data distribution, respectively. The relevant total cost \(TC(E_{i})\) can be defined as  
\[TC(E_{i})=C(E_{i})+P_{i}. \tag{4}\]  
The communication cost can be defined as  
\[C(E_{1},E_{2},\dots,E_{n})=\sum_{t=1}^{N}bits\Big{(}w_{(t+1,i)}\Big{)}+bits(ID _{i}), \tag{5}\]  
where \(bits(w_{(t+1,i)})\) is the bits required to transmit the model parameters \(w_{(t+1,i)}\) from device \(i\) to the central server, and \(bits(ID_{i})\) is the number of bits required to transmit a unique ID from device \(i\) to the central server, which is specific to our proposed method. The computational energy consumption can be defined as  
\[p_{i}=D_{i}\epsilon_{i}E_{i}m_{i}f_{i}^{2}l_{i}, \tag{6}\]  
where \(D_{i}\) is the client's data sample, \(\epsilon_{i}\) is the required central processing unit for handling each data sample, \(E_{i}\) is the number of epochs, \(m_{i}\) is the size of the deep learning model implemented in each client, \(f_{i}\) is the central processing unit frequency, and \(l_{i}\) is the complexity of the learning task. The challenge is to design an adaptive and efficient algorithm that can learn the optimal number of epochs for each device, based on its local characteristics and the current state of the model, while minimizing communication and computation costs. This can be formulated as UA-FedRL, where the agent learns a policy for each participating IoT device that maps the current state (e.g., the model's performance on its local data) to the action (the number of epochs to select). The objective is to achieve the highest predicted cumulative reward, which is a function of the learned model's quality while reducing communication and computing costs. Weight parameter aggregation also poses a great challenge to obtain the optimal performance of the global model in FL. A device with larger local datasets and higher computational resources may take more epochs to converge compared to a device with a smaller dataset and lower computational resources. The optimal number of epochs for each device \(D_{i}\) to train a model depends on numerous factors such as quantity of data, quality of the data, computational resources, size of the model, etc. Therefore, it poses a challenge to aggregate the global model from heterogeneous industrial IoT devices and achieve an accurate global model accuracy. Let \(w_{i}^{(e_{i})}\) be the model parameters of device \(i\) after \(e_{i}\) training epochs, and let \(p_{i}\) be the computational resources available on device \(i\).