## 4 Construction of the UA-FedRL Method
### Uncertainty-Aware Reinforcement Learning-Based Optimal Epoch Selection
#### 4.1.2 Proposed Method  
The optimum epoch selection can be represented as an MDP problem using the equation \(M=(S,A,P,R,\gamma)\). UA-RL is then used to explore the action to select the optimal epoch for each IoT device to acquire the best accuracy. To achieve lower latency, SDN and MEC are employed in the core of heterogeneous industrial IoT networks. SDN facilitates the separation of the control plane from the data plane within the communication network. The control plane handles decision-making related to network traffic, while the data plane is responsible for the actual forwarding of this traffic. This study utilized multi-agent reinforcement learning approach, where each client is assigned to a UA-RL agent to select epochs for each device. Figure 3 shows the detailed workflow of the UA-RL method in selecting local epochs for each client considering the status of each client. After monitoring the current state \(S(t)\) of the assigned device at time step \(t\), each agent performs an action \(a(t)\). The central server calculates team rewards \(R(t)\) considering the accuracy, available resources, and total cost of the training model. The aim of the UA-FedRL method is to acquire the maximum reward by exploring an optimal policy. The state space, action space, and reward can be defined as follows,  
1. State space: The state space of the UA-FedRL method at time step \(t\) consists of available computation resources \(G(t)=G_{1}^{t},G_{2}^{t},\dots,G_{i}^{t}\), model state, \(M(t)=M_{1}^{t},M_{2}^{t},\dots,M_{i}^{t}\), and epoch selection state, \(E(t-1)=E_{1}^{(t-1)},E_{2}^{(t-1)},\dots,E_{i}^{(t-1)}\). The state space of UA-FedRL can be defined as \[S=\{(e_{1},e_{2},\dots,e_{N},m,g_{1},g_{2},\dots,g_{N})\mid\] \[e_{i}\in E\text{ for }i=1,2,\dots,N,m\in M,\] (16) \[g_{i}\in G\text{ for }i=1,2,\dots,N\}.\]
2. Action space: The action of the UA-FedRL method is to select epochs for each IoT device at each time step to find the optimal epoch for each device. The action space of the UA-FedRL model can be defined as \[A= \{a_{1},a_{2},\dots,a_{N}\mid\] \[a_{i}\in E,1\leq a_{i}\leq 30\text{ for }i=1,2,\dots,N\},\] (17) where \(E\) is the set of possible epochs for each IoT device, and \(a_{i}\) is the next epoch number for the \(i\)th device. 3. Reward space: The reward function is used by the agent to evaluate the action taken to find the optimal action, i.e., the optimal epoch number. The UA-FedRL model's reward is defined as \[R(s,a,s^{\prime})=\frac{k_{i}}{c_{i}\times TC_{i}},\] (18) where \(k\) is the model accuracy, \(c_{i}\) is the loss value of the training model from the \(i\)th IoT device, and \(TC_{i}\) is the total cost of training the model on the \(i\)th IoT device. The global reward can be defined as \[R_{t}=\frac{1}{N}\sum_{i=1}^{N}R_{i}(s,a,s^{\prime}).\] (19)The detailed procedure of the proposed UA-FedRL to select local training epochs can be seen in Algorithm 1. The algorithm starts by initializing the weights \(\theta\) of the neural network with a Gaussian prior \(\theta_{\pi}\). Additionally, we initialize the hyperparameters \(\alpha\), \(\beta\), \(\gamma\), and \(N\), where \(N\) represents the number of agents. During the training phase, the algorithm iterates over multiple episodes. Within each episode, the algorithm loops through several time steps \(t\). At each time step, it observes the current state \(s_{t}\) and iterates on each agent \(i\). For each agent, we sample the weights \(\theta_{i}\) from the approximate posterior distribution \(q_{\theta_{i}}(\theta_{i})\) and compute the action value estimates \(Q_{i,t}(s_{t},a;\theta_{i})\) for all possible actions a. We then calculate the upper confidence bounds \(U_{i,t}(s_{t},a)\) for all actions and select the action \(a_{i,t}\) that has the highest upper confidence bound. After determining the actions for all agents, we perform the composite action \(a_{t}=[a_{1,t},a_{2,t},\ldots,a_{N,t}]\) and observe the reward \(r_{t}\).