# Safe Deep Model-Based Reinforcement Learning with Lyapunov Functions

Harry Zhang
Course Project for 16.332

###### Abstract

Model-based Reinforcement Learning (MBRL) has shown many desirable properties
for intelligent control tasks. However, satisfying safety and stability
constraints during training and rollout remains an open question. We propose a
new Model-based RL framework to enable efficient policy learning with unknown
dynamics based on learning model predictive control (LMPC) framework with
mathematically provable guarantees of stability. We introduce and explore a
novel method for adding safety constraints for model-based RL during training
and policy learning. The new stability-augmented framework consists of a
neural-network-based learner that learns to construct a Lyapunov function, and
a model-based RL agent to consistently complete the tasks while satisfying
user-specified constraints given only sub-optimal demonstrations and sparse-
cost feedback. We demonstrate the capability of the proposed framework through
simulated experiments. Code and data are available here.

Summary: Reinforcement learning (RL) for control systems with unknown
nonlinear dynamics and safety constraints has been difficult due to unintended
behavior from imperfect cost function and dynamic uncertainty, which could
lead to instability breaking down the system. We propose a new RL framework to
enable efficient policy learning with unknown dynamics based on the recently
proposed learning model predictive control framework [1, 2, 3, 4, 5], with
provable guarantees of stability. The new stability-augmented framework
consists of a neural-network-based learner that learns a Lyapunov function,
and a model-based RL agent to consistently complete the tasks while satisfying
user-specified constraints given only sub-optimal demonstrations and sparse-
cost feedback.

Our work is based on Safety Augmented Value Estimation from Demonstrations
(SAVED) [6, 7, 8, 9, 10]. SAVED is a sample-efficient method that can
consistently complete tasks with complex safety constraints using only non-
expert, sub-optimal demonstrations and sparse rewards from simple supervision
identifying task completion. By adding safety constraints during the training
of policy learning, SAVED ensures a high probability of being able to plan
back to the safe region. However, it has no stability guarantees as the
learned value function is an arbitrary neural network that is not constrained
to be a Lyapunov function. To address the stability issues which could cause
states to blow up and catastrophic consequences in practice, we use the
learned Lyapunov network for the model-based RL agent as the value function,
which will act in the environment in tandem with the safety-augmented agent in
order to achieve both stability and safety.

To construct and learn a Lyapunov-constrained value function approximator, we
train a Bayesian Neural Network with a loss function adapted from Mittal et
al. [11], Zhang et al. [12], Eisner et al. [13], Avigal et al. [14, 15]. The
output from the network is defined as $V_{Net}(x)$, a $n_{V}\times n_{x}$
matrix, and the actual Lyapunov function used as the value function can be
derived as $V(x)=x^{T}(l_{l}I+V_{net}(x)^{T}V_{net}(x))x$, where $x$ is the
state and $l_{l}$ is a tunable parameter.

We train a 3-layer neural network with 500 hidden units per layers and swish
activations, as in [6, 16, 17, 18] and evaluate the performance of the SAVED
algorithm with and without Lyapunov-constrained value functions in four
simulated environments and point-mass navigation tasks. The effectiveness of
our method and its effects on the stability during training as well as the
optimality of the learned policy are shown and summarized in Figure 2. From
the results, we found that the Lyapunov value function significantly improved
the controller performance by training a value function that was monotonically
decreasing along the planning trajectories, which prevented the agent from
being stuck in local minima around the initial states and resulted in more
optimal trajectories with less variance across runs as well as higher task
completion and constraint satisfaction rates.

Our contribution with this study is twofold: first, the implementation of
learning a Lyapunov function as the value function of SAVED; second, the
empirical validation of the theoretical performance improvement in terms of
stability, task completion rates, and constraint satisfaction rates. Some
future work of this study includes the physical experiments on a robot and the
study of stronger stability guarantees such as asymptotic stability.

##  1 Introduction

Modern control problems require a proper understanding of the system. However,
the dynamic model of some tasks could be very complicated with thousands of
parameters which makes it impossible to deploy system identification.
Furthermore, some tasks require additional safety constraints such as the
collision avoidance in navigation and manipulation, velocity and position
limitation in human robot interaction, which makes the problem harder. In
addition, a poorly designed controller could lead to unstable behavior
resulting in catastrophic consequences.

Deep model-based reinforcement learning (MBRL) has the advantage of using a
deep neural network to express the dynamic model with sample efficiency.
However, the performance of deep model-based RL usually depends on the design
of the cost function which could lead to unintended behavior [6, 19, 20, 21].
Prior work has been done to relax the requirement of the cost function design
by using the task completion as the supervision with some sub-optimal
demonstrations [22, 6]. To ensure the agent staying in the safety region,
safety constraints are added during the training the RL agent.

One concern of deep model-based RL is the stability of the closed-loop system
with the learned policy. Since the dynamic model is not perfectly modeled and
the outputs from the learned policy are not explicitly constrained, there is
no stability guarantee. A common verification of stability of the dynamic
system is the existence of the Lyapunov function, which implies the
convergence to the equilibrium point of the resulting trajectory. A Lyapunov
function for a system with unknown dynamics can be estimated using a neural
network [23]. [24, 25, 26] have shown that the learned Lyapunov neural network
can provide stability guarantees and improve the performance.

In this study, we focus on policy learning for unknown systems with safety
constraints and stability guarantees. We use a learned Lyapunov function as
the value function for Safety Augmented Value Estimation from Demonstrations
(SAVED) [6], a learning method which is able to complete the tasks efficiently
with safety constraints using simple task completion supervision and sub-
optimal demonstrations. The Lyapunov function is learned using a neural
network with a loss function proposed in [11]. Our experiments using the
Safety Augmented Lyapunov Value Estimation from Demonstrations (SALVED) show
the performance improvement brought by the Lyapunov value function, which
helps to increase the stability of the MPC controller and lower the
probability of the controller getting stuck in sub-optimal trajectories by
forcing the trained value function to be monotonically decreasing along the
closed-loop trajectories and reducing the number of local minima in the value
function. The results from SALVED also show higher task completion and
constraint satisfaction rates and more optimal trajectories with less
variance.

The contributions of this study consist of two parts: (1) the implementation
of the Lyapunov value function in SAVED and (2) the experiments validating the
benefits and stability of the new framework SALVED.

##  2 Related Work

Significant interest in model-based planning and data-driven model-based RL
[27, 28, 29, 30, 31] has been sparked in research due to the improvements in
sample efficiency when planning over learned dynamics compared to model-free
methods for continuous control [32, 33]. Moreover, model predictive control
(MPC) with learned model through data-driven approach has been shown to be
promising as it has been applied in various tasks such as deformable objects
manipulation [34], and car racing [1].

The safety of a control system implies the ability of the system to stay in
the safe region and converge to the region of attraction. In reinforcement
learning, the safety of the closed loop system with agent is a challenging
problem due to unknown system dynamics. Safe RL has risen broad interests in
the past few decades [35]. [22, 6] focuses on the problem of ensuring the
states lie in safety set with high probabilities for system with unknown
dynamics. [22, 6] solve the problem using sparse cost and probabilistic
constraints during exploration. [36, 37] study the problem of exploration with
explicit constraints satisfied.

The study on the stability of closed loop system with RL agents mostly adapts
the idea of using Lyapunov stability criterion. To our knowledge, [38] first
used Lyapunov function in reinforcement learning to guarantee the closed loop
system stability. The construction of the Lyapunov function for an unknown
system is challenging. [23] used a neural network to estimate the Lyapunov
function. [39] proposed a Lyapunov based approach for constrained Markov
decision process (CMDP) by constructing the Lyapunov function using an LP-
based algorithm. [11] proposed a method using a learned Lyapunov Neural
Network as the value function estimation to be the terminal cost of the Model
predictive control (MPC) and showed an increased stable horizon compared to
the baseline MPC. [26] learned the safety set from the Lyapunov function. The
Lyapunov network proposed in [26] is used in [11]. In this paper, we use the
same Lyapunov network from [26] to learn a Lyapunov function as the value
function.

##  3 Preliminaries

In this section, we introduce some key concepts that this paper builds upon.

###  3.1 Control Lyapunov Functions

A Lyapunov function $V(x)$ is a positive scalar function that takes as input a
state and outputs a positive scalar value:

$V(x):\;\mathbb{R}^{n}_{x}\to\mathbb{R}^{+}$

, and it should satisfy the following constraint that it monotonically
decreases along the closed-loop trajectory of a controlled system until a
target point (or set) is reached. The existence of such a function is a
necessary and sufficient condition for stability and convergence of a
dynamical system [40]. Control Lyapunov functions can be used to design
control policies.

A control Lyapunov function should be upper and lower bounded by strictly
increasing, unbounded, positive functions $\mathcal{K}_{\infty}$ [40].
Consider the optimal control problem with cost-to-go function defined as:

$c(x,u)=x(t)^{T}Qx(t)+u(t)^{T}Ru(t)$

, where $x$ is the state and $u$ is the input. The $\mathcal{K}_{\infty}$ can
be designed as the scaled, squared $l_{2}$ norm of the state:

$l_{l}||x(t)||_{2}^{2}\leq V(x(t))\leq L_{V}||x(t)||_{2}^{2}$

where $l_{l}$ is the minimum eigenvalue of $Q$ and $L_{V}$ is a local
Lipschitz constant for $V$.

Furthermore, we want the Lyapunov function values to decrease monotonically
along the closed-loop system’s trajectory. The amount of the decrease is
usually chosen to be quantified by the stage cost (cost-to-go) value incurred
at the current time step. Mathematically, this condition can be written as:

$V(x(t+1))-V(x(t))\leq-c(x(t),u(t)).$

###  3.2 Model Predictive Control

Model predictive control (MPC) is a control strategy that iteratively plans
ahead and solves for a sequence of actions to take that incurs the minimal
cost-to-go values in the plan horizon. The agent will only take the first
action in the planned sequence in order to minimize the deviation caused by
model inaccuracy. The iterative optimization problem that MPC tries to solve
is described as follows:

$\displaystyle J^{*}(x(t))$ $\displaystyle=\min_{u_{t}\dots
u_{t+N-1}}\sum_{k=0}^{N-1}c(x_{t+k},u_{t+k})$ $\displaystyle\text{s.t.
}x_{t+k+1}$ $\displaystyle=f(x_{t+k},u_{t+k})$

, where $N$ is the planning horizon of the MPC and we only take the first
element in the planned the sequence of inputs $u_{t}\dots u_{t+N-1}$. After
taking this step, we result in the next state and re-plan according to the
above optimization problem.

##  4 Methods

The main novelty of our work is that we are trying to design a Lyapunov-
constrained value function which the model-based RL agent can leverage in
order to achieve stable performance. Specifically, we plan to improve the
SAVED algorithm, which can perform poorly if the learned value function
contains many local minima as shown in Figure 1, using a Lyapunov neural
network that can guide the agent to correctly plan the action to take at each
time step.

Figure 1: Planning trajectories of the MPC policy in SAVED in Pointbot 1
(point mass navigation) environment, where the value of the heat map is the
predicted cost-to-go (negative returns) according to the value function at
each iteration. In the plots, the $x$-axis is [-130, -80] and the $y$-axis is
[-10, 10]. The starting position $x_{0}$ is $(-100,0)$ and the goal is at the
origin, $(0,0)$. We see that even in an environment with no obstacles or
constraints, SAVED struggles to complete the task due to local minima.

###  4.1 SAVED Algorithm

At each time step, SAVED solves the optimization problem in Equation 1, where
$C(x_{t+i},u_{t+i})$ is the sparse cost function indicating task completion,
$V_{\phi}^{\pi}(x_{i+H})$ is the value function and $f_{\theta}$ is the
learned dynamic model. $\rho_{\alpha}$ is the safety density model
constraining the exploration by requiring $x_{t+H}$ to be in the region with
high task completion probability. And $\beta$ is the probability of
$x_{t:t+H}$ falling in the feasible region $\mathcal{X}^{H+1}$, where
$\mathcal{X}^{H+1}$ is the set of $H+1$ length state sequences [23].

$\displaystyle
u^{*}_{t:t+H-1}=\mathop{\operatorname{argmin}}_{u_{t:t+H-1}\in\mathcal{U}^{H}}\mathbb{E}_{x_{t:t+H}}[\sum_{i=0}^{H-1}C(x_{t+i},u_{t+i})+V_{\phi}^{\pi}(x_{i+H})],$
(1) $\displaystyle\text{s.t.}\quad x_{t+i+1}\sim
f_{\theta}(x_{t+i,u+i})\forall i\in\\{0,\dots,H-1\\},$
$\displaystyle\quad\quad\rho_{\alpha}(x_{t+H})>\delta,\mathbb{P}(x_{t:t+H}\in\mathcal{X}^{H+1})\geq\beta$

###  4.2 Lyapunov Value Function

We plan to constrain the value function used by the model-based RL to be a
control Lyapunov function. To achieve this, we use the function introduced in
[26]:

$s_{v}V(x)=s_{v}\cdot
x^{T}\left(l_{l}I+V_{\text{net}}(x)^{T}V_{\text{net}}\right)x,$ (2)

where $V_{\text{net}}$ is a neural network that produces a $n_{V}\times n_{x}$
matrix. Here $n_{V}$, $s_{v}$, and $l_{l}>0$ are hyper-parameters. In the
original introduction of the value function in [26], $s_{v}$ is fixed to be 1,
but in our experiments, we have found that with $s_{v}=1$, $V(x)$ tends to be
numerically unstable, causing the loss to blow up, so we scale down the value
function output significantly (e.g. $s_{v}=0.001$).

We use the same loss function as described in [11], which enforces the value
function to be control Lyapunov. The objective of the agent is to minimize the
following loss in order to train the Lyapunov neural network:

$\min_{V_{\text{net}},l_{s}}\mathbb{E}_{(x,u,x^{\prime})}\left[J(x,u,x^{\prime})\right]$

where $J(x,u,x^{\prime})$ is defined as

$J(x,u,x^{\prime})=\frac{\mathcal{I}(x)}{\rho}J_{s}(x,u,x^{\prime})+J_{\text{vol}}(x,u,x^{\prime})$

and

$\displaystyle\mathcal{I}(x)$ $\displaystyle=0.5(\text{sign}[l_{s}-V(x)]+1)$
$\displaystyle J_{s}(x,u,x^{\prime})$ $\displaystyle=\frac{\text{ReLU}[\Delta
V(x)]}{V(x)+\epsilon}$ $\displaystyle J_{\text{vol}}(x,u,x^{\prime})$
$\displaystyle=\text{sign}[\Delta V(x)](l_{s}-V(x))$ $\displaystyle\Delta
V(x)$ $\displaystyle=V(x^{\prime})-\lambda V(x)+v\mathcal{I}(x)$

.

In the definitions above, $\mathcal{I}(x)$ is an indicator function
representing the step cost that the agent incurs at a given state $x$.
$J_{\text{vol}}$ is a classification loss that computes the correct boundary
between the stable and unstable points, which also helps to increase the safe
set volume [11]. The scalars $\epsilon>0$, $\lambda\in[0,1)$, $v\in[0,1]$, and
$\rho>0$ are hyper-parameters. $l_{s}$ is a trainable variable, and the role
of $l_{s}$ is to provide an upper bound for $V(s)$ such that the states can
remain in an invariant set.

###  4.3 Safety-Augmented Lyapunov Value Estimation from Demonstration
(SALVED)

By substituting the Lyapunov value function $V$ calculated from Equation 2 in
Equation 1, we have SALVED to be SAVED with a Lyapunov value function, which
provides an estimate for the terminal cost for the current policy at the end
of the planning horizon:

$u^{*}_{t:t+H-1}=\displaystyle\mathop{\operatorname{argmin}}_{u_{t:t+H-1}\in\mathcal{U}^{H}}\mathbb{E}_{x_{t:t+H}}\left[\sum_{i=0}^{H-1}C(x_{t+i},u_{t+i})+s_{v}V(x_{i+H})\right],\\\
$ (3)

According to Lemma 1 from [11], for a given horizon length $N$ and contraction
factor $\lambda$, there is minimum scaling of Lyapunov function $V$, which is
calculated from Equation 2, and a lower bound on the discount factor $s_{v}$
such that the system under MPC is stable.

When the learned dynamic model $f_{\theta}$ is perfect, the state would
converge to the equilibrium point. In practice, however, the learned dynamics
model may be imperfect for various reasons, such as due to aleatoric
uncertainty in the actual dynamics and the epistemic uncertainty as a result
of insufficient information about the environment, but the system can still
stable when error between the learned dynamic model and the actual model is
bounded, in which case the state would converge to a bounded area around the
equilibrium point, whose size depends on the model error, $N$ and $s_{v}$
[11].

##  5 Experiments

###  5.1 Experimental Setup

For all of our SALVED experiments, we used Safety-Augmented Value Estimation
from Demonstrations (SAVED) as our baseline. In order to test whether the new
Lyapunov-based value functions can provide additional stability to model-based
RL methods that are based on the learning MPC framework, such as SAVED, we
benchmarked our proposed method, SALVED, and compared its performance with
that of the baseline on four 4-dimensional simulated navigation tasks in which
a point-mass navigates from a random starting position to a goal set, which is
defined as a unit ball centered around the origin:

  * •

Pointbot 1: long navigation task with no obstacles

  * •

Pointbot 2: medium-length navigation task with a single large obstacle along
the $y$-axis

  * •

Pointbot 3: medium-length navigation task with two medium-sized obstacles and
an open channel along the $x$-axis

  * •

Pointbot 4: long navigation task with obstacles of varying sizes surrounding
the goal and a small opening for entry

###  5.2 Implementation Details

The dynamics models for the baseline and SALVE have the same network
architecture, which is a probabilistic ensemble of five 3-layer neural
networks with 500 hidden units per layer and swish activations. For SAVED, the
value function has the same network architecture as its dynamics model,
whereas for SALVED, the value function is a neural network with the same
number of layers, same number of hidden units per layer, and the same
activations but constrained to be a Lyapunov function.

###  5.3 Results

(a) Pointbot 1

(b) Pointbot 2

(c) Pointbot 3

(d) Pointbot 4

Figure 2: Pointbot Navigation Tasks: SALVED is evaluated on four Pointbot
navigation tasks. Pointbot 1 has no obstacles or safety constraints, Pointbot
2-4 contain obstacles with increasingly more complex constraints. We observe
that SALVED has a lower iteration cost than baselines throughout all four
tasks and shows considerably more stability during training and evaluation.

(a) Pointbot 1 (SAVED)

(b) Pointbot 1 (SALVED)

(c) Pointbot 2 (SAVED)

(d) Pointbot 2 (SALVED)

(a) Pointbot 3 (SAVED)

(b) Pointbot 3 (SALVED)

(c) Pointbot 4 (SAVED)

(d) Pointbot 4 (SALVED)

Figure 4: The plots show evaluation trajectories combined across 100
iterations. The left plots show that the control policies learned by SAVED are
considerably less optimal and have lower task completion rates across the
experiments. The right plots demonstrate higher task completion rates and
suggest that the policies learned by SALVED are significantly more optimal and
consistent as the resulting trajectories are also more dense.

##  6 Conclusion

In this paper, we present SALVED, a model-based reinforcement learning
algorithm with provable stability guarantees by learning Lyapunov neural
network value functions. We empirically evaluated SALVED on four simulated
benchmarks and found that SALVED provided significant performance improvements
over the baseline in all four environments. SALVED is able to maintain the
sample efficiency of model-based RL algorithms such as SAVED and the iterative
improvement guarantees of the Learning MPC framework, while providing
additional Lyapunov stability guarantees. Results not only show higher task
completion and constraint satisfaction rates, but also better-performing
learned control policies with more optimal trajectories with less variance
across runs and controllers that are less prone to local minima and
catastrophic failures from constraint violations due to the stability
guarantees from the Lyapunov-constrained value function approximation.

## References

  * Rosolia and Borrelli [2017] U. Rosolia and F. Borrelli, “Learning model predictive control for iterative tasks. a data-driven control framework,” _IEEE Transactions on Automatic Control_ , vol. 63, no. 7, pp. 1883–1896, 2017.
  * Rosolia et al. [2017] U. Rosolia, X. Zhang, and F. Borrelli, “Robust learning model predictive control for iterative tasks: Learning from experience,” in _2017 IEEE 56th Annual Conference on Decision and Control (CDC)_. IEEE, 2017, pp. 1157–1162.
  * Rosolia and Borrelli [2019] U. Rosolia and F. Borrelli, “Sample-based learning model predictive control for linear uncertain systems,” in _2019 IEEE 58th Conference on Decision and Control (CDC)_. IEEE, 2019, pp. 2702–2707.
  * Pan et al. [2022] C. Pan, B. Okorn, H. Zhang, B. Eisner, and D. Held, “Tax-pose: Task-specific cross-pose estimation for robot manipulation,” _arXiv preprint arXiv:2211.09325_ , 2022.
  * Zhang [2016] H. Zhang, “Health diagnosis based on analysis of data captured by wearable technology devices,” _International Journal of Advanced Science and Technology_ , vol. 95, pp. 89–96, 2016.
  * Thananjeyan et al. [2020a] B. Thananjeyan, A. Balakrishna, U. Rosolia, F. Li, R. McAllister, J. E. Gonzalez, S. Levine, F. Borrelli, and K. Goldberg, “Safety augmented value estimation from demonstrations (SAVED): Safe deep model-based RL for sparse cost robotic tasks,” _IEEE Robotics and Automation Letters_ , vol. 5, no. 2, pp. 3612–3619, 2020.
  * Sim et al. [2019] K. C. Sim, F. Beaufays, A. Benard, D. Guliani, A. Kabel, N. Khare, T. Lucassen, P. Zadrazil, H. Zhang, L. Johnson _et al._ , “Personalization of end-to-end speech recognition on mobile devices for named entities,” in _2019 IEEE Automatic Speech Recognition and Understanding Workshop (ASRU)_. IEEE, 2019, pp. 23–30.
  * Elmquist et al. [2022] A. Elmquist, A. Young, T. Hansen, S. Ashokkumar, S. Caldararu, A. Dashora, I. Mahajan, H. Zhang, L. Fang, H. Shen _et al._ , “Art/atk: A research platform for assessing and mitigating the sim-to-real gap in robotics and autonomous vehicle engineering,” _arXiv preprint arXiv:2211.04886_ , 2022.
  * Zhang et al. [2020] H. Zhang, J. Ichnowski, Y. Avigal, J. Gonzales, I. Stoica, and K. Goldberg, “Dex-net ar: Distributed deep grasp planning using a commodity cellphone and augmented reality app,” in _2020 IEEE International Conference on Robotics and Automation (ICRA)_. IEEE, 2020, pp. 552–558.
  * Zhang et al. [2021] H. Zhang, J. Ichnowski, D. Seita, J. Wang, H. Huang, and K. Goldberg, “Robots of the lost arc: Self-supervised learning to dynamically manipulate fixed-endpoint cables,” in _2021 IEEE International Conference on Robotics and Automation (ICRA)_. IEEE, 2021, pp. 4560–4567.
  * Mittal et al. [2020] M. Mittal, M. Gallieri, A. Quaglino, S. S. M. Salehian, and J. Koutník, “Neural lyapunov model predictive control,” _arXiv preprint arXiv:2002.10451_ , 2020.
  * Zhang et al. [2023] H. Zhang, B. Eisner, and D. Held, “Flowbot++: Learning generalized articulated objects manipulation via articulation projection,” _arXiv preprint arXiv:2306.12893_ , 2023.
  * Eisner et al. [2022] B. Eisner, H. Zhang, and D. Held, “Flowbot3d: Learning 3d articulation flow to manipulate articulated objects,” _arXiv preprint arXiv:2205.04382_ , 2022.
  * Avigal et al. [2020] Y. Avigal, S. Paradis, and H. Zhang, “6-dof grasp planning using fast 3d reconstruction and grasp quality cnn,” _arXiv preprint arXiv:2009.08618_ , 2020.
  * Avigal et al. [2021] Y. Avigal, V. Satish, Z. Tam, H. Huang, H. Zhang, M. Danielczuk, J. Ichnowski, and K. Goldberg, “Avplug: Approach vector planning for unicontact grasping amid clutter,” in _2021 IEEE 17th International Conference on Automation Science and Engineering (CASE)_. IEEE, 2021, pp. 1140–1147.
  * Pan et al. [2023] C. Pan, B. Okorn, H. Zhang, B. Eisner, and D. Held, “Tax-pose: Task-specific cross-pose estimation for robot manipulation,” in _Conference on Robot Learning_. PMLR, 2023, pp. 1783–1792.
  * Lim et al. [2021] V. Lim, H. Huang, L. Y. Chen, J. Wang, J. Ichnowski, D. Seita, M. Laskey, and K. Goldberg, “Planar robot casting with real2sim2real self-supervised learning,” _arXiv preprint arXiv:2111.04814_ , 2021.
  * Lim et al. [2022] ——, “Real2sim2real: Self-supervised learning of physical single-step dynamic actions for planar robot casting,” in _2022 International Conference on Robotics and Automation (ICRA)_. IEEE, 2022, pp. 8282–8289.
  * Devgon et al. [2020] S. Devgon, J. Ichnowski, A. Balakrishna, H. Zhang, and K. Goldberg, “Orienting novel 3d objects using self-supervised learning of rotation transforms,” in _2020 IEEE 16th International Conference on Automation Science and Engineering (CASE)_. IEEE, 2020, pp. 1453–1460.
  * Yao et al. [2023] Y. Yao, S. Deng, Z. Cao, H. Zhang, and L.-J. Deng, “Apla: Additional perturbation for latent noise with adversarial training enables consistency,” _arXiv preprint arXiv:2308.12605_ , 2023.
  * Shen et al. [2024] S. Shen, Z. Zhu, L. Fan, H. Zhang, and X. Wu, “Diffclip: Leveraging stable diffusion for language grounded 3d classification,” in _Proceedings of the IEEE/CVF Winter Conference on Applications of Computer Vision_ , 2024, pp. 3596–3605.
  * Thananjeyan et al. [2020b] B. Thananjeyan, A. Balakrishna, U. Rosolia, J. E. Gonzalez, A. Ames, and K. Goldberg, “ABC-LMPC: safe sample-based learning MPC for stochastic nonlinear dynamical systems with adjustable boundary conditions,” _arXiv preprint arXiv:2003.01410_ , 2020.
  * Berkenkamp et al. [2017] F. Berkenkamp, M. Turchetta, A. Schoellig, and A. Krause, “Safe model-based reinforcement learning with stability guarantees,” in _Advances in neural information processing systems_ , 2017, pp. 908–918.
  * Bobiti and Lazar [2016] R. Bobiti and M. Lazar, “Sampling-based verification of lyapunov’s inequality for piecewise continuous nonlinear systems,” _arXiv preprint arXiv:1609.00302_ , 2016.
  * Bobiti [2017] R. Bobiti, “Sampling–driven stability domains computation and predictive control of constrained nonlinear systems,” Ph.D. dissertation, Department of Electrical Engineering, Oct. 2017, proefschrift.
  * Gallieri et al. [2019] M. Gallieri, S. S. M. Salehian, N. E. Toklu, A. Quaglino, J. Masci, J. Koutník, and F. Gomez, “Safe interactive model-based learning,” _arXiv preprint arXiv:1911.06556_ , 2019.
  * Chua et al. [2018] K. Chua, R. Calandra, R. McAllister, and S. Levine, “Deep reinforcement learning in a handful of trials using probabilistic dynamics models,” in _Advances in Neural Information Processing Systems_ , 2018, pp. 4754–4765.
  * Deisenroth and Rasmussen [2011] M. Deisenroth and C. E. Rasmussen, “PILCO: A model-based and data-efficient approach to policy search,” in _Proceedings of the 28th International Conference on Machine Learning, ICML_ , 2011, pp. 465–472.
  * Fu et al. [2016] J. Fu, S. Levine, and P. Abbeel, “One-shot learning of manipulation skills with online dynamics adaptation and neural network priors,” in _2016 IEEE/RSJ International Conference on Intelligent Robots and Systems, IROS_. IEEE, 2016, pp. 4019–4026.
  * Lowrey et al. [2018] K. Lowrey, A. Rajeswaran, S. Kakade, E. Todorov, and I. Mordatch, “Plan online, learn offline: Efficient learning and exploration via model-based control,” _arXiv preprint arXiv:1811.01848_ , 2018.
  * Nagabandi et al. [2018] A. Nagabandi, G. Kahn, R. S. Fearing, and S. Levine, “Neural network dynamics for model-based deep reinforcement learning with model-free fine-tuning,” in _2018 IEEE International Conference on Robotics and Automation (ICRA)_. IEEE, 2018, pp. 7559–7566.
  * Fujimoto et al. [2018] S. Fujimoto, H. Van Hoof, and D. Meger, “Addressing function approximation error in actor-critic methods,” _arXiv preprint arXiv:1802.09477_ , 2018.
  * Haarnoja et al. [2018] T. Haarnoja, A. Zhou, P. Abbeel, and S. Levine, “Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor,” _arXiv preprint arXiv:1801.01290_ , 2018.
  * Yan et al. [2020] W. Yan, A. Vangipuram, P. Abbeel, and L. Pinto, “Learning predictive representations for deformable objects using contrastive estimation,” _arXiv preprint arXiv:2003.05436_ , 2020.
  * Garcıa and Fernández [2015] J. Garcıa and F. Fernández, “A comprehensive survey on safe reinforcement learning,” _Journal of Machine Learning Research_ , vol. 16, no. 1, pp. 1437–1480, 2015.
  * Achiam et al. [2017] J. Achiam, D. Held, A. Tamar, and P. Abbeel, “Constrained policy optimization,” _arXiv preprint arXiv:1705.10528_ , 2017.
  * Li et al. [2018] Z. Li, U. Kalabić, and T. Chu, “Safe reinforcement learning: Learning with supervision using a constraint-admissible set,” in _2018 Annual American Control Conference (ACC)_. IEEE, 2018, pp. 6390–6395.
  * Perkins and Barto [2002] T. J. Perkins and A. G. Barto, “Lyapunov design for safe reinforcement learning,” _Journal of Machine Learning Research_ , vol. 3, no. Dec, pp. 803–832, 2002.
  * Chow et al. [2018] Y. Chow, O. Nachum, E. Duenez-Guzman, and M. Ghavamzadeh, “A lyapunov-based approach to safe reinforcement learning,” in _Advances in Neural Information Processing Systems_ , 2018, pp. 8092–8101.
  * Khalil [2014] H. K. Khalil, _Nonlinear Control_. Pearson Higher Ed, 2014.


# WeiAvg: Federated Learning Model Aggregation Promoting Data Diversity

Fan Dong1, Ali Abbasi1, Steve Drew1, Henry Leung1, Xin Wang2, Jiayu Zhou3
1Department of Electrical and Software Engineering, University of Calgary,
Calgary, AB, Canada
2Department of Geomatics Engineering, University of Calgary, Calgary, AB,
Canada
3Department of Computer Science and Engineering, Michigan State University,
East Lansing, MI, USA
{fan.dong, ali.abbasi1, steve.drew, leungh<EMAIL_ADDRESS><EMAIL_ADDRESS>
###### Abstract

Federated learning provides a promising privacy-preserving way for utilizing
large-scale private edge data from massive Internet-of-Things (IoT) devices.
While existing research extensively studied optimizing the learning process,
computing efficiency, and communication overhead, one important and often
overlooked aspect is that participants contribute predictive knowledge from
their data, impacting the quality of the federated models learned. While
FedAvg treats each client equally and assigns weight solely based on the
number of samples, the diversity of samples on each client could greatly
affect the local update performance and the final aggregated model. In this
paper, we propose a novel approach to address this issue by introducing a
Weighted Averaging (WeiAvg) framework that emphasizes updates from high-
diversity clients and diminishes the influence of those from low-diversity
clients. Specifically, we introduced a projection-based approximation method
to estimate the diversity of client data, instead of the computation of an
entropy. We use the approximation because the locally computed entropy may not
be transmitted due to excess privacy risk. Extensive experimental results show
that WeiAvg converges faster and achieves higher accuracy than the original
FedAvg algorithm and FedProx.

###### Index Terms:

federated learning, weighted averaging, heterogeneity, entropy

##  I Introduction

In the past decade, Artificial Intelligence (AI) has seen a tremendous impact
on our daily life thanks to the substantial amount of data generated by mobile
[1, 2] and edge [3] devices to train predictive models. Traditional training
methods collect and train data on a centralized server, introducing high
communication costs and privacy concerns. There also has been growing
recognition and increasing legislation efforts in various jurisdictions to
protect the privacy of their residents, e.g., GDPR [4] in EU. Federated
learning (FL) [5] has emerged as a promising solution to training models on
distributed data while preserving data security and privacy, where clients
train the model locally and send only the model parameters to the central
server. By keeping data local and transmitting only model updates, FL balances
collaborative learning and preserving the confidentiality of sensitive
information.

Despite its potential, FL still faces the challenges of statistical data
heterogeneity [1, 6, 7]. Theoretical analysis [8] has proved that such
heterogeneity would increase the convergence and jeopardize the model
performance. Several approaches have been proposed to mitigate the effects of
statistical heterogeneity. FedProx [9] introduces a proximal term to the
objective function, which helps prevent over-fitting during model training.
Reinforcement learning was adopted in [10] to counterbalance non-i.i.d data
distribution by proactively selecting a subset of clients in each global
round. FedAdp [11] accelerates the convergence of FL training by weighted
averaging based on cosine similarity with the constraint requiring moving
average values to stabilize the angle between local update and global update,
resulting in the selection of all clients in each global epoch and posing a
restriction to the potential application scenarios.

Moreover, the definition of “data heterogeneity” is not equivalent in
different research. In FedAvg [5], the authors discussed the non-i.i.d case by
assigning each client two classes of samples with almost the same number of
data samples. The sample split method would yield an almost identical
diversity of classes for all clients, which is not representative from the
perspective of data label heterogeneity.

Figure 1: (a) Two clients, A and B, with the same number of samples but
different diversity. Intuitively, we assign more weight to B when aggregating
model updates of the two clients. (b) The projection of a local update
$m_{k}^{i}$ from Client $k$ at the $i$th round onto the global update from
$m^{i}$ to $m^{i+1}$.

In this paper, we propose a projection-based weighted aggregation in FL to
mitigate the impact of data heterogeneity by using the measurement of entropy,
which commonly exists in the real world. We also use an intuitive example to
reveal why this weighted averaging would work. Given a dataset split and
assigned to multiple clients, represented by the colored dots shown in Fig. 1
(a), where each dot represents a data sample. The color of a dot represents
its respective label. Suppose there are two clients A and B, each assigned 9
samples. Despite both clients having the same number of samples, those on
Client B demonstrate significantly higher diversity. Naturally, we tend to
assign more weight to Client B when aggregating their model updates at the
central server. We quantify the diversity with information entropy [12].
However, the requirement of uploading entropy to the central server may
compromise the privacy of clients. As an alternative, we find the strong
correlation between entropy and projection of local updates on the global
update and use the projection as an approximation to perform the weighted
averaging, as shown in Fig. 1 (b). By this approach, we significantly improve
the convergence speed of FL. Furthermore, we explore integrating WeiAvg with
proximity terms to penalize clients with larger data heterogeneity and receive
and enhance the performance of WeiAvg. Our contributions are listed below:

  * •

We link the clients’ data label diversity to the global model performance in
FL. We verify this hypothesis by using the label entropy to represent the
client data diversity. To our knowledge, this is the first work to consider
label diversity for weighted averaging in FL.

  * •

We substitute the entropy information with a projection-based weighted
averaging algorithm to eliminate privacy concerns. The projection-based metric
proves to be highly correlated to label diversity represented by entropy. An
algorithm called WeiAvg is proposed for projection-based weighted averaging.

  * •

We conduct an extensive empirical study to evaluate the performance of WeiAvg.
The numerical results demonstrate the superiority of WeiAvg in various
scenarios.

##  II Related Work

In this section, we introduce two categories of existing studies relevant to
our proposed method: entropy-based contribution evaluation schemes and
different variants of weighted averaging.

###  II-A Entropy-based Contribution Evaluation

First proposed by Google [5], FL offered a promising paradigm to meet the
conflicting demands of privacy protection and a vast volume of data for model
training. The model convergence time, communication cost, and privacy
enhancement [1, 6, 13] remain critical research topics in FL. One key
challenge in FL is the data heterogeneity among different clients. In [8], the
negative impact of heterogeneity on FL has been theoretically analyzed.

FedProx [9] aimed to mitigate the statistical heterogeneity by introducing a
proximal term in its optimization objective to regularize the training
process, preventing the local training from over-fitting. A client selection
method was adopted to counterbalance the bias introduced by non-i.i.d. data
distribution [10]. In [14], an efficient knowledge transmission method was
proposed for heterogeneous FL using a public dataset and weighted ensemble
distillation scheme.

###  II-B Variants of Weighted Averaging

There have also been studies for weighting model aggregation in FL. FedAT [15]
used a straggler-aware, weighted aggregation heuristic to steer and balance
the training for further accuracy improvement. FedAdp [11] adopted the
weighted averaging method to accelerate the convergence based on the angle
between the local gradient and global gradient, which is only suitable for
cross-silo situations.

To achieve better convergence performance in cross-device scenarios, we
proposed weighted averaging based on the entropy of label distribution.
Furthermore, we demonstrate the strong correlation between projection of local
updates on global updates and entropy, making it possible to perform weighted
averaging solely based on the local updates without asking for additional
client information. We compare the performance of our design with FedAvg [5]
and FedProx [9] and their combined approach, proving our algorithm’s
effectiveness.

##  III Problem Formulation

This section introduces our intuition and formally articulates the research
problem.

###  III-A Intuition: The Importance of Data Diversity

The intuition of our proposed method can be sourced to acknowledge the
benefits of diversity. Participants with different cultures and academic
backgrounds will spark innovations and bring novel discussions in a research
or working environment. The importance of diversity was also systematically
summarized in [16]. To our knowledge, diversity has not been widely used in
FL. In FedAvg [5], the local updates are treated equally during model
aggregation if they possess the same amount of samples. However, the quality
of samples located on each client can be different. We can verify our
hypothesis if diversity can be quantified for each client.

###  III-B Quantifying Diversity by Entropy

To measure the class diversity of each client, we turn to the entropy of the
labels. A higher entropy indicates more diverse samples on a client, while a
lower entropy represents a more skewed dataset. Based on this entropy-based
diversity, we can increase the impact of the highly diversified clients by
assigning higher weights to their corresponding updates according to their
entropy.

Let $X$ be a discrete random variable such that $x\in\mathcal{X}$ and the
probability density function

$p(x)=Pr\\{X=x\\},x\in\mathcal{X}$ (1)

The entropy of variable X is then defined by $H(X)$, where

$H(X)=-\sum_{x\in\mathcal{X}}p(x)\log p(x)$ (2)

To verify our hypothesis that clients with high-entropy classes improve the
model performance, we naturally propose an algorithm named WeiAvg, shown in
Algorithm 1. This algorithm defines the total number of clients as $C$. The
initial global model is $m^{0}$. Suppose there are $N$ clients in total. In
each round, $n$ clients are selected for model updates. The training will last
a total of $I$ rounds. A hyperparameter, denoted by $\lambda$, controls the
degree of emphasis given to high-entropy updates. WeiAvg constructs a positive
correlation between the weight assigned to $c_{i}$ and $H(X_{i})$ when $c_{i}$
contributes its local update for aggregation. To simplify the experiment
setting, we distribute the same amount of samples to each client. In Step
(3.2), we perform a small shift of the entropy vector by the LinearTransform
to ensure each value is larger than 0, in case zero-entropy updates are
discarded completely. In Step (3.3), we apply a non-linear transformation to
control the weights assigned to clients with different entropy. When $\lambda$
equals 0, the weight of each client would be the same. And the WeiAvg would
degenerate to FedAvg.

Algorithm 1 WeiAvg By Entropy

0: $m^{0}$, $N$, $n$, $I$, $C$.

for global epoch $i$ in $[0,1,2,3,\cdots,I-1]$ do

(1) Randomly select $n$ clients from the $C$ and store them into
$\boldsymbol{c}=[c_{1},c_{2},\cdots,c_{n}]$.

(2) Send the current global model $m^{i}$ to $\boldsymbol{c}$. Perform local
training and upload the local updates
$\boldsymbol{u^{i}_{\cdot}}=[u^{i}_{1},u^{i}_{2},\cdots,u^{i}_{n}]$ to the
central server.

(3.1) Compute the entropy of each client in $\boldsymbol{c}$ as
$\boldsymbol{e}$.

(3.2) $\boldsymbol{z}$ $\leftarrow$ $LinearTransform(\boldsymbol{e})$.

(3.3) $\boldsymbol{z^{\prime}}$ $\leftarrow$ $\boldsymbol{z}^{\lambda}$.

(3.4) Calculate the weight of each client in $\boldsymbol{c}$:
$\boldsymbol{w}=\boldsymbol{z^{\prime}}/\sum\boldsymbol{z^{\prime}}$.

(4) Conduct the weighted averaging to obtain the weighted global update:
$\boldsymbol{U^{i}}=\sum(\boldsymbol{w}\cdot\boldsymbol{u^{i}_{\cdot}})$.

(5) Update the global model: $m^{i+1}\leftarrow m^{i}+\boldsymbol{U^{i}}$.

end for

return $m^{I}$

###  III-C Using Projection As An Approximation

Despite that weighted averaging based on entropy could improve the performance
compared to FedAvg. Algorithm 1 has a strong assumption that clients must
upload their entropy which indicates their data diversity, besides uploading
model updates. This is not practical in FL settings due to privacy concerns,
as entropy would expose information on data distribution.

During the training process, we found that the projection of local updates on
aggregated global updates is highly correlated with entropy. The calculation
of projection is shown in Fig. 1 (b), where $m^{i}$ represents the flattened
global model parameters of the $i_{th}$ round before sending to clients,
$m_{k}^{i}$ represents the updated flatten model parameters on client $k$, and
$m^{i+1}$ represents the simple averaged global model after clients send their
local model updates to the central server, note that $m^{i+1}$ is the same
with the aggregated model in FedAvg. The projection norm shown in Fig. 1 (b)
could be calculated as:

$p_{k}=\frac{(m_{k}^{i}-m^{i})\cdot(m^{i+1}-m^{i})}{\lVert
m^{i+1}-m^{i}\rVert}$ (3)

Figure 2: Correlation between projection and entropy with FedAvg and WeiAvg on
different weight skew values ($p$) on MNIST, FashionMNIST, and CIFAR10.

We build a correlation heatmap as illustrated in Fig. 2, wherein for each
grid, we concat all projection values in an epoch and average them based on
the client ID. Then we calculated the Pearson correlation coefficient and the
p-value of the hypothesis test with the null hypothesis that the underlying
sample distributions are not correlated and normally distributed. We performed
such calculations on each dataset as shown in each row in Fig. 2. Within each
dataset, we also applied different $\lambda$ in WeiAvg to further test the
correlation between projection and entropy, as shown in each column in Fig. 2.
The number in the parentheses of xticks denotes the value of $\lambda$. The
number in each grid represents the correlation coefficient (above) and the
corresponding p-value (below). The results show a strong positive correlation.
Based on this finding, we are confident to use the projection values as an
approximation of entropy to infer the data diversity and perform weighted
averaging, which will not impose additional requirements of uploading entropy
for the clients.

Algorithm 2 WeiAvg By Projection

0: $m^{0}$, $N$, $n$, $I$, $C$.

for global epoch $i$ in $[0,1,2,3,\cdots,I-1]$ do

(1) Randomly select $n$ clients from the $C$ and store them into
$\boldsymbol{c}=[c_{1},c_{2},\cdots,c_{n}]$.

(2) Send the current global model $m^{i}$ to $\boldsymbol{c}$. Perform local
training and upload the local updates
$\boldsymbol{u^{i}_{\cdot}}=[u^{i}_{1},u^{i}_{2},\cdots,u^{i}_{n}]$ to the
central server.

(3.1) Simple average $\boldsymbol{u^{i}_{\cdot}}$ to obtain the temporal
global update $\boldsymbol{u^{i+1}}$.

(3.2) For each local update of $\boldsymbol{u^{i}_{\cdot}}$, calculate its
projection on the temporal global update, which would yield a projection
vector $\boldsymbol{p}=\\{p_{1},p_{2},\cdots,p_{n}\\}$.

(3.3) $\boldsymbol{z}$ $\leftarrow$ LinearTransform($\boldsymbol{p}$).

(3.4) $\boldsymbol{z^{\prime}}$ $\leftarrow$ $\boldsymbol{z}^{\lambda}$.

(3.5) Calculate the weight of each client in $\boldsymbol{c}$:
$\boldsymbol{w}=\boldsymbol{z^{\prime}}/\sum\boldsymbol{z^{\prime}}$.

(4) Conduct the weighted averaging to obtain the weighted global update:
$\boldsymbol{U^{i}}=\sum(\boldsymbol{w}\cdot\boldsymbol{u^{i}_{\cdot}})$.

(5) Update the global model: $m^{i+1}\leftarrow m^{i}+\boldsymbol{U^{i}}$.

end for

return $m^{I}$

Figure 3: WeiAvg based on entropy. WeiAvg with a wide range of $\lambda$ shows
superiority over vanilla FedAvg. Large $\lambda$ values may introduce a large
advantage initially while causing instability in the later epochs.  Figure 4:
WeiAvg based on projection with similar patterns compared to Fig. 3.

##  IV Experiment Results

Since the degree of difficulty of training models on different datasets
varies, we adopted slightly different settings for each dataset we use. We
distribute $M$ samples across a total of $N$ clients. The entropy distribution
among $N$ clients is controlled by a hyperparameter $p$. In our experiment, we
take $p$ being $0.8$, and the distribution could be seen as Fig. 5, where
$N=100$. We will discuss the effect of $p$ on our algorithm later.

Figure 5: Distribution of user entropy.

We select $n$ clients in each global epoch to perform local training. During
local training, the batch size is set to be $B$, the learning rate is denoted
as $lr$, and the number of local epochs is $E$. We applied the same “sgd”
optimizer with a momentum of 0.9 for each dataset. The hyperparameter
$\lambda$ is consistent across all datasets, set to 1e-4. The configurations
for each dataset are summarized in Table I. Due to the stochastic nature of FL
in the client selection step, we employ 20 different random seeds to obtain an
average performance for comparison in the following experiment. This approach
not only ensures a more consistent evaluation but also enables us to draw
conclusions with statistical confidence.

TABLE I: Hyperparameters for each dataset. | $M$ | $N$ | $n$ | $B$ | $E$
---|---|---|---|---|---
MNIST | 3000 | 100 | 10 | 64 | 10
FashionMNIST | 4000 | 100 | 10 | 32 | 10
CIFAR10 | 4000 | 100 | 10 | 32 | 20

###  IV-A Comparing WeiAvg and FedAvg

As shown in Fig. 3, we performed weighted averaging using entropy with
different $\lambda$ on MNIST, FashionMNIST, and CIFAR10. WeiAvg with a wide
range of $\lambda$ shows superiority over vanilla FedAvg. Large $\lambda$ like
10 may introduce a large advantage initially while causing instability in the
later epochs. Later, we used projection as an approximation of entropy to
perform the same experiment as shown in Fig. 4. The accuracy line shows a
similar pattern with Fig. 3. This proved that our algorithm works and could
achieve better convergence performance over FedAvg.

Figure 6: WeiAvg and FedProx based on entropy. For all datasets, WeiAvg and
FedProx could achieve consistently better performance compared with FedAvg.
While the performance between WeiAvg and FedProx differs on different
datasets. Figure 7: WeiAvg and FedProx based on projection. The performance
pattern is similar to the experiment shown in Fig. 6.

###  IV-B Exploring the Combined Approach: WeiAvg and FedProx

Another effective algorithm trying to mitigate the impact of heterogeneity in
federated learning is FedProx [9]. In this experiment, we are going to take
FedProx as a comparison. After fine-tuning hyperparameter $\mu$ in FedProx,
which controls how much we want to punish those updates far away from the
previous global model, we found 0.05 is the optimal value for $\mu$. Although
FedProx and our WeiAvg algorithm can mitigate heterogeneity’s impact, they are
not mutually exclusive. They can be combined synergistically to leverage the
advantages of each method.

We first conduct the experiment based on entropy. As shown in Fig. 6, for all
of the datasets, WeiAvg and FedProx could achieve consistently better
performance compared with FedAvg. While the performance between WeiAvg and
FedProx differs on different datasets. For the MNIST dataset, WeiAvg is better
than FedProx all way long. For the FashionMNIST dataset, WeiAvg is also better
in most of the beginning parts. For the CIFAR10 dataset, WeiAvg could prevail
over FedProx in the initial epochs. After combining WeiAvg and FedProx
together, the performance becomes significantly better than any of them.

We replaced entropy with projection and conducted the same experiment again.
As shown in Fig. 7, the performance pattern is similar to the experiment shown
in Fig. 6.

###  IV-C Impact of Distribution Changes on WeiAvg

Before conducting the experiments, we must split the dataset into $N$ clients
to simulate federated learning. We use a hyperparameter $p$ to control the
distribution of entropy. When $p$ is 0, the distribution becomes more even.
When increasing $p$, the distribution becomes more skewed. The entropy
distributions with different $p$ are shown in Fig. 8. The experimental results
on FashionMNIST are presented in Fig. 9, demonstrating the superior
performance of our algorithm across various distributions. Similar patterns
were observed in the experiments conducted on MNIST and CIFAR10. Due to
limited space, we omit their display for brevity.

Figure 8: Entropy distribution with different $p$. Figure 9: Impact of
distribution changes on WeiAvg. WeiAvg demonstrates superior performance in
various distributions.

###  IV-D Identifying Limitations: Cases When WeiAvg Falls Short

The WeiAvg algorithm works conditional on the strong correlation of projection
with entropy. When we decrease the local epochs or by any means cause the
local model under-fitted a lot, the correlation of projection with entropy may
decrease, which will lead the WeiAvg based on projection could not achieve
better performance over FedAvg. The experimental results on FashionMNIST, as
depicted in Fig. 10, demonstrate the effect of under-fitting. It can be
observed that reducing local epochs to 1 or 2 weakens the correlation between
projection and entropy. Consequently, in Fig. 11, WeiAvg tends to exhibit
poorer performance than FedAvg. Conversely, when the number of local epochs is
increased to 5 or 10, the correlation becomes stronger, leading to enhanced
convergence performance for WeiAvg.

Figure 10: Correlation between projection and entropy with different local
training epochs (1, 2, 5, and 10). Higher correlations are observed with a
higher number of local epochs. Figure 11: Convergence line with different
local training epochs. WeiAvg shows an improved performance with a higher
number of local epochs.

##  V Conclusion

This paper introduces Weighted Federated Averaging (WeiAvg), a novel approach
that addresses the issue of data heterogeneity and diversity in federated
learning. By assigning weights based on the entropy of label distribution,
WeiAvg emphasizes updates from high-diversity clients and diminishes the
influence of low-diversity clients by utilizing a projection-based
approximation method to estimate client data diversity without compromising
individual client information. This ensures that the weighted averaging based
on diversity can be performed while preserving privacy. Additionally, we
explore the combination of WeiAvg with the FedProx algorithm, which mitigates
the impact of statistical heterogeneity by adding a proximal term to the
objective. The combined approach achieves improved convergence speed and
accuracy, further enhancing the performance of FL.

## References

  * Bonawitz et al. [2019] K. Bonawitz, H. Eichner, W. Grieskamp, D. Huba, A. Ingerman, V. Ivanov, C. Kiddon, J. Konečnỳ, S. Mazzocchi, B. McMahan _et al._ , “Towards federated learning at scale: System design,” _Proceedings of machine learning and systems_ , vol. 1, pp. 374–388, 2019.
  * Hard et al. [2019] A. Hard, C. M. Kiddon, D. R. Ramage, F. S. Beaufays, H. Eichner, K. Rao, R. Mathews, S. Augenstein, and S. Ramaswamy, “Federated learning for mobile keyboard prediction,” 2019.
  * Abreha et al. [2022] H. G. Abreha, M. Hayajneh, and M. A. Serhani, “Federated learning in edge computing: a systematic survey,” _Sensors_ , vol. 22, no. 2, p. 450, 2022\.
  * [4] “General Data Protection Regulation (GDPR).” [Online]. Available: https://gdpr.eu/
  * McMahan et al. [2017] B. McMahan, E. Moore, D. Ramage, S. Hampson, and B. A. y Arcas, “Communication-efficient learning of deep networks from decentralized data,” in _Artificial intelligence and statistics_. PMLR, 2017, pp. 1273–1282.
  * Guerraoui et al. [2018] R. Guerraoui, S. Rouault _et al._ , “The hidden vulnerability of distributed learning in byzantium,” in _International Conference on Machine Learning_. PMLR, 2018, pp. 3521–3530.
  * Wu et al. [2023a] J. Wu, S. Drew, and J. Zhou, “FedLE: Federated Learning Client Selection with Lifespan Extension for Edge IoT Networks,” in _Proceedings of 2023 IEEE International Conference on Communications (ICC’23)_ , 2023.
  * Li et al. [2019] X. Li, K. Huang, W. Yang, S. Wang, and Z. Zhang, “On the convergence of fedavg on non-iid data,” _arXiv preprint arXiv:1907.02189_ , 2019.
  * Li et al. [2020] T. Li, A. K. Sahu, M. Zaheer, M. Sanjabi, A. Talwalkar, and V. Smith, “Federated optimization in heterogeneous networks,” _Proceedings of Machine learning and systems_ , vol. 2, pp. 429–450, 2020.
  * Wang et al. [2020] H. Wang, Z. Kaplan, D. Niu, and B. Li, “Optimizing federated learning on non-iid data with reinforcement learning,” in _IEEE INFOCOM 2020-IEEE Conference on Computer Communications_. IEEE, 2020, pp. 1698–1707.
  * Wu and Wang [2021] H. Wu and P. Wang, “Fast-convergent federated learning with adaptive weighting,” _IEEE Transactions on Cognitive Communications and Networking_ , vol. 7, no. 4, pp. 1078–1088, 2021.
  * Jost [2006] L. Jost, “Entropy and diversity,” _Oikos_ , vol. 113, no. 2, pp. 363–375, 2006.
  * Wu et al. [2023b] J. Wu, S. Drew, F. Dong, Z. Zhu, and J. Zhou, “Topology-aware federated learning in edge computing: A comprehensive survey,” _arXiv preprint arXiv:2302.02573_ , 2023.
  * Deng et al. [2023] Y. Deng, J. Ren, C. Tang, F. Lyu, Y. Liu, and Y. Zhang, “A hierarchical knowledge transfer framework for heterogeneous federated learning,” in _IEEE INFOCOM 2023-IEEE Conference on Computer Communications_ , 2023.
  * Chai et al. [2020] Z. Chai, Y. Chen, L. Zhao, Y. Cheng, and H. Rangwala, “Fedat: A communication-efficient federated learning method with asynchronous tiers under non-iid data,” _ArXivorg_ , 2020.
  * Gong et al. [2019] Z. Gong, P. Zhong, and W. Hu, “Diversity in machine learning,” _IEEE Access_ , vol. 7, pp. 64 323–64 350, 2019.