## 4 Construction of the UA-FedRL Method
### Uncertainty-Aware Reinforcement Learning-Based Optimal Epoch Selection  
#### 4.1.1 Preliminaries  
This section is dedicated to describing the preliminaries of the uncertainty-aware reinforcement learning (UA-RL) method, which is proposed in this paper.  
We assume that the agent's policy is represented by a neural network that takes the state and action as input and outputs a distribution over the number of epochs. Specifically, let \(a_{i,t}\) be the action (number of epochs) selected by agent \(i\) at time step \(t\), and let \(s_{i,t}\) be the state observed by agent \(i\) at time step \(t\). The agent's policy is given by:  
\[\pi_{\theta\pi}(a_{i,t}\mid s_{i,t})=\text{Softmax}(z_{i,t}), \tag{9}\]  
where \(z_{i,t}\) is the output of a neural network with weights \(\theta\), which takes \(s_{i,t}\) as input and outputs a distribution over the number of epochs. We assume that the neural network has one hidden layer with hidden units and a Gaussian prior over the weights:  
\[\theta_{\pi}\sim N\Big{(}\mu_{\pi},\text{diag}\Big{(}\sigma_{\pi}^{2}\Big{)} \Big{)}, \tag{10}\]  
where \(\mu_{\pi}\) and \(\sigma_{\pi}^{2}\) are learnable parameters that are updated during training. Bayes by backprop with variational inference can be used to approximate the posterior distribution over the weights given to the data. We assume a Gaussian distribution over the weights with a mean and variance that are parameterized by the neural network as follows,  
\[q_{\theta}(\theta)=\mathcal{N}\Big{(}\mu\theta,\text{diag}\Big{(}\sigma\theta ^{2}\Big{)}\Big{)}. \tag{11}\]  
We aim to learn the parameters \(\theta_{\pi}\) and the weights \(\theta\) of the neural network with variational inference using the following loss equation:  
\[\begin{split} L=-\frac{1}{N}\sum_{i=1}^{N}\sum_{t=1}^{T_{i}} \log\pi_{\theta\pi}(a_{i},t\mid s_{i},t)\\ \Big{[}\frac{1}{\sigma_{i,t}}&(\sigma_{i,t}-\mu_{i, t})^{2}+\log\sigma_{i,t}\Big{]}\\ +&\beta D_{KL}(q_{\theta}(\theta)\|p(\theta)).\end{split} \tag{12}\]  
The first term corresponds to the expected negative log-likelihood of the actions under the policy and the assumed Gaussian distribution over the number of epochs. The second term is the Kullback-Leibler (KL) divergence between the approximate posterior \(q_{\theta}(\theta)\) and the prior \(p(\theta)\). The hyperparameter \(\beta\) controls the strength of the regularization, which can be defined as follows,  
\[\beta=\min\left(\frac{\text{epoch}}{\frac{\text{num\_epochs}}{4}},1\right). \tag{13}\]  
The upper confidence bound (UCB) algorithm is used to select actions that maximize a combination of the expected reward and the uncertainty of the estimate of the value of the action. Uncertainty is captured by the variance of the approximate posterior distribution over the weights.  
The uncertainty can be computed by sampling \(\theta^{*}\) from the approximate posterior distribution \(q_{\theta}(\theta)\) and computing the estimate of the action value \(Q(s,a,\theta^{*})\) for each action a using the weighted neural network. The relevant equation of variance of the action value estimate is defined as  
\[\begin{split}\text{Var}_{\theta}[Q(s,a,\theta)]&= \text{Var}_{\theta}\Big{[}z^{T}(s,a)\theta\Big{]}\\ &=z^{T}(s,a)\text{Var}_{\theta}[\theta]z(s,a),\end{split} \tag{14}\]where \(z(s,a)\) is the feature vector for the state-action pair \((s,a)\). The UCB \(U(s,a)\) for each action is then calculated as  
\[U(s,a)=Q(s,a,\mu_{\theta})+\beta\sqrt{\text{Var}_{\theta}[Q(s,a,\theta)]}, \tag{15}\]  
where \(\mu_{\theta}\) is the mean of the approximate posterior distribution \(q_{\theta}(\theta)\), and \(\beta\) is a hyperparameter that controls the balance between exploration and exploitation. The action with the highest upper confidence bound is selected as the next action to take.