On the Linear Speedup of Personalized Federated Reinforcement Learning with Shared Representations
Authors: GUOJUN XIONG
Shufan Wang
Daniel Jiang
Jian Li
Venue: ICLR 2025 Poster (2025)
Keywords: personalized federated reinforcement learning
shared representations
stochastic approximation
Primary area: reinforcement learning

Abstract: Federated reinforcement learning (FedRL) enables multiple agents to collaboratively learn a policy without needing to share the local trajectories collected during agent-environment interactions. However, in practice, the environments faced by different agents are often heterogeneous, but since existing FedRL algorithms learn a single policy across all agents, this may lead to poor performance. In this paper, we introduce a personalized FedRL framework (PFedRL) by taking advantage of possibly shared common structure among agents in heterogeneous environments. Specifically, we develop a class of PFedRL algorithms named PFedRL-Rep that learns (1) a shared feature representation collaboratively among all agents, and (2) an agent-specific weight vector personalized to its local environment. We analyze the convergence of PFedTD-Rep, a particular instance of the framework with temporal difference (TD) learning and linear representations. To the best of our knowledge, we are the first to prove a linear convergence speedup with respect to the number of agents in the PFedRL setting. To achieve this, we show that PFedTD-Rep is an example of federated two-timescale stochastic approximation with Markovian noise. Experimental results demonstrate that PFedTD-Rep, along with an extension to the control setting based on deep Q-networks (DQN), not only improve learning in heterogeneous settings, but also provide better generalization to new environments.

Published as a conference paper at ICLR 2025
ON THE LINEAR SPEEDUP OF PERSONALIZED FED- ERATED REINFORCEMENT LEARNING WITH SHARED REPRESENTATIONS
Guojun Xiong'; Shufan Wang", Daniel J Jian Li² Harvard University, 2tony Brook University, Meta
ABSTRACT
Federated reinforcement learning (FedRL) enables multiple agents to collabora- tively learn a policy without needing to share the local trajectories collected dur- ing agent-environment interactions. However, in practice, the environments faced by different agents are often heterogeneous, but since existing FedRL algorithms learn single policy across all agents, this may lead to poor performance. In this paper, we introduce personalized FedRL. framework (PFEDRL) by tak- ing advantage of possibly shared common structure among agents in heteroge- neous environments. Specifically, we develop a class of PFEDRL algorithms named PFEDRL-REP that learns (I) shared feature representation collabora- tively among all agents, and (2) an agent-specific weight vector personalized to its local environment. We analyze the convergence of PFEDTD-REP, a particu- lar instance of the framework with temporal difference (TD) learning and linear representations. To the best of our knowledge, we are the first to prove a linear convergence speedup with respect to the number of agents in the PFEDRL set- ting. To achieve this, we show that PFEDTD-REP is an example of federated two- timescale stochastic approximation with Markovian noise. Experimental results demonstrate that (FEDTD-RER along with an extension to the control setting based on deep Q-networks (DQN), not only improve learning in heterogeneous settings, but also provide better generalization to new environments.
INTRODUCTION
Federated reinforcement learning (FedRL) (Nadiger et al. 2019 Liu et al. 2019 Xu et al. 2021
ZZhag et al.] 2022a Jin al. 2022 Khodadadian al. 2022 Yuan al. 2023 Salgia & Chi 2024 Woo et al. 2024 Zheng et al.
et al.) 2024) has recently emerged as promising framework that blends the distributed nature of federated learning (FL) (McMahan et al. 2017) with reinforcement learning's 's (RL) ability to make sequential decisions over time (Sutton & Barto 2018 In FedRL, multiple agents collaboratively learn single policy without sharing individual trajectories that are collected during evcet-environment interactions, protecting each agent's 's privacy.
One key challenge facing FedRL environment heterogeneity, where the collected trajectories may vary to large extent across agents. To illustrate, consider few existing applications of FL: on- device NLP applications (e.g., next word prediction, sentence completion, web query suggestions, and speech recognition) from Internet companies (Hard et al.) 2018 Yang et al. 20188] Wang al.) 2023b), on-device recommender or ad prediction systems (Maeng et al. 2022 Krichene et al. 2023). and Internet of Things applications like smart healtheare smart thermostats (Nguyen et al. 2021
Imteaj et al. 2022 Zhang et al.) 2022b; Boubouh et al. (023). Note that all of the above (1) exist in settings with environment heterogeneity (heterogeneous users, devices, patients, homes) and (2) could potentially benefit from an RL problem formulation.
As result, if all agents collaboratively learn single policy, which most existing FedRL frameworks do, the learned policy might perform poorly on individual agents. This calls for the design of personalized FedRL (PFEDRL) framework that can provide personalized policies for agents in different environments. Nevertheless, despite the recent advances in FedRL, the design of PFEDRL
"hii work done when G. Xiong was PhD student Stony Brook University.

---

Published as a conference paper at ICLR 2025
Table 1: Comparison of existing FedRL frameworks in terms of noise; environments (Homo: ho- mogeneous, Hetero: heterogeneous); using representation learning (Rep.L) or not; timescale (TS), single or two-TS (two) updates; multiple local updates or not; personalization across agents or not; and with or without linear convergence speedup guarantee.
Method
Noise
Ens.
Rep.I.
TS
Local updates
Personalized
Linear speedup
Fo & FedQ (Khodaddiinn et 2022
Markov
Homo
Single Single
X
FodID (Dal Fabbeo
F wang er ni 20258
Markov
Hetero.
Single
QAvg PAVy
1.d.
Hotem.
Single
FedQ WWoo 6 N013
Markov
Hetero.
Single
/
A3C
Markov
Homo.
Two
K
FedSARSK (Zhang 21ai 2024
Markow
Notern.
Single
K
P.FDRL-RRP
Markov
Hetero.
Two
and its performance analysis remains, to large extent, an open question. Motivated by this, the first. inquiry we aim to answer in this paper is:
Can design a PFEDRL framework for agents in heterogeneous environments that not only collaboratively learns a useful global model without sharing local trajectories, but also learns personalized policy for each agent?
We address this question by viewing the PFEDRL problem in heterogeneous environments N par- allel RL tasks with possibly shared common structure. This is inspired by observations in centralized learning (Bengio al. [2033] LeCun et al.] (2155) and federated or decentralized learning (Collins) iet al., [021] Tziotis et al., 2023 Xiong et al. 2024), where leveraging shared (low-dimensional) representations can improve performance.
theoretical understanding of using shared representa- tions amongst heterogeneous agents has received recent emphasis in the standard supervised FL (or
decentralized learning) setting (Collins al. 2021 Tziotis al. 2023 Xiong et al. 2024).
However, a theoretical analysis of PFEDRL with shared representations is more subtle because each agent PFEDRL collects data by following its own policy (thereby generating a Markovian trajectory) and simultaneously updates its model parameters. This is in stark contrast to the standard FL paradigm, where data is typically collected in an i.i.d. fashion. Our second research question
How do the shared representations affect the convergence of PFEDRL under Markovian noise, and is it possible achieve N-fold linear convergence speedup?
Despite the recent progress in the standard supervised FL setting (Collins et al. 2021 Tziotis al, 2023 Xiong et al. 2024), to the best of our knowledge, this question is still open in the context of learning personalized policies in FedRL under Markovian noise (see Table 1.. Motivated by these open questions, our main contributions are:
PFEDRL framework with shared representations. PFEDRL-REE learns a global shared feature representation collaboratively among agents through the aid of central server, along with agent-specific parameters for personalizing to each agent's 's local environment. The PFEDRL-RE: framework can be paired with a wide range of RL algorithms, including both value-based and policy-based methods with arbitrary feature repre- sentations.
PFEDRL-REP framework. We propose PFEDRL-REP
Linear speedup for TD learning. We then introduce PFEDTD REP, an instantiation of the above PFEDRL-REP framework for TD learning (Sutton & Barto 2018). We analyze its con- vergence linear representation setting, proving the convergence rate of PFEDTD-REP to be O(N-2/3(T + 2) -223,),, where N is the number of agents and 7 is the number of communication rounds. This implies linear convergence speedup for PFEDTD-RED with respect to the number of agents, highly desirable property that allows for massive parallelism in large-scale systems. To our knowledge, this is the first linear speedup result for PFEDRL with shared representations under Markovian noise, providing a theoretical answer to the empirical observations in Mnih et al. (2016) that federated versions of RL algorithms yield faster convergence. To show this result, make of wwo-timescale stochastic approximation theory and address the challenges of Markovian noise through Lyapunov drift approach.
2

---

Published as a conference paper at ICLR 2025
2 PROBLEM FORMULATION
In this section, we first review the standard FedRL framework and then introduce our proposed PFEDRL-REP framework, which incorporates personalization and shared representations. Let N and 7'' be the number of agents and communication rounds, respectively. Denote [N] as the set of integers ..) and || || as the 12-norm. We use boldface denote matrices and vectors.
2.1 PRELIMINARIES: FEDERATED REINFORCEMENT LEARNING
A FedRL system with N agents interacting with N independent heterogeneous environments is modeled as follows. The enviromment of agent i [N] is a Markov decision process (MDP) M = (S. A, R², Pi, PP)Y). 7). where S and A are finite state and action sets, R' is the reward function, P is the transition kernel, and y E (0.1) is the discount factor. Suppose agent i is equipped with a policy n' S + ((A) (a mapping from states probability distributions over A). At each time step k agent i is in state S and takes action dj according (-())))) resulting in reward R' (s), at) The environment then transitions to a new state + according to P (-]sk, a()) This sequence of states and actions forms Markoy chain, the source of the aforementioned Markovian noise. In this paper, this Markov chain is assumed to be unichain, which is known to asymptotically converge to a steady state. We denote the stationary distribution as p
The value of n' in environment M is defined as Vi (s)= realistic problems with large state spaces, it is infeasible to store VA, (s) for all states, so function
= E, [En= (s a()) 16 = s]. In
approximation is often used. One example is Vi.= (s) ~ (s) 0. where $ E RlS)xd xd is state feature representation and E Rd is an unknown low-dimensional weight vector.
One intermediate goal in RL is to estimate the value function corresponding to policy using trajectories collected from the environment. This task is called policy evaluation, and one widely used approach is temporal difference (TD) learning (Sutton 1988). The FedRL version of TD learning is called FedTD (Khodadadian et al.] 2022: Dal Fabbro et al. 2023 Wang et al. 2023a), where N agents collaboratively evaluate single policy и by learning common (non-personalizee) weight vector 0, using trajectories collected from N different environments. More precisely, we have n' and 0 0, vi [N]. Given feature representation (s), Vs, this can be formulated as the following optimization problem:
min
1 N
Σ 1=1
E Exngie 0 (asj0-vir (s]
(I)
Due to space constraints, we focus our presentation on the policy evaluation problem. Note that policy evaluation an important part of RL and control, since is a critical step for methods based on policy improvement. Our proposed PFEDRL framework (see Algorithm TT can be directly
applicd to control problems as well, but we relegate these discussions to Scction and Appendix C
2.2 PERSONALIZED FEDRL WITH SHARED REPRESENTATIONS
Since the local environments are heterogeneous across the N agents, the aforementioned FedRL methods (in Section 2 that aim to learn common weight vector 0 may perform poorly on in- dividual agents. This necessitates the search for personalized local weight vectors 0 that can be learned collaboratively among N agents in N heterogeneous environments (without sharing their locally collected trajectories). As alluded to earlier, we view the personalized FedRL. (PFEDRL) problem N parallel RL tasks with possibly shared common structure, and propose that the agents collaboratively learn a common features representation & addition to a personalized lo- cal weights 0². Specifically, the value function of agent is approximated as Vi f'(Oi Φ), where f' (...) is a general function parameterized by these two unknown parameters [ The policy
'The approximation ("(0", 1100,90 2) is general and can take on various forms, including as linear approximations neural networks. For instance, it can he represented linear combination of & and 0¹ i.c., f' +) ==
Φ0' in TD with linear function approximation (Bhandari et al. (0188). In addition, f' ; b) can represent a deep neural network; see, e.g., our extension of PFFDDRL to control problems and its instantiation with DON (deep Q-networks) (Miih et al. 2015) in Section 5 and Appendix
3

---

Published as a conference paper at ICLR 2025
Algorithm 1 PFEDRL-REP: A General Description
Input: Sampling policy x', vi E [N];
1: Initialize the global feature representation $ and local weight vector jj. Vi [N] randomly;
do 2: for round 0, 1, .....
3: for agent 1, ..... N do
4: 0++1 WEIGHT JUPDATE , 0i at, K);
5: -++1/2 = FEATURE (PTTUTEE.UDDTTE(( 2. 1: B:);
6: end for
7: Server computes the new global feature representation L.+1 Eim1 Φ i/41/2"
8: end for
evaluation problem of | can be updated for this new setting as:
snnus ot |i
{(s)) (())))(ss) vix (s) (s)
min
min E vi
where N agents collaboratively learn shared feature representation + via server, along with personalized local weight vector {0r, Vij using local trajectories at each agent.
Remark 2.1. The learning of shared feature representation + in PFEDRL is related ideas from representation learning theory (Agarwal et al. 2020 2023), and this is believed to achieve better generalization performance with relatively small training data. In conventional FedRL, the feature representation & is given and fixed. Indeed, as we numerically verify in Section 4.3 presents better generalization performance new environments.
our PFEDRL
PFEDRL-REP ALGORIHHM
We now propose a class of algorithms called PFEDRL-REP that realize PFEDRL with shared rep- resentations. PFEDRL-REF alternates comprises of three main steps for each agent at each commu- nication round: (1) local weight vector update; (2) a local feature representation update; and (3) global feature representation update via the server.
Steps and 2: Local weight and feature representation updates. At round agent performs an update on its local weight vector given its current global feature representation &; and local weight vector 8 We allow each agent to perform K steps of local weight vector updates. Once the updated local weight vector 0 is obtained, each agent executes a one-step local update on its feature representation to obtain $ 1+1/2'
We represent these updates using the following generic notation:
1+1 WEIGHT UPDATE L, 0 Or,, K) and क t+1/2 = FEATURE UPDATE o 0 (3)
where a₁ and 3, learning rates for the weight and feature updates, respectively. The generic functions WEIGHT UPDATE and FEATURE UPDATE will be specialized to the particulars of the un- derlying RL algorithm: in Section we discuss the of TD with linear function approximation in detail, and in Appendix
C
we show instantiations of Q-learning and DQN in our framework.
Step 3: Server-based global feature representation update. The server computes an average of the received local feature representation updates Φ //1//2 from all agents obtain the next global feature representation #.+1 as
N
$
(4)
N
1+122
The PFEDRL-REP procedure repeats 3 and (4) and is summarized in Algorithm 1 and Figure We emphasize that because PFEDRL-REP operates an RL setting, there is no ground truth for the value function and learning occurs through interactions with an MDP environment, resulting in non-i.i.d. data. In contrast, in the standard FL setting (where shared representations have been investigated), there exists known ground truth and training data are sampled in i.i.d. fashion (Collins et al. 2021 Tziotis et al.) 2023) Xiong et al. 2024). The non-i.i.d. (Markovian) data the main technical challenge that need overcome.
4

---

Published as a conference paper at ICLR 2025
Sinver
Serve
Server
Agpin11
Agrit
Agmit.3
Apenl
AApent 2 (1h,
Agent 3 (4,,
Agpni 1 (Orey, Di
Adpent 2
Agrin
(l)
(0,,))
(9,,0)
(0..)l.)
(Phuu .
Faw:
Ens.
Rav,
Atthebaghooingg roundtt
Welght Liccite
Reprossniation Ljcktte
(a) At the beginning of round 1, 2, 3 has a local weight vector 0ü and global feature representation Φ, (b) Using ( 0.)) each agent performs Kstepp update to obtain 0/+1 in (3. Note that o, remains unchanged at this step. (c) Agent i updates the feature representation by executing a one-step update obtain $ 1+1 as (3), which depends on both 01+1 and & Finally, each agent shares & 1+1/2 with the server, which then executes an averaging step as in (4) to produce L Updated parameters highlighted in red, while shared parameters (the global feature representation) are in blue.
Figure 1:
An illustrative example of PFEDRL-REP for 3 agents.
I, each agent
4 PFEDTD-REP WITH LINEAR REPRESENTATION
We present PFEDTD-RE REP, instance of PFEDRL-REP paired with TD learning and analyze its convergence a linear representation setting.
4.1 PFEDTD-REE ALGORITHM DESCRIPTION
Here, the goal of N agents is to collaboratively solve problem when the underlying RL algorithm TD learning. We first need to specify WEIGHT. JUDDATE and FEATURE_UPDAYS of Algorithm for the case of TD. At time step k, the state of agent is sf. and its value function can be denoted as V(st) (( 0 in linear representation setting. By the standard one-step Monte Carlo approximation used in TD, we compute ((s = r + (s The TD error is defined as
c = ((s)) v(s.) = r + 2 (
क (
(5)
The goal of agent is to minimize the following loss function for every of ES
C' ( (s 0-) = ((s)) - vssil
(6)
with V(s() treated as constant. We now denote the Markovian observations of agent i at the k-th time step of communication round as Xi tion sequences (X/,A> Yt, k- differ across agents in heterogeneous environments. We assume that {Xi,, t,k, VI, k- are statistically independent across all agents.
(h. Note that the observa-
Local weight vector update. As in line of Algorithm [ given the current global feature represen- tation s each agent takes K local update steps its local weight vector 0 as
0 02 + ai g g0;,,_,, 4., Xixx1);
for k E [K], where g(0) -- s x xi is the negative stochastic gradient of the loss function L (s, 0(,,. with respect 0, given the current feature representation 1
g(0 2. Xi,x-)) == = LoL (Po(si, 1). 1-) (
(8)
Since there are K steps of local updates, we denote 01+1 = 0- We further add norm-scaling (i.e.. clipping) step for the updated weight vectors 0/ i.c., enforcing | < B, to stabilize the
update. This is essential for the finite-time convergence analysis in Section 4.2 and this technique is
widely used in conventional TD learning with linear function approximation [hhandari et al. 2018).
Local feature representation update. As in line 5 of Algorithm I given the updated local weight vector i agent then executes a one-step local update on the global feature representation:
c + 1 1+ L, {Xi 1k-1)
(9)
++1/
5

---

Published as a conference paper at ICLR 2025
where h(0)41 L2: (xx 11 is the negative stochastic gradient of the loss L (' with respect to the current global feature representation 2,, satisfying
h(0 +1 1, xix-1) == Vo ( ( 1).0 = 5% 0/+1 T
(10)
Server-based global feature representation update. As in line 7 of Algorithm | the server then averages the received local feature represeniation updates in (9) to obtain the next global feature representation:
N h(0 Le, 1 18
@ 2++ = 0 + +1+
Σ
(11)
N
The full pseudo-code of PFEDTD- REP is given in Appendix
4.2 CONVERGENCE ANALYSIS
The coupled updates in and can be viewed as a federated nonlinear wo-timessale stochastic
approximation (2TSA) DDan 2021) with Markovian noise, with 0 updating on a faster timescale
and C slower timescale. We aim establish the finite-time convergence rate of the 2TSA coupled updates and This is equivalent to finding solution pair ( (Bir, Vi)) such tha
E Rot
g0" Φ", x;) = - and E Be ee P (le))) [h "., ";,x))] =0 (12)
hold for all Markovian observations X; Here, 2² is the unknown stationary distribution of state sf of agent at t, and p is the transition kernel of agent i under policy n².
Although the root (Φ*, { ( Vi}) of the nonlinear 2TSA in 7 and is not unique due to simple permutations (rotations), is proved in Tsitsiklis & Van Roy 1996) that the standard TD iterates converge asymptotically to a vector 0" given fixed feature representation & almost surely, where 6" the unique solution of certain projected Bellman equation. Hence, for agent i, in order to study the stability of 0 when the feature representation + is fixed, note that there exists mapping that maps & to the unique solution of Bei Bel P(,
$ X;)]= = -.
Inspired by Doan (2020), the finite-time analysis of a 2TSA boils down to the choice of two step sizes {ax, Br, V and Lyapunov function that couples the two iterates in and (11 We first define the following two error terms:
+ $ and 0 = 0 (((),), Vi E [N],
(13)
which together characterizes the coupling between (0 ViJ and c.. If (0.11: ViJ and é, go to zero simultancously, the convergence of (10 Vi). Φ.) to (10 vi,, x*) can be established. Thus, to prove the convergence of ((0i+1, vi], 4.) of the 2TSA in 7 and [1] to its true value (Ø", (0Sr, (v)),, we define the following weighted Lyapunov function to explicitly couple the fast and slow iterates
3 at
M M
1 1+1 V.).
Σ 10 - 2
(14)
Remark 4.1. Note that the Lyapunov function (4) for 2TSA does not inherently require the solution to be unique. If multiple solutions or equilibria exist, the Lyapumov function should still be able to show that the system will converge to of these possible equilibria, ensuring that the system's state does not diverge and eventually stabilizes at some equilibrium point, which highly depends on the initialization of o To clarify this, in the rest of this paper, we use & to clearly indicate the dependence of the initialization of Φ and x- in (I4) is interchangeable with D which denotes the optimum close the initial point.
"The root (Φ' ( Vi)) of the nonlinear 2TSA in and can be established by using the ODE method following the solution of suitably defined differential equations (Don 2021: 2020 Chen al. 2019 in

---

Published as a conference paper at ICLR 2025
Our goal is characterize the unite-time convergence of E[M Vi], 6:)]. the Lyapunov function in (14). We start with some standard assumptions first.
Assumption 4.2. The learning rates a, and B, satisfy the following conditions: (i Σi=o ar = oc, (ii) Et= of o0,
(ii) E Pt se. (iv) Ei= Bf oc, (v) 3:/a; is non-increasing and
(vi) lim 3 0.
Assumption 4.3. Agent i's Markov chain {x;} is irreducible and aperiodic. Hence, there exists a unique stationary distribution μs (Levin & Peres 2017) and constants C > 0 and p E (0, 1) such that (P distance (Levin & Peres 2017).
x). S Cp", vk 2 0, r E R. where drv (:; is the total-variation (TV)
Remark 4.4. Assumption 4.3 implies that the Markov chain induced by n' admits unique station- ary distribution μ³. This assumption is commonly used in the asymptotic convergence analysis of stochastic approximation under Markovian noise (Borkar 2009
Chen al. 2019).
We can define the steady-state local TD update direction as
g(0', +) := Bei at (11)) (g(())((0 (0', ④ X;)],
h(0' ②) = E, Bap -t (-+4))
[h(O 2 x;)].
(15)
Definition 4.5 (Mixing time, similar to Chen et al. (019 First, define the discrepancy term
(0', Φ, x) = max{ 2, x:) Xu= - 8 R) [E[[(("', 2 x;) Xo xu==]-h((( .)
For 0, the mixing time defined as
T8 max min 2 1: ( Φ x) S 8( 4'1 + 10 n 4 || 1), V (0 ,,,)} :E[N]
which describes the time it takes for all agents' trajectories (Markov chains) to be well-represented by their stationary distributions.
Lemma 4.6. g(0, x in
(8)
is globally Lipschitz continuous w.r.t 0 and & uniformly in X, i.e, < L, 021 + 0 22ll), VX E X.
(g 1 X) x)0 22,
X
Lemma 4.7. h(0, 2 X) in
IO
is globally Lipschitz continuous w.r.t 0 and & uniformly in X, iee [ 021 | 22ll), VX X.
||h(0, 1 X) h(0,, 22
Lemma 4.8. y (). Vi is Lipschitz continuous in b i.e., ||v (.,) y ( - Lyy | $2
For notational simplicity, we let L == max Ln, Ly and assume that I. is the common Lipschitz constant in Lemmas 4.8 in the following.
Remark 4.9. The Lipschitz continuity of guarantees the existence of a solution & to the equilib- rium 1 for fixed f, while the Lipschitz continuity of g and y² ensures the existence of solution 0² of 12 when fixed.
Lemma 4.10. There exists a w such that V Φ, and V
1 4ú ()) 2)) < - 1
(0 n (Q:-1) g())))) 2;))) < -w y ||
Remark 4.11. Lemma 4.10 guarantees the stability of the two-timescale update in 7 and and can be viewed as the monotone property of nonlinear mappings leveraged in Doan (2020); Chen et al. 2019).
Lemma 4.12. Under Assumption 4.3 and Lemma 4.6 and 4.7 there exist constants C > 0, p E
(0, 1) and L1 max(Ly, Ln, maxx (0 Φ" X), maxx h 0 8² X)) such that
log(1/8) log (2L Cd) and
T8 S
lim 878 0. 6
log(1/p)
4.2.1 MAIN RESULTS
We now present our main theoretical results in this work. Theorem 4.13. Ler T 2 2T. for some 5 > O. Suppose that the learning rates are chosen as
at 0n/ (t 2/5/ and B, = Bo/ (t. 2),

---

Published as a conference paper at ICLR 2025
where og 21/(2L/2 Bo S w/2, and L max Ln Ly). We have
M( 20)
C1
M(10 0T+1) 5
+
(T 2)
(T + 2)2/33
N
C2
R 00 Σ
1/ (B6)1)))
(16)
(T' 2)2/3
:=1
with C1 = (4a03, (38 B + + 20 + 3K 22222B3) 88008009)
and C2 = (14473 12 1.202 4T.2 ao0
The first term of the right-hand side of (16) corresponds to the bias due to initialization, which goes to zero at a rate (1) The second term is due to the variance of the Markovian noise. The third term corresponds to the accumulated estimation error of the two-timescale update. The second and third terms decay at rate (1/T2//5) and hence dominate the overall convergence rate in (16.
Remark 4.14. Doan 2020) provided the first finite-time analysis for general nonlinear 2TSA under i..d. noise, and then extended it to the Markovian noise setting under the assumptions that both g and h functions are monotone in both parameters (Doan 2021). Since Doan (2021) leverages the methods from Doan (2020), it needs a detailed characterization of the covariance between the error induced by Markovian noise and the residual error of the parameters in (13), rendering the convergence analysis much more intricate. To address this, we take inspiration from (Srikant & Ying 2019) (which operates the single-timescale SA setting) and use a Lyapunov drift approach to capture the evolution of two coupled parameters under Markovian noise. Characterizing the impact of a norm norm-scaling step further distinguishes our work.
Corollary 4.15. Suppose that 3o o o(N22/3) and that T > N. Then, we have
1 N2/3(T + 2)2/3
M(10 8 S
Remark 4.16. Corollary 4.15
indicates that attain accuracy, it requires T= ON-
steps. In this sense, we prove that PFEDTD- REP achieves a linear convergence speedup with re- spect to the number of agents N, i.e., the number of steps until e-convergence is multiplied by factor of N-1 In other words, we can proportionally decrease T as N increases. To our knowl- edge, this is the first linear speedup result for personalized FedRL with shared representations under Markovian noise, which is highly desirable since it implies that one can efficiently leverage the massive parallelism
large-scale systems. Recently, [Shen al. (2023) considered 2TSA in federated RL setting and achieved convergence rate of O (T 270 and thus sample complexity of O (e -5/2). In contrast, our method can converge quicker and enjoys a lower sample complexity, and the convergence speed matches the best-known convergence speed for non-linear 2TSA under even i.i.d. noise (Doan 2020). In addition, we note that single-timescal. (SA) methods may enjoy a faster convergence speed and lower sample complexity. However, the analysis of the 2TSA setting is more involved, as there are two parameters to be updated in a coupled and asynchronous manner. To our knowledge, there are no existing works in the 2TSA settings that achieve the same conver- gence rate or sample complexity as those in the SA settings. It may be an interesting direction to investigate for the community. Finally, similar to FL settings (Collins al. 2021) the local step does not hurt the global convergence with a proper learning rate choice.
Remark 4.17. The gradient tracking technique discussed in Zeng & Doan 2024 could potentially be effective in handling Markovian sampling and improving the current convergence rate to OTT-1). However, it is unclear if it can be applied PFEDRL-REP, because Zeng & Doan (0244) only considers single-agent settings under i.i.d. noise. Furthermore, Zeng & Doan (2024) assumes a second-order variance bound of the stochastic function, while our analysis does not include such an assumption. Investigating these is out of the scope of this work, which already considers a very challenging setting.
4.2.2 INTUITIONS AND PROOF SKETCH
We highlight the key ideas and challenges behind the convergence rate analysis of PFEDTD- REP with two coupled parameters, which is example of federated nonlinear 2TSA. With the defined
8

---

Published as a conference paper at ICLR 2025
Campanson Value Function Estimotes by Different Alporthins
Negative Expected Valye
10
10
TedTO
Profto Mot
5
TO
FedID
5:31
0
100
200
0
100
200
100
200
Episode
Agent
Value function estimates.
(b) Convergence speed of value function learning.
Figure 2: Comparisons in CliffWalking Environment with 3 agents.
Lyapunov function in (14), the key is to find the drift between M((0;+1 vi). +F) in the f-th com- munication round and M ( Vij, 1)= in the (t - 1)-th communication round. To achieve this, we separately characterize the drift between Lu+1 and i. and the drift between 0; and 0j, vi. We emphasize the three main challenges in characterizing the drift: (i) how to bound the stochastic gradient with Markovian samples; (ii) how to leverage the mixing time T to handle the biased pa- rameter updates due to Markovian noise; and (iii) how to deal with multiple local updates for the local weight vector 0'
By the mixing time property of MDPs, we have that the gap between the biased gradient at each time step and the true gradient can be bounded when the time step exceeds the mixing time т, as defined in Definition 4.5
To characterize the effect of local updates, the key idea is to bound the gradient at the initial local step and the gradient at the final local steps, which can be done by leveraging the Lip- schitz property of those gradient functions in Lemmas 4.6 4.7 and 4.8 See Appendix [.1.] and Ap- pendix F.1.2 for details. Once we establish the drift of the Lyapunov function, the remaining task is to select suitable dynamic two-timescale learning rates (or, Vt) and {3,, Vt} for the weight vector up- date in ( and the feature representation update in (9) respectively, See Appendix F.1.3 1.3 for details.
4.3 NUMERICAL EVALUATION
We empirically evaluate the performance of PFEDTD-REP. We consider a tabular CliffWalking environment (Brockman et al. (016) with a 4 x 12 grid world, where 3 agents evaluate 3 different policies. The dimension for the feature representation and weight vector is set to be 6. We compare PFEDTD-REP with (i) "TD": each agent independently leverages the conventional TD without com- munication; and (ii) "FedTD" without personalization (Khodadadian et al. 2022 Dal Fabbro et al.) 2023) listed in Table ! As shown in Figure 2a PFED ensures personalization among all agents while FedTD tends to converge uniformly among all agents. Further, PFEDTD-REP at- tains values much closer to the ground-truth achieved by TD for each agent compared to FedTD; and PFEDTD-RED converges much faster than TD. For instance, agent only needs 50 episodes to converge under PFED
while it takes more than 150 episodes converge under TD. as illus- trated in Figure 26 The improved convergence performance of PFEDTD-REE further supports our theoretical findings that leveraging shared representations not only provides personalization among agents in heterogeneous environments but yield faster convergence.
5 APPLICATION TO CONTROL PROBLEMS
In this section, we briefly discuss how our proposed PFEDRL-RE: REP framework can be applied to the control problems in RL. More details are provided Appendices C.3 (i.e., Algorithm and G
"FEDDQN-RUS REP (an instance of PFEDRL- REP paired with DQN) leverages shared representations to learn a common feature space that captures the underlying dynamics and features relevant across the target network is
different but related tasks encountered by various agents. In PFEDDQN- a critical component that provides stability to the learning process by serving as a relatively static benchmark for calculating the loss during training updates (Mnih et al. 2015). The target network's 's architecture mirrors that of the main network, including the shared representation model. However, its parameters are updated less frequently. This setup ensures that the calculated target values, which guide the policy updates, are based on a consistent representation of the environment's 's state, as encoded by the shared representation model. The synergy between the target network and the repre-

---

Published as a conference paper at ICLR 2025
PracDoN-KR M DON
-100"
Rebum
Retum
Robum
Retum dis
0.5
400
509
300
402
500
102
pioode
Epsode
pole length Generatzation pole length
[a] Retur length
(b) G Generaizzzion poie length
(a) Cartpole environment.
(b) Acrobot environment.
Figure 3: Comparisons in control problems.
sentation model is thus central to achieving stable and convergent learning. In Line 13 of Algorithm the algorithm performs a scheduled update of the shared representation & of the main network's parameters with the guidance of the target network. In Line 18 of Algorithm [ every Turget steps, the algorithm performs a scheduled update of the target network's parameters by copying over the parameters from the main network. This step is essential for maintaining the stability of the learning process, as ensures that the target values against which the policy updates are computed remain consistent and reflects the most recent knowledge encoded in the shared representation. The update frequency is carefully chosen to balance learning stability with model adaptivity.
Numerical evaluation. We consider a modified CartPole environment (Brockman et al. 2016) by changing the length of pole to create heterogenous environments (Jin et al. 2022). Specifically, we consider 10 agents with varying pole lengths from 0.38 to 0.74 with step size of 0.04. We compare PFEDDQN REP with (i) a conventional DQN where each agent learns its own environment independently; (ii) federated version of DQN (FedDQN) that allows all agents to collaboratively learn a single policy (without personalization); (iii) two federated algorithms without personalizing, FedQ-K (Khodadadian et al. 2022) and L.FRI (Liu et al., 2019); and (iv) two personalized algorithms, PerDQNAvg (Jin et al., (202)) and FedAsynQ-ImAv 2023). We randomly choose one agent and present its performance in Figure Bal We observe that our PFEDDQN-REP obtains larger reward than the baselines without personalization and achieves the maximal return faster than existing personalized algorithms. We further evaluate the effectiveness of the shared representation learned by PFEDDQN- REP when generalizing to a new environment. As shown Figure 3a PFEDDQN-REF generalizes quickly to the new environment. Similar observations can be made from Figure 3b using Acrobot environments (see details Appendix G). In summary, the significance of our PFEDRL-REI framework lies in its superior performance in heterogeneous environments compared to existing algorithms that do not incorporate personalization. Additionally, our PFEDRL-REE framework also enables quick adaptation to new, previously unobserved environments.
(Woo et al.)
Limitations and open problems. In this paper, we characterize the finite-time convergence rate for PFEDTD-Re REP with linear feature representation. However, analysis is not directly applicable to control problems. Consider Q-learning. The difficulty arises because the update for Q-learning not linear operation with respect the shared representation and local weights. Additionally, Q-learning is typically combined with deep neural networks, where the Q-function is approximated by neural network as in PFEDDQN- REP.. The complexity is further compounded in personalized federated RL frameworks, where multiple agents share common representation while maintaining personalized local weights. Given the promising experimental results on control, whether we can extend our theoretical results to the control setting remains open problem.
ACKNOWLEDGEMENTS
This work was supported in part by the National Science Foundation (NSF) grants 2148309, 2315614 and 2337914, and was supported in part by funds from OUSD R&E, NIST, and industry partners as specified in the Resilient & Intelligent NextG Systems (RINGS) program. Any opinions, findings, and conclusions or recommendations expressed this material are those the authors and do not necessarily reflect the views of the funding agencies.
10

---

Published as a conference paper at ICLR 2025
REFERENCES
Alekh Agarwal, Sham Kakade, Akshay Krishnamurthy, and Wen Sun. Flambe: Structural com- plexity and representation learning of low rank mdps. Advances neural information processing systems, 33:20095-200 20107, 2020.
Alckh Agarwal, Yuda Song, Wen Sun, Kaiwen Wang. Mengdi Wang, and Xuezhou Zhang, Prov- able benefits representational transfer reinforcement learning. The Thirty Sixth Annual Conference Learning Theory, pp. 2187-2187. PMLR, 2023.
Manoj Ghuhan Arivazhagan, Vinay Aggarwal, Aaditya Kumar Singh, and Sunav Choudhary, Fed- crated learning with personalization layers. arXiv preprint arXiv:1912.00818, 2019.
Yoshua Bengio, Aaron Courville, and Pascal Vincent. Representation learning: A review and new perspectives. IEEE transactions on pattern analysis and machine intelligence, 3598),1798-1888, 2013.
El Houcine Bergou, Konstantin Burlachenko, Aritra Dutta, and Peter Richtárik. Personalized feder- ated learning with communication compression. arXiv preprint arXiv:2209.053 0414 2022.
Jalaj Bhandari, Daniel Russo, and Raghav Singal. A finite time analysis of temporal difference learning with linear function approximation. In Conference on learning theory, pp. 1691- PMLR, 2018.
Vivek S Borkar. Stochastic approximation: dynamical systems viewpoint, volume 48. Springer, 2009.
Karim Boubouh, Robert Basmadjian, Omid Ardakanian, Alexandre Maurer, and Rachid Guerraoui. Efficacy of temporal and spatial abstraction for training accurate machine learning models: A case study in smart thermostats. Energy and Buildings, 296: 1 13377, 2023.
Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojciech Zaremba. Openai gym. arXiv preprint 1rXiv: 2016.
Fei Chen, Mi Luo, Zhenhua Dong, Zhenguo Li, and Xiugiang He. Federated meta-learning with fast convergence and efficient communication. arXiv preprint arXiv: /002:08800.07855 2018.
Zaiwei Chen, Sheng Zhang, Thinh Doan, Siva Theja Maguluri, and John-Paul Clarke. Perfor- mance of q-learning with linear function approximation: Stability and finite-time analysis. arXiv preprint 1rXiv:1905.11422 pp. 4, 2019.
Liam Collins, Hamed Hassani, Aryan Mokhtari, and Sanjay Shakkottai. Exploiting shared repre- sentations for personalized federated learning. In International conference on machine learning, pp. 2089-2099. PMLR, 2021.
Nicolò Dal Fabbro, Aritra Mitra, and George J Pappas. Federated td learning over finite-rate erasure channels: Linear speedup under markovian sampling. IEEE Control Systems Letters, 2023.
Thinh T Doan. Nonlinear two-time-scale stochastic approximation: Convergence and finite-time
performance. arXiv preprint arXiv:2011.01866 2020.
Thinh T Doan. Finite-time convergence rates of nonlinear two-time- scale stochastic approximation under markovian noise. arXiv preprint arX1v:2104.01627, 2021.
Alireza Fallah, Aryan Mokhtari, and Asuman Ozdaglar Personalized federated learning: A meta- learning approach. arXiv preprint arXiv:2002.07948, 2020.
Xiaofeng Fan, Yining Ma, Zhongxiang Dai, Wei Jing. Cheston Tan, and Bryan Kian Hsiang Low. Fault-tolerant federated reinforcement learning with theoretical guarantee. Advances in Neural Information Processing Systems, 34: 1007-1011 2021.
Andrew Hard, Kanishka Rao, Rajiv Mathews, Swaroop Ramaswamy, Françoise Beaufays, Sean Augenstein, Hubert Eichner, Chloé Kiddon, and Daniel Ramage. Federated learning for mobile keyboard prediction. arXiv preprint arXiv: a 18/118//10000 2018.

---

Published as a conference paper at ICLR 2025
Ahmed Imteaj, Khandaker Mamun Ahmed, Urmish Thakker, Shiqiang Wang, Jian Li, and M Hadi Amini. Federated learning for resource-constraing IoT devices: Panoramas and state of the art. Federated and Transfer Learning. pp. 7-27, 2022.
Chi Jin, Zhuoran Yang, Zhaoran Wang, and Michael Jordan. Provably efficient reinforcement learning with linear function approximation. In Conference on Learning Theory, pp. 2137-2143. PMLR, 2020.
Hao Jin, Yang Peng, Wenhao Yang, Shusen Wang, and Zhihua Zhang. Federated reinforcement learning with environment heterogeneity. In International Conference on Artificial Intelligence and Statistics, pp. -8. PMLR, 2022.
Sajad Khodadadian, Pranay Sharma, Gauri Joshi, and Siva Theja Maguluri. Federated reinforcement learning: Linear speedup under markovian sampling. In International Conference on Machine Learning. pp. 10997-11057. PMLR, 2022.
Walid Krichene, Nicolas Mayoraz, Steffen Rendle, Shuang Song, Abhradeep Thakurta, and Li Zhang. Private learning with public features. arXiv preprint arXiv:2310.1545. 2023.
Guangchen Lan, Dong-Jun Han, Abolfazl Hashemi, Vaneet Aggarwal, and Christopher G Brinton. Asynchronous federated reinforcement learning with policy gradient updates: Algorithm design and convergence analysis. arXiv preprint unXiv:2404.08005 2024.
Yann LeCun, Yoshua Bengio, and Geoffrey Hinton. Deep learning. nature, 521(7553):436 444, 2015.
David A Levin and Yuval Peres. Markov chains and mixing times, volume 107. American Mathe- matical Soc., 2017.
Boyi Liu, Lujia Wang, and Ming Liu. Lifelong federated reinforcement learning: learning ar- chitecture for navigation in cloud robotic systems. IEEE Robotics and Automation Letters, 4(4): 4555 4565-4562, 2019.
Kiwan Maeng, Haiyu Lu, Luca Melis, John Nguyen, Mike Rabbat, and Carole-Jean Wu. Towards fair federated recommendation learning: Characterizing the inter-dependence of system and data heterogeneity. Proceedings of the 16th ACM Conference on Recommender Systems, pp. 156- 167, 2022.
Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas. Communication learning of deep networks from decentralized data. In Artificial intelli- gence and statistics, pp. 1273-1282. PMLR, 2017.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Mare G Belle- mare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, et al. Human-level control through deep reinforcement learning. nature, 51 518(7540):529-533 2015.
Volodymyr Mnih, Adria Puigdomenech Badia, Mehdi Mirza, Alex Graves, Timothy Lillicrap, Tim Harley, David Silver, and Koray Kavukcuoglu. Asynchronous methods for deep reinforcement learning. International conference on machine learning, pp. 1928-1937. PMLR, 2016.
Aditya Modi, Jinglin Chen, Akshay Krishnamurthy, Nan Jiang, and Alekh Agarwal. Model-free representation learning and exploration in low-rank mdps. arXiv preprint arXiv:2/02.07035, 2021.
Chetan Nadiger, Anil Kumar, and Sherine Abdelhak. Federated reinforcement learning for fast personalization. In 2019 IEEE Second International Conference Artificial Intelligence and Knowledge Engineering (AIKE), pp. 123-127. IEEE, 2019.
Dinh C Nguyen, Ming Ding. Pubudu Pathirana, Aruna Seneviratne, Jun Li, and H Vincent Poor. Federated learning for Internet of Things: A comprehensive survey, IEEE Communications Sur- veys & Tutorials, 23(3):1622-11)
2021.
Jiaju Qi, Qihao Zhou, Lei Lei, and Kan Zheng. Federated reinforcement learning: Techniques, applications, and open challenges. arXiv preprint arXiv:2108. 1887, 2021.
12

---

Published as a conference paper at ICLR 2025
Sudeep Salgia and Yuejie Chi. The cumple-communiciin complexity trade-off in federated q- learning. arXiv preprint arXiv:2408. 2024.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06344 2017.
Han Shen, Kaiqing Zhang, Mingyi Hong, and Tianyi Chen. Towards understanding asynchronous advantage actor-critic: Convergence and linear speedup. IEEE Transactions on Signal Processing, 2023.
David Silver, Guy Lever, Nicolas Heess, Thomas Degris, Daan Wierstra, and Martin Riedmiller. Deterministic policy gradient algorithms. In International conference on machine learning, pp- 387-395. Pmlr, 2014.
Virginia Smith, Chao- Kai Chiang, Maziar Sanjabi, and Ameet S Talwalkar. Federated multi-tass learning. Advances in neural information processing systems, 30, 2017.
Rayadurgam Srikant and Lei Ying. Finite-time error bounds for linear stochastic approximation andtd learning. Conference on Learning Theory, pp. 2803-2830. PMLR, 2019.
Richard Sutton. Learning to predict by the methods of temporal differences. Machine learning, 3:9 44, 1988.
Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT press, 2018.
John Tsitsiklis and Benjamin Van Roy. Analysis of temporal-diffferene learning with function approximation. Advances in neural information processing systems, 9, 1996.
Isidoros Tziotis, Zebang Shen, Ramtin Pedarsani, Hamed Hassani, and Aryan Mokhtari. Straggler- resilient personalized federated learning. Transactions Machine Learning Research, 2023.
Han Wang. Aritra Mitra, Hamed Hassani, George J Pappas, and James Anderson. Federated tem-
poral difference learning with linear function approximation under environmental heterogeneity. arXiv preprint 230.022300..2222.. 2023a.
Sid Wang. Ashish Shenoy, Pierce Chuang, and John Nguyen. Now it sounds like you: Learning personalized vocabulary on device. arXiv preprint arXiv:2305.03584 2023b.
Christopher JCH Watkins and Peter Dayan. Q-learning. Machine learning, 8:279-292. 292, 1992.
Jiin Woo, Gauri Joshi, and Yuejie Chi. The blessing of heterogeneity in federated q-learning: Linear speedup and beyond. In International Conference Machine Learning. PMLR, 2023.
Jiin Woo, Laixi Shi, Gauri Joshi, and Yuejie Chi. Federated offline reinforcement learning: Col- laborative single-policy coverage suffices. In Forty-first International Conference on Machine Learning. 2024.
Guojun Xiong, Gang Yan, Shiqiang Wang, and Jian Li. Deprl: Achieving linear convergence speedup in personalized decentralized learning with shared representations. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 38, pp. 16103-16111, 2024.
Minrui Xu, Jialiang Peng, BB Gupta, Jiawen Kang, Zehui Xiong. Zhenni Li, and Ahmed Abd El- Latif. Multiagent federated reinforcement learning for secure incentive mechanism in intelligent cyber-physical systems. IEEE Internet of Things Journal, 9(2)):2095--22109 2021.
Timothy Yang, Galen Andrew, Hubert Eichner, Haicheng Sun, Wei Li, Nicholas Kong, Daniel Ra- mage, and Françoise Beautays. Applied federated learning: Improving google keyboard query suggestions. arXiv preprint arXiv: 1812.02903 2018.
Zhenyuan Yuan, Siyuan Xu, and Minghui Zhu. Federated reinforcement learning for generalizable motion planning. In 2023 American Control Conference (ACC), pp. 78-83. IEEE, 2023.
Sihan Zeng and Thinh Doan. Fast two-time-scale stochastic gradient method with applications in reinforcement learning. In The Thirty Seventh Annual Conference on Learning Theory, pp. 5166- 5212. PMLR, 2024.
13

---

Published as a conference paper at ICLR 2025
Chenyu Zhang, Han Wang, Aritra Mitra, and James Anderson. Finite-time analysis of on-policy heterogeneous federated reinforcement learning. arXiv preprint arXiv:2401.15273, 2024.
Kaiqing Zhang, Zhuoran Yang, and Tamer Başar. Multi-agent reinforcement learning: A selective overview of theories and algorithms. Handbook of reinforcement learning and control, pp. 321- 384, 2021.
Sai Qian Zhang, Jieyu Lin, and Qi Zhang. A multi-agent reinforcement learning approach for efficient client selection in federated learning. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 36, pp. 90919,, 2022a.
Tuo Zhang, Lei Gao, Chaoyang He, Mi Zhang, Bhaskar Krishnamachari, and A Salman Avestimehr. Federated learning for the Internet of Things: Applications, challenges, and opportunities. IEEE Internet of Things Magazine, 5(1):24 29, 2022b.
Xuezhou Zhang, Yuda Song, Masatoshi Uchara, Mengdi Wang, Alckh Agarwal, and Wen Sun. Efficient reinforcement learning in block mdps: A model-free representation learning approach. In International Conference on Machine Learning, pp. 26517-26547. PMLR, 2022c
Zhong Zheng, Fengyu Gao, Lingzhou Xue, and Jing Yang. Federated q-learning: Linear regret speedup with low communication cost. In The Twelfth International Conference on Learning Representations, 2024.
14

---

Published as a conference paper at ICLR 2025
A RELATED WORK
Single-agent reinforcement learning. RL is a machine learning paradigm that trains agents to make sequences of decisions by rewarding desired behaviors and/or penalizing undesired ones in a given environment (Sutton & Barto 2018). Starting from Temporal Difference (TD) Learning (Sut-) ton 1988). which introduced the concept of learning from the discrepancy between predicted and actual rewards through episodes, the widely used Q-Learning (Watkins & Dayan [992) emerged, advancing the field with an off-policy algorithm that learns action-value functions and enables pol- icy improvement without needing a model of the environment. Later on, the introduction of Deep Q-Networks (DQN) (Mnih et al. 2015) marked a sigmificant leap, integrating deep neural networks with Q-Learning to handle high-dimensional state spaces, thus enabling RL to tackle complex prob- lems. Subsequently, policy-based algorithms such as Proximal Policy Optimization (PPO) (Schul-) man et al., 2017) and deep Deterministic Policy Gradients (DDPG) (Silver et al.] 2014). leverage the Actor-Critic framework to provide more stable and robust ways to directly optimize the policy, overcoming challenges related to action space and variance.
Federated reinforcement learning. Jin et al. (2022) introduced a FedRL framework with N agents collaboratively learning policy by averaging their Q-values or policy gradients. Khhoadadiann al. (2022) provided a convergence analysis of federated TD (FedTD) and Q-learning (FedQ) when N agents interact with homogeneous environments. A similar FedTD was considered in [aal Fabbro et al. (2023). and expanded to heterogeneous environments in WWnn et al. (2023a). Woo et al. (2023) analyzed (a)synchronous variants of FedQ in heterogencous settings, and an asynchronous actor-critic method was considered in Shen et al. (2023) with linear speedup guarantee only under i..dd samples. Zhang et al. (2024) provided a finite-time analysis of FedSARSA with linear func- tion approximation (i.e.. fixed feature representation). To facilitate personalization in heterogeneous settings, Din et al. (2022) proposed a heuristic personalized FedRL method where agents share com- model, but make use of individual environment embeddings. There is also related paper Fan let al. (202T). which considers a special setting where each agent can be Byzantine and suffers ran- dom failure in every round. In Fan et al. (2021). convergence was established based on i.idd. noise.
Personalized federated learning (PFL). In contrast to standard FL where a single model is learned, PFL aims to learn N models specialized for N local datasets. Many PFL methods have been developed. including but not limited to multi-task learning (Smith et al.] 2017), meta-learning (Chen 2018). and various personalization techniques such as local fine-tuning (Fallah al., 2020 layer personalization (Arivazhagan et al.) 2019), and model compression (Bergou et al., 2022. J. Another line of work (Collins et al.) 2021: Xiong al., (2024) leveraged the common representation among agents in heterogeneous environments to guarantee personalized models for federated supervised learning.
Representation learning in MDP. Representation learning aims to transform high-dimensional ob- servation to low-dimensional embedding to enable efficient learning, and has received increasing at- tention in Markov decision processs (MDP) settings, such as linear MDPs (Jin al. 2020), low-rank MDPs Modi et al. | Agarwal et al. 2020) and block MDPs (Zhang et al.] 2022c). However, open in the context of leveraging representation learning in PFedFL... In this work, prove that representation augmented PFedFL. forms a general framework as a federated two-timescale stochas- tic approximation with Markovian noise, which differs significantly from existing works, and hence necessitates different proof techniques.
Multi-agent reinforcement learning versus federated reinforcement learning. The advent of Multi-Agent Reinforcement Learning (MARL) expanded RL's applications, allowing multiple agents learn from interactions in cooperative. competitive, or mixed settings, opening new avenues for complex applications and research (Zhang et al., 202T). Multi-agent Reinforcement Learning (MARL. addresses scenarios where multiple agents operate within a shared or interrelated environ- ment, potentially engaging in both cooperative and competitive behaviors. The complexity arises from each agent needing to consider the strategies and actions of others, making the learning pro- cess highly dynamic. Federated Reinforcement Learning (FedRL) (Qi et al. 2021), contrasts with MARL by focusing on raw data. Instead, these agents might contribute towards centralized learning model without com- promising individual data privacy, addressing the unique challenges of learning from decentralized data sources.
distributed learning across agents that do not share their
15

---

Published as a conference paper at ICLR 2025
Algorithm 2 PFEDTD-REP
Input: Sampling policy T', Vi E [N]]
2: Initialize 0 Sy Vi [N], and randomly generate + IilSidd with each row being
unit-norm vector;
do 4: for i 1 N do
3: for 0,
5: for k = ,,..,
6:
Sample observations X 1-
7: Set 0/ 0 - + ag x- 2,, Xi x));
8:
end for
9: Scale 10 to B 0+11
> B, otherwise keep it unchanged;
10: Set 3
1/ $ + 3.
L. (Xi k 11k=
1
1+
11: Normalize Φ
as Φ
1/2
1+1/2
1+1/2
12: end for
Σi $
1+1 =
1+1/2"
14: end for
B PSEUDOCODE OF PFEDTD-RES -EP
In this section, we present the pseudocode of PFEDQ-REP as summarized in Algorithml 2
C APPLICATION TO CONTROL TASKS IN RL
The Q-function of agent i in environment M under policy x' is defined Q (s,a) E T (sf. )( 8, ab = a]. When the state and action spaces are large, it is com- putationally infeasible to store Qi (s, a) for all state-action pairs. One way to deal with to approximate the Q-function as Qixx' (s, = ( 0, where + E R|S]x|A[ is a feature repre- sentation corresponding to state-actions, and E Id a low-dimensional unknown weight vector. When & is given and known, this falls under the paradigm of RL or FedRL with function approxi- mation.
C.I PRELIMINARIES: CONTROL IN FEDERATED REINFORCEMENT LEARNING
Another task in RL is to search for optimal policy, which is called control problem, and one commonly used approach is Q-lcarning WWatkins & Dayan 1992). Under the FedRL framework, the goal of control problem is to let N agents collaboratively learn a policy x² that performs uniformly well across N different environments, i.e., x arg max, N ΣN E. Vii (s))]s)
do). where do the common initial state distribution in these N environments. Similar to this can be formulated the optimization problem in (17) to collaboratively learn common (non- personalized) weight vector 8 Vi E [N] when the feature representation $ ((8, a), Vs, a are given.
(())
min
N
1
Σ
E
imn"in" arall (+]a)
( Qi (s.a))l
(17)
1=1
Again, the superscript i highlight heterogeneous environments p among agents.
16

---

Published as a conference paper at ICLR 2025
C.2 CONTROL PERSONALIZED FEDRL WITH SHARED REPRESENTATIONS
The control problem 17 aims to learn t and {0', Vij simultaneously among all N agents via solving the following optimization problem:
r' (s. ('(0',0(s,,0) Qi (s.0)l² "
C (0,,))))i))
min
min E vi
(18)
smp souphtti
ams --1s) (11)
C.3 ALGORTHHM
In this subsection, we present two realizations of our proposed PFEDRL-REP in Algorithm one is PFEDQ-REP summarized in Algorithm federated Q-learning with shared representations, and the other is PFEDDQN-REP as outlined in Algorithm federated DQN with shared representations.
Algorithm 3 PFEDQ- REP
Input: Sampling policy x², Vi E [M]
Initialize 8 0, and sh. Vi [N], and randomly generate RIS||AA with each row being unit-norm vector.
2: for 0, ...,. 1.., 3: for 1, N do
do
4: for k = ,, do
5:
Sample observations X f- (s), ( 1
6:
With fixed i,, update 0i +- 0i 1 + ar (r/R-1 + max, 2r (
(s 10 end for
1- 1,
(s 1 aii
7:
8:
Scale ++11l
B if 01 > B, otherwise keep unchanged.
9:
if (s, a) E X 3k E (0....,F. K 1) then
10:
Update Φ (11/2(5,0) = ( a) (r,,rrsssss a) + max, 2,(5', as
0 t1 else
11:
12:
$ (1122(8, (8
end if
14: Normalize $ 1+1/2 as $
141/22
+
[ |
$1+1/3_
15: end for
T 1++1 f- E Φ 1/ Vi [N].
17: end for
D FIGURE ILLUSTRATIONS
We present some figures to further highlight the proposed personalized FedRL (PFEDRL) frame- work with shared representations.
Schematic framework of conventional FedRL. framework (Khodadadian et al. 2022), where N agents collaboratively learn a common policy (or optimal value functions) via server while engaging with homogeneous environments. Each agent generates independent Markovian trajectories, as depicted in Figure 4
We begin by introducing the conventional FedRL.
Schematic framework for our proposed PFEDRI. with shared representations. We introduce proposed personalized FedRL (PFEDRL) framework with shared representations in Figure [5] In PFEDRL, N agents independently interact with their own environments and execute actions according to their individual RL component parameterized by & and 0' Each agent performs local update on local weight vector 0, while jointly updating the global shared feature representation & through the server. Similarly, the update follows the Markovian trajectories.
Motivation of personalized FedRL. In the following, we also want to provide some examples showing that the conventional FedRI. framework may fail, as depicted in Figure 6 In Figure 6a
17

---

Published as a conference paper at ICLR 2025
Algorithm 4 PFEDDQN-REP
Initialize: The parameters ( O*) for each Q network (s, a), the replay buffer R', and copy
the same parameter from network initialize the target network (*, for agent Vi [N];
1: for episode = .,... E do
2: Get the initial state of the environment;
3: for
T 1 do
4: for i = . do
5: for k = -,.==...... .,..., K do
6:
Select action dz according to e-greedy policy with the current network
Q'(s1,6-1, a);
7:
Execute action d receive the reward 15, and the environment state transits to 8t,k;
8:
St,k) into the replay buffer Ri; Sample N data tuples from the replay buffer R²;
Store the tuple 1...-11.0,A--1 P(. ",,, d1,k-,
9:
10:
Update the local weight (1, k) by minimizing the loss compared with the target network Q with fixed representation c.; end for
11:
12:
Sample data tuples from replay buffer R':
13:
Update representation model locally by minimizing the loss compared with the target network Q with fixed weights 0e,.1 and yield t (+1/2' end for Average the representation model from all agents, i.e., P.t1 = E Φ 14122 end for
14:
17:
if mod(t, Tiarget) then
18:
update the target network Qi+ be copy the up-to-date parameters of Q network Q', Vi [N]
19:
end if
20:
end for
V/N TH1
Vt
Agent 1
Agent
Agent N
Action agent
Observation
Action of agent N
Observation
Environment
Figure 4: Schematic representation of FedRL, where N agents interact with homogeneous environ- ments.
we provide an example where three agents assess distinct policies within the same environment. In the traditional FedRL framework, agents exchange the evaluated value functions via a central server, leading unified consensus on value functions for three different policies. This enforced consensus on value functions, despite the diversity in policies, is not optimal. In another scenario depicted in Figure (bb] three agents each interact with their unique environments. The objective for each agent is to learn an optimal policy tailored to its specific environment. However, within the traditional FedRL framework, the central server mandates a uniform policy across all three agents, which clearly contradicts the intended goal of achieving environment-specnn optimization. This highlights the necessity for personalized decision decision-making, a feature that conventional FedRL frameworks do not accommodate.
18

---

Published as a conference paper at ICLR 2025
Environment
Enviromment
Enviromment N
Observation
Observation
Observation
Action of agent
Action agent
Action of agent
20
Companent with paramelers (9t,
Component parameters (9r, o
Component parameters (4, 0
G-1
1
Figure 5: Our proposed PFEDRL-REP framework where N agents independently interact with their own environments and take actions according to their individual RL component parameterized by + and 0: Agent locally update weight vector 0, while jointly updating the shared feature representation c through the central server. The update follows the Markovian trajectories.
Agann
Aggint
AppnT
(b) Agents learn optimal policies for heterogeneous environments.
(a) Agents evaluate difference policies in the same environment.
Figure 6: An illustrative example with three agents that demonstrutes the conventional frame- work fails to work.
Example of RL components that fit the proposed PFEDRL with shared representations. In the following, aim to showcase examples of RL components that are compatible with our proposed PFEDRL framework featuring shared representations. An illustrative example of this framework is presented in Figure 7
that both the DQN architecture in Figure 7 and the policy gradient (PG) approach in Figure 7b seamlessly integrate into our proposed framework. This integration is achieved by designating the parameters of the feature extraction network as the shared feature representation & and the parameters of the fully connected network, which either predict the Q-values or determine the policy, as the local weight vector 0. This arrangement underscores the adaptability of our framework to various RL methodologies, facilitating personalized learning while maintaining a common foundation of shared representations.
important
E PROOF OF LEMMAS IN SECTION 4.2
E.
PROOF LEMMA 4.6
Proof. Recall that for any observation X (s),,')), the function g(0, 4 X) defined in expressed as
g is
g(,,5, = (r(s, a) + (s') (s) - (s)
19

---

Published as a conference paper at ICLR 2025
Fectur Henwark
Pextare
Dossacted
Dorection
Connectei
Vakae
Festure Vecter
Teurure Vicao
DQN
Policy Network
(a) When DQN meets the proposed framework.
(b) When PG meets the proposed framework.
Figure 7: An illustrative example for the proposed framework. Notice that both the DQN in (a) and policy gradient (PG) in (b) can be fitted into the proposed framework by treating the parameters of the feature extraction network as the shared feature representation & and the parameter of the fully connected network which maps to the Q value of policy as the local weight vector 0.
and hence we have the following inequality for any parameter pairs (01. 21 and (02, with X = (s, a, s'),
(gg(0, $ X) g(02, 2. X)
[[(r(s,)) 1 (s)00)) - (r(s, 2 -2(s)02) 22(s)
(01)
( 1 क (s) (10 (s) .1)
+ +1( (1))) ()) (1 22()) 22())]
(o2) [1(701 (s')0 (s())) (, (s)02)) [s)]
+ |1 02 91 (s) (102())00 (2())2)) (s)
(nx)
(1 (1+7) |0 021 + (8))02 02) (8) (- 2 (8)
(04)
(1+ 7) |0 021 + 11 4 1 (s) (77) 2
+ +1( (8') (8)02)
(2))) - (11922'000 -06))))=0(s00) (s)
(ax)
(1+7) 021l (7 (5 (s)02))
(s)()) (s))
+ 11(> 1 (s' 02 (s)02) (7 (s')0 22))0 [2(s))
(as)
(1 + 7) |0: 02
+ ||| (7 (s')0 Φ (s) 02)
( 22())
+ (s') 2 12 + (1) 2 101
(1 + 7) 110 021l
+ (2 2+) 02ll
[ 221
(a7)
S Ly (1 - 021 | 2211)
(a1) is due to the fact that |xx + yll < |x|| + | Vx, E Rd. (az) holds due to ||x y/ <
|x[| ||y Vx,y E Rdd (aa) comes from the fact and | [+) < [ I Ivs. (a4) (ac)
holds for the same reason (a1) (az). The last inequalty (a7) comes from the fact that 0 is
bounded by norm and by setting Lg == max(1 + + (2 B).
E.2 PROOF OF LEMMA 4.7
Proof. Recall that for any observation X = (s,a,s') the function h(0, 2, defined in 10 is expressed as
h(0, x) = (r(s, 7 ())0)) AT,
and hence we have the following incquality for any parameter pairs (0,, 21) and (2)) with X (s, s'),
[h(0,, 21, h(0,, $2, X)II ((((s,)) + (s')0) @ (s)00) .T (r(s, a) + 72 22())22) 0311
20

---

Published as a conference paper at ICLR 2025
(6))
1
(1())01) 0 (n - (s))) 011
+ 12772(2 (5)) - (2))))) 10 ( ())0 02())02) 0
(4))
(s)00)) (7 (2))) |011
+ +11 - (s()))) 01 ( 22)))22) 0.
(b)
S
(1 + ++ 112 1| 221 + (1 22 (2()))1) 01 2 00 22())02) 110
(44)
x (1 + +
112
| 1 + 11 ()))) 0 ( (2 (b)))) 0
+ 2 (s())))) 1o ( (2())02) 05
(k)
3 (1 11)1001111
| 1 2
2
(s) (()) 110 021l
+ ( (s')0 |0 (2()))) (n 22 |022l
(4))
2 (1 + 7)110,11
||0 22 (1+
100 021l + 10 110, -2
(1 + 7) 10000111 112
| 221 + +1++) ||021)) 100 021l
(Z) (10 - -021 |1 221l).
(b2) is due the fact that ||x + yll < ||x|| + ll Vx,y E Rd, (by) holds due ||x yl S
|x]| -lyyll Vx, E Rd, (bs) comes from the fact and [1(s)]] < 1, [1))l I IYs. (ba) - (bb)
holds for the same reason as (b,) (bg). The last inequalty (by) comes from by setting Ln
max((1 + B², (2 + B).
,
E.3 PROOF OF LEMMA 4.8
Proof. Due to the norm-scale step (step 9) in Algorithm
2 we have
+ y (2)) S
max
|10 10-0'1 < 2B.
(19)
(1
(1)))))
Since the representation matrices $ and o are of unit-norm in each row, there exists positive constant Ly such that
|| (1'( 1/ (02) L 2211-
(20)
E.4 PROOF LEMMA 4.10
Proof. In the TD learning setting for our PFEDTD-REP, at time step k. the state of agent i is sf. and its value function can be denoted as V(s{) (st ) in a linear representation, where ((sf) is a feature vector and 0² is a weight vector. The goal of agent i is to minimize the following loss function for every st S:
(0(()))))) = ((s() - v(6)))
with ((+)) = r > ) being constant. Therefore, to update ((s) and 0, we just take the natural gradient descent. Specifically, we update 0 according (7) by taking a gradient descent step with respect to 0, with fixed $ Similarly, we update (s) according to (g) by taking gradient descent step with respect (s). with fixed 0.
Next, we show the convexity of the loss function L' ( 0¹) with respect to the feature representa- tion ((f.)) under fixed 0² Since the estimated value function is approximated as V(s))) = $ (( e where 6¹ is a fixed parameter. Taking the second-order derivative involve TT which is positive semi-definite matrix as long as 0³ y 0. Positive semi-definiteness of the Hessian implies convexity. Hence, L (*)()) (,))) 0') is convex (ssf) under fixed 0', This property holds vice versa, i.e., (((tt) 8() ()))))))))))))) is convex on 0 under fixed ((s)).
( ( ), 8¹) w.n.t. ((s will
21

---

Published as a conference paper at ICLR 2025
Recall that the optimal solution i and 0" is defined the set of possible values that make the expectation of stochastic gradient g and tends to be O, as defined (12). which is analogy to make the first-order gradient of loss function be 0 and achieve the local minima. The inequalities in 1emma|4.110)er denote that the updates made to the feature matrix for fixed 0 in the first equation and the parameters 0 for fixed + in the second equation directed towards reducing the deviation from the optimal solutions close to initial point. As we only care about the solution to make stochastic gradients be for fixed 0, the loss function L is convex w.r.t. L, the learning process of $ is guaranteed to move towards decreasing the difference from an optimal point. This also holds for the update of 0.
E.5 PROOF OF LEMMA 4.12
Proof. Under Lemma 4.6 we have
0g T X) ggy' + )
x)Il - L([e- v'( + + $- I). Vi [N].
(21)
Similarly, under Lemma 4.7
we have
[h(0, Φ X) h(y' (+)
X) < L( y'('')|
$ (]) Vi E [N].
(22)
Let 11- max maxx g(y"
+ X), maxx h(y (+) $ X)), then according to 21
we have
(g 3 - L ( (0''') + 0 $ + 1),
and
[h(0, 6)1| - L1 (1 y (0"")
+ $ 1).
Denote h² (0,06 X) as the j-th element of h(0, क X). Following Chen et al. (019)) we can show that 0 E Rdd $ E xd and E X,
| Φ X = x| - 1h + xX]II
ET 10 Xo = 1 0,, X)]|
10 Φ X)
< 2L, (110 - 1 ( Φ ||| o 61 + 1)
IE
No
22 ('()))
1 X+
j=1
h0 & X))
En,
2L @
y¹ [+)
|| X+ +1]
1)dCrrh.
where the last inequality holds due to Assumption
4.3 with constants C₁ > 0 and E (0, 1). To
guarantee 2/, (I|0 y' (0*)
+
क +111+1)CC1
8(1 y ( )| 1 $01 I
1), we have
log(1/6) (2L
TB v log(1/P1)
(23)
Using the same procedures we can show that
( कb X) x]
Eg क. X)]II 21, (1)0 y'(0") + 11 $ 1)dCapp
hence we have
log(1/6) + log(2L,
T8
(24)
log(1/ (2)
By setting т8 the largest value in
23 and
24
arrive at the final result in Lemma 4.12
22

---

Published as a conference paper at ICLR 2025
F PROOFS OF MAIN RESULTS
F.1 PROOF OF THEOREM 4.13
For notational simplicity, in the proofs, we use h(0 2, denote h(0 2t, (xi 1 and g( 1, p.) to denote g( or x In the following, we first focus on the update of the global representation 4 and characterize the drift of it.
F.1.1 DRIFT OF Q:
The drift of Φ given following lemma.
Lemma F.1. The drift between D441 and 2, is given by
B c ++1
111
E Φ x01 112
E
h 1+
2
+ 23,
(
h(O E+1
21))
N
Term 2
N
+ 23,
$
h( 1+1 p.) - h(0 ++))
(25)
N
Term 3
Proof. Based the update of q in 11
We have the following equation
[
EE[
111
E[
1 (+1
2( 2+11) E
0 1 1
1
[[
[2]
B[1 1 2 1 [+] Φ.)]
E[(
L,, Pet1 + k.)] - 2 0,+1
T.)]
B[(00+1 (+1
2 r+1 b.) + 2E 1 r+1
1, Φ
2 2 2,+1 x1)1
88 E
h(tt c+ Φ.)
23,
h(O +1 R.))
(26)
N2
which directly leads to
E
1 1-1 -1
32 N²
N
27
N
= =E Φ क-
11
Σ
h
1+
Φ,)
23,E
h
(0
2+1
+
(27)
Rearranging the last term yields the desired result.
In the following, separately bound Term to Term 3. We first bound Term as follows.
Lemma F22 For any 2 we have
4 N
Term 1 < 488 1
+11
E
0
(
+ 4,3,82
(28)
Proof. Note that
32
N
N Σ
N (2). Φ*) + >
Term = E N2
Σ
h 1:
2,)
h(p
(Z))
h(y ( "-)
[=1
1=1
23

---

Published as a conference paper at ICLR 2025
triangle inequality
2
Σ h 1+1 2,)
Σ h(yy(()) (R.), (*)
N2
Lipschitz property of h
2 N
E
Σ h ( 25)
(a)) 2
2NV
[0 -2 (
2N 11 1²²
N2
2
+ N22
Σ
h(y (o.),
Σ h(y $ )" x
h(y 4 ),
4.3
2e1121
E
0 10 - (2)))
N
4
N
N
+
h(y (*))
N2
Lipschitz h,
4
N
+
Σ
h o
N2
(02)
1.
N
I 48
2E -1,11²++ +
E
10 10
y
N
27
4,6
+ L4 110 - $0113
+
N2
Σ
h(
-
h( (
x²)
i=1
1=1
(0g)
4.3
NV
(1. + -+)
1111
E
10
v0 + 4,378
N
where the (a1) is due to 1 12 < functions and y² and (aa) holds based on the mixing time property in Definition 4.3.
Σ
||x]|] (az) is due to the Lipschitz property of
;=
Next, we bound Term 2 in the following lemma.
Lemma F.3. We have
3 N
N y y²
Term 2 B.(L/ar 2w)
2.1131
+
E
10
(29)
Proof. We have
1
Term 23,E
(
1
h(0 11 21))
N
N
1
= 25,E
(
$
h
(F.),
+1))
N
24

---

Published as a conference paper at ICLR 2025
23,E
2 t
h (2t), 2t) hi 211 2())
Lipachatz of h
< 2
(Φ
21 N
h (o)) 2())
LB
(Φ
22 N
(y ( 0(+1))
(b))
2
(Φ
h(y ( 2.))
+ B,L/ -1
$111"1
N
==1
Becr, I
E
0 t+1 (..)))1
N
(b2)
2.
$
h (+), 1())
/T 0.113)
3,arL
+
10 1/ (0,)
N
BlarL E N
S Bi(L/a - 2wd 2
10 i+1 - (8,))1
where (b1) holds because y S
|||x||2
1/
43 0, (b2) is due Ei= x
N WE == ||x and the last inequality is due to Assumption 4.10
Next, we bound Term 3 in the following lemmas.
Lemma E4. For all 2 T have
Term S (73t/at + 2. 6,, B : $
+ (63//a, + 63 (1 L3) + 43, BL - $
+ 163,a, + 63,a,0
EE
Σ 1 - -0++1l?
+ +18,a002.
(30)
N
Proof. We first decompose Term as follows
Term 3 = 2
$
h(0 1+1 Φ.) - -h(0 (1 Ø.)1
2=1
( 2,,
h(0 10 (+ Φ.) h(0141, R1))
N
+ 23,E
h(0 11, 2.) -h(0 Φ.))
N
2
h(0
1+ 2.)
h 2())
N
1=1
N
N
+ 2,,EE
h(0
1+1 +
Σ
h F,r))
N
N
im
iml
25

---

Published as a conference paper at ICLR 2025
N
N
+
2
( (
h(O
p(-r)
h(O
21 $
N
N
N Σ
+ 2
(
h
1+1 D(-r)
h(0 21)
N
N
1=1
t=
Next, we bound C1
N
C = 23,E
(
h(0
22) -h(0 1, 21))
N
1=1
N
$ 33/0,B111
1 8 2.--1191 + 3,
h(O ++ 21) - -h(0)+1. 21) + (2))
N
N
S 33//aBB1l0..
Φ - + 2B
Σ
h0
++1 p:)
+ 2
h(y (Ø"), )1)")
h(0*
2+1 2.)
N
23,00 E N²
-B/o,BB100
$ $--11171
Σ
h(O h(0++1+ 2+1 3
1
+ 28,a,B
(0",,"))
h( ++1 2.)
N
8
$
---11
+ 8/ [I2 + I,*)E|
g
$
88 T.2
E
Σ
10 (+1 (0,))1
+ 8,,0000
N
1
+ 23,,,
Σ
h(y'((())) (6) (*) h(0 2
N
Lipschitz
$ 3./o, Φ. o 11 + 88 (I. I.a T.²)B[[15 [| $
8, +
N
E
10 (+ (0)211 Σ
+ 8,,00002
N
+ T.
2( ②.) + +200
4+ 5/(00 (Φ)))
N
$ 3//
| ---1"
+ 88,
(L2 L |B[ 1
80,0, +
N
E
1 F+1 (1)
+ 88,0,02
N
4 N
N
+ 4,3,0,
Φ 11 +
E
110 :+ 0'
3./o
Φ.
88
I,
-
26

---

Published as a conference paper at ICLR 2025
L2
N
+
E
Σ
10
(8,
+ 88,00002 + 1. - -.112
N
4
L2
N
+
E
Σ
(+1 (0,,) + yy 5 00
N
B. 0. Φ |12] + 88,00
[1?
I, | $.1131
L2
N
+
E
10 [+1
y (B:)11
+ 8,,,00022
N
83,a L2
N
+ 143,a, 1
0.112] +
E
110
3 (B,)))
N
+ +83, LB 1 - - $11121
/,/a,BI100 2 11 + 4,,a, 23 + 4L 2 -
16
E
Σ 10'+1 9 (1.))
88,000²
where the last inequality is due to the Lipschitz of the function y Next, we bound C2 as follows.
N
N h(6
C2 = 2
h(O 141
D.)
1+1 $(-r))
N
1=1
N
N
$511 + 3rO F
h( +11
②.)
h
10 2 2(-r)
N
N
Lipschitz of h
< B1/a, 8 - 65113 + 3,a, L || 0 $
B|/a, p,- 4 + 1 4 $11 + 3:0; L. || 0 $ $
< 22
$ a
8 111 + 3, L B1 o 11
(2.r/aa + B,a,
11 / Ø: $ $1
Similarly, C4 is bounded exactly same as Cz,
C < (23:/ar + S L)) 1 23c aaE [ $-11
Next, we bound C3 follows.
N R.-r)
N N
C3 = 23,
h(0
10 o Φ--r))
N
N
27
& / B
Φ $11 +
BeCs N2
E
Σ hee 1+1
o Φ--))
Σ
h( 1+1 o F--7)
Detinition 4.3
Br/anEill p.-- 11
(N6|10 $61 N8 + Σ
N
27
+ 3,0 N2
0 0 - y/ (8*))]
1=1
< /.//a p,--
$" $11
+ 8
|
$-1²2
+ 0,,
8
N
+ N
Σ
0 1+1 - 3 (²(0))11
27

---

Published as a conference paper at ICLR 2025
=
/B/a,Bl1
2-1 + 3, 'E | - -011² + +
0/ + N
N
1i+1 - (2 + ) - -'(0))
S
$11
+ 3,000, 2 |0, 1 - Φ-1² + 3
63 +
E
0 h+1
1 (0:))1
+63
| - $11
N
(23:/a + 63,a,0
E
2e11 + (28r/ar 63
+ +, 1228² ] 0011
)
63,0002 + 33,028² +
N
E
Σ
00 s (
N
where the last inequality comes from B|I Hence, we can write Term 3 as follows
111 s 2 2 2.11 + | -11
Term Ci + C2 + C3 + C4
$ B
11
+ 43,a, 4L ||| - 2c11
1 163,a,
N
+
E
110: y' 21)
+ 83,
N
11=
+ (23r/ar + 3, L2) D.-r $ $
23:/az a(I -
+ (23//a: 6,,a,
11
+ (29//r. 63,,0,0² 8 63tar 1.22 1 $ 111
+ +3 + 6,,a,82
Σ
0 - (0,)))
+ +223//a, + 3, L E ,- 1 23,/a, E|| Φ x "111
< (78./a, + 28,a, + 63, 8² /B 0. $ 111
+ (63//or + 8 (1 + 1 + 43,0r 1,2 1 + A1.2) E 171
16 + 63,
E
1lo (Φ.) 0 01 1
+11,,0,0²
N
which completes the proof.
To bound Term 3, we need to bound [[ 2.-r
which is shown in the following lemma.
Lemma F.5. We have VI. > 2r
El| の 1 < 4r 2 /// E E| o"
1 8,3,0272.
(31)
Proof.
The proof is similar that of Lemma 3 in Dal Fabbro al. (2023). Starting with
32 2
② ++1
2
21 +
Σ
h 2+1
2.)
2 Φ f
h( 10 2.))
(3:00 + {B
S (1 + 8,/a0)) x 111122 2
h
0 + +
N²
23,00
< (1 +
0000 1+
1
(32)
N2
28

---

Published as a conference paper at ICLR 2025
where the first inequality holds due to 2x"y I +||x|| By > O, and the second inequality holds since $ea0 2 S. We then have the following inequality according to Lemma F.2
E ||
L++1
1| S (1 + Be/an 8 2 + 12)) ~| $
L3
2 E
N
0 - (2,)
+ 83,0000
N
$ (1 + B /o + 8 (1 + L )) -1 - -1,1' + 88 + +82).
(33)
By letting og 1 21. 321/2(1111) 1.)) we have 81/a0 2 83,00 + L2), and hence
E [[ $ r+1
||! S (1 206 on) [1]Ø*
+.11" - 8:000 (L3B²
(34)
Therefore, for all t' such that t T s st.
T-1
[|| 0- T [1]] < (1 + 286/on) TE|l 0- Φ 83000 (T, + 82)
Σ
(1 /00)
t=0
(35)
Using the fact that (1 (1+ 2)) c Dal Fabbro al. 2023 if we let Bo/ao S J. we have
8,
(1+ 230/a0)* < (1 200/o0)
0.25 2,
and
(1 + 32 2T.
8=0
Hence, we have
0-
2. ||
< 2E| 0- -, + 168 + +
Since | - = T - @0et1
Φe|12 = T 8² No El= E/ h; &e)1 when
t 2 27, we have l 2 and thus
B|| ; 1 r-- [|
g2
t-1
N
I T N2
O2 22)1
t=
i=1
ST
Σ
((45) 22 + 10
t=
< 4L6 + (2 -0,-- -0- + 43 4338872
< 72 / 4. - ., 11 115=111++342128+44444
< 2T / E| x- + 1 2 1 || 4 L2 4888222
Since 2T / /0/ < 1/2 when Bo/00 S 57 we have
E + 1
< 81 888854..
This completes the proof.
Lemma E6. Term 3 is bounded follows
Termy & (78./o, + + (3 -| +
-
+ 108//a++633.rm + (1 + L²) + 43,0r, L2 + 4L |||| 8 - 11
163,a,L2 + ,3
N
E
Σ
| 0 1141l2 2
+ +11,00002
N
29

---

Published as a conference paper at ICLR 2025
Proof. Substituting the bound of _ - in 31 into Term in Lemma F.4 yield the final results.
Term in Lemma R3 and Term in Lemma R6 we have the following lemma characterize the drift between Φ.+1 and +,.
Provided Term Lemma F22
Lemma E.7. For t 2 2T, the following holds
+++1/1 |] x |
(1 + 43/ (L2 + ( (
+ (63:/a: 6
- 4L²)) Bi(L/ac ( 2w)) $-1
L3 + +6 E
10 0z+11 11
N
+ + 63,0 (3
8 + 4,8 + 11
Proof. Substituting Term,, and Termy back into Lemma
we have
@
1
4
N
S O
1 4,
1.1121 +
Σ
10+11 (8))) 4
BrarL N
N
+ BelB/o: 2w) E|| $.11 +
|10 1+1 y¹ (
+ (73:/a, + 23,a + 63 (47 /
$ 111 B2
(63./a, + 63,00, ( +1+1000 1.3) + 4,,
41.3 ( - $11131
163,o, + 6
E
0 1.t1l1 + +11,,0,,0
(1+ 48 1.4) (73,/a, 23, + 63, 82 )4
+ (6.3//r + (3) + 4I²) + BL( - B[)) $1
4.8 + L 163 +
N
0 - - Σ
N
(73r/a, + 23,,, 63:a
8800773 + +3300² 113,0,0
This completes the proof.
F.1.2 DRIFT OF 0 Vi. Next, we characterize the drift between 0 and 0 Lemma F.8. The drift between 0:11 and 0; Vi is given by
B +1 (1 = E
0 (0 at
g 1 2.)
+ [1 (110(00,1 (w.1) - y (
Terms
Terma
K
+ 2B
0: (r( +
0
1-- 1.). v'(3,)
k=
Terms
(36)
30

---

Published as a conference paper at ICLR 2025
Proof. According to the update of 0 in
7 have
B| y'(+)) E
0: y2 (
04
1,6-1 21) + s (1(111) v0))))
k=1
E
0: y y/ (ic-1) at
g(0: i Lc)
E
11 (f(-1) y @+))
Terms
Terma
+ 2F
0: v2 (-(-)) 04 g(0 k- 10).1(001))) -y((0,)
Terms
(37)
where the second inequality holds due to ||x + y/1² ||x| + +
We next analyze each term in (37). First, bound Term in the following lemma. Lemma F.9. With 2 we have Term 4 bounded as
Term < (1 23 - Ka,Kw)B |10 y' v'(
111
+ (22
1 || 1. 2.-1
+ 6 (1+ B²) +2) + 21 K2 B² / B + at. a3k 22 +3K28).
(38)
Proof. According to the definition of Term 4, we have
K
Term = E
0 (' + a
Σ
D.)
K
E |10 1 (Ø (0.-1)11²)
+ aZ
Σ
g 0
1.k
K
+ 20, 0 y'( 1), Σ
g k Φ
k=1
E
[10:
2d (0, (0.1))
+ 20,
(0: 0i- 1 ($,1) 1),
(0)
1.k 21)
k=1
K
K
Σg
0 1.
2 Σe
(
2 +
( 1.k-1
Σ
g y
+
k=
k=
K
<B
0
+ 2
Σ
0
+.)
1.k-1
1
A=
k=
Mixing property Delinition
K
K
27
+ 20
Σ
gly ( 2
k
31

---

Published as a conference paper at ICLR 2025
( 0: y p 1),
K
+ 20
y'
g 1 B2))
I E 0 y (0.-1)) (0,-1)]1²
+ 6a K
+
( 0; y2( 1 1),
20
8 1 $1)
Torm 4.1
(
g( k.k $ |
+ 2
0
0
k-
+
Σ k=
(39)
Term 4.2
where the first inequality holds due to the fact that ||x + yl1 S 2 2|xx|| + +2|ly|l and the second
inequality is due to the mixing time property of function g as in Definition 4.3.
Next, we bound Term 4.1
K
Term 4.1 = 2
0
g
k-1
$
k=1
0
1),
B( $ Φ
K
K
+ 20,
0 8 (
1)
g( 1: L1)
Σ
g D 2--1)
1 --2a, KWE
[100: V @
K
20,1F
0: v @ 1),
Σ
( 1,k-1 b
Σ ( 1
k=1
20 KWIE [110:
y' ( (1-1))1
1 1 E 000 y
(
az 1
Σ g(0
1 2.)
g
(40)
In partieular, bound E
D g(0 1- 8.)
ΣK (00) $
||]
K
E
0 f-k-1
g
A=1
1=1
K
< 22 2
| 1 11
+ 2L
0 0,
K
< 2L'
|0 Φ 1-11 1
+ 2I
10 0.11
k=1
< 2L
$-11
12
+ +
(41)
Substituting 41 back into 40 we have Term 4.1 bounded as
Term 4.1 -2a, KwE |||0 y ( Br 1/oBB
|10: -'(0
a2// 2E | & 2-11?
k2L2
(42)
32

---

Published as a conference paper at ICLR 2025
We next bound Term 4.2
( 0; y' ())1,,
8 00  11)1
Term 4.2 = 2a,B
1), Σ g( 1, 0.)
Σ
k=1
- Bt-1/00
|10 y ( (0.-
+ 0 +3K28 +3K 8 $-1
- 3
|10: 1 p 11
+ og (3K +3
6K a2 / 13 .- - $ + 6K 6k / 1
$ $-111 (43)
Substituting Term 4.1 and Term 4.2 back into
39 we get the final result
Term 4 < E [100 -y' (0.-1) [|] 6028"k*
K3 $0 Φ T1]
BL) -20
+ Term 4.1 + Term 4.2
E |10: y (0.-1)
|. 1-1
16,20K2211 B²) + 2o
20, KwlF [110]
y ④ - )|?
+ 1-.11//0p1
|10: - -
+ 12//8--122
| .--1 |
x222K3B2)
91-1/00 E
|10: y (9.-1) 1|1
+ (32
+ 6K* 6 a. 3
2 1-1 - 6 & / 12-1
$ (1+ (1+23/2//a 20, KW)
|00 - v (0-1))
+
+
02 |E|
c 1 1-11
+ +
+ (3K
(44)
This completes the proof.
Next, we bound Term 5 in the following lemma.
Lemma F.10. With 2 T, we have Term 5 bounded as
4,,_
L
Term 5 < ( L E
E
10:
(0)
4L22
(45)
Proof. We have
Term 5 E
[1 ①
(0.-)) 1) y (
L'E [1 o
$ 12]
1338-11
N
Σ iml
h $1-1)
N²
4,82 L
N
< <437_1L4 + + E - 2 1
E
10
+ + 32 18
N
(46)
where the last inequality holds due to Lemma F.2
Next, we bound Term 6 in the following lemma.
33

---

Published as a conference paper at ICLR 2025
Lemma F.11. We have Term 6 bounded
Term
S B 4 + at/3 Term
(47)
Proof.
(
Term 6
2E
0
+
04
g( 1,k-1 (.).)'(()))) (0 v²(())
k=
0; y¹ 1 1- 1) + 04
8(0 2-1
+ 04/ 3 EE [1y ( (o.-1) y (Φ:)
Term
4
(48)
Providing Term in Lemma F9 Term in Lemma F11 and Term Lemma F.11 we have the following result.
Lemma F.12. For & 2 the following holds
E El
y(( [2]
(1 30-1/0))
(1+23/-1/0) - k2)
+ (12a x2 21 +2800 482_,422
48% L + +111 (1+ at (1+aa/BL-1)
(1
N
E [100 y' (8 (0)111
+ (1+ (1+B (1+B1-1/a1) (1202
+ 12229998-4200//33353333332221111 1111
+ (1+ 100)
+ (1+ (1++ (120 +43
+ 6a?52 + 2o + - +32283)
+ - (1-a///41) 41. 12
(49)
Proof. According 336.. have
E 0 +1 y¹ (0)) = Terma + Terms + Terms
Lemmn MI
(1 81-1/aa)Teem, + (1 0z/ Terms
< (1+ + (1+B/11/a))
(1+28--1/a: - -ay 0 s (0
+ +
+ /12033/4-21//0/551003388//////
o -
+ (2)
/3/ + B3K +3
34

---

Published as a conference paper at ICLR 2025
+ (1 + a11B1-1)
1 + [.)
-12-
4,6% L
N
E
100 1/ 0/( +
N
iill
Lemma E.0
(1 11-11/aa))
(1+23,-1/44) 20, Kw,Kw) Kw)B
|10: y0 (0 (0.11))
+ (12a70 + 6K 6K /
$11
+ 2
4,3 2
(40 (I.² B|| 3 25
12
E
0. 5'
+ +488_
N
+ (60 + B²)
20 K2
Br
+332283)
+ +1+ (1+0t/BB-1)
( + E : 2-11
48 L
N
E
110 y 01 + Σ
N
(1 81-1/ar)
(1 (1+ (1+23(-1/a4 - -2a(kk
1,3
+ (1 (1+aa/Br-1) 1
+ +122200K2 +2
-E =1100 y (0)1))
+ (1+ B0-1/a)) (1203882++6(886)//441))
+ + |,0 L11)
+ -0++00//--)) 1.0)
+ +1+ 1 (120 2
+ 60,28 22 (1 +
"" B3
+ (1 + at / 11- 4L 3,,3,
(50)
This completes the proof.
F.1.3 FINAL STEP OF PROOF FOR THEOREM 4.13
Now, we are ready to proof the desired result in Theorem 4.13 According to the definition of Lyapunov function in
14 We have
Be
x = +1 +1112+
0 1/ (01+1)12
1++1
N
1=1
(1 + (T2 =.A) (73./o, 23. + 6,3.0
) 82/03
35

---

Published as a conference paper at ICLR 2025
66/ 4 48fI² + 3, + 16 63, E
+ 4L²)) + (L(aa 2w)) $-1""
N
10 1+1 (21)2
+ (73e/ar 23,a, L² 63,0,, + 8,3882 + 43,88 11,,,a,
(1+ [1 (1+8//at41) (1+28 20 Kw)
+
0++1
+ (12a 8² + +2L 0 /3 6k / 43%
(1 a 1/+1/3)) 43f
N
N
Σ -
- 2
N
+
(1 +1+ Be/oet1)
(120
+ 11 X22 + a 62 / (B)) + +())
+ (1 + a1+11 + L9) 1B $.11
+ (1 / 1 82 K2 + 22L2 6 /
+ , K2 + 33- +3x26)
+ (1 (1+ (1+0++1/8)) 3282
(51)
To simplify the notations, we define
Di = (4) + 0a0p://+022200000000000001000013
+ (63./a4 + 6, + 4I2)) 3(L/ai)
+ 3x
(1 (1+3//a++1)
+ + (1224
0/++
+
+ 2L -
+ (1 (1+14433))43 + L8)
(52)
and
D2 == -/8 a af 16a 1 68
(28r/at+1) + (12a K2 21 2 +6K / 43,1.²
4,3
(1 + 04+1/3.)
Be/ou41 (1+28r/0041 20.411uu Kw)
43,L2 N
48
(1202 x² 2T?
6K 6K38
+11 18,)
(53)
N
36

---

Published as a conference paper at ICLR 2025
Since D1 is of higher orders of ((3)) and D2 is of higher order of o(ar+1), we can let D1 S wBt and D2 I Kwa++1- Therefore, we have
M((0++2) 11)) $ wB, (10 +11, 1,, +.)
1
1 4+1 - [()
- / + 4 Bi1 4 111+
E
Σ
N
+ +40++1
+ B²) + - B²
3K3 8²
+
B²) + 8a
S 51 (1
N Σ
+ +11447K23300 - 4L2 /L2//N 3,00. B11 $11 +
1
E
00 (0)))
+
+2 3K
2L² + 8,, (54)
where the first inequality holds by omitting the higher order of learning rates, and the second in- equality holds due to the decreasing learning rates of at.
We now set the proper decaying learning rates. Let a+ o + 2) and Br Bo/ + 2). We then have
(t + 2)² (1 - w = + ( (t+ - (t + 1)²
(55)
if wB, < 2. In addition, we have the following inequalities
(t ()+ 3 ag + 2)
(1+ 2)² a at 2)²
Hence, multiplying both sides with (t + 2)2, we have
(I+ ()+2)) 0 = (1 + 1) ( R.)
B[[ 1
E
0 i+1
y
N
+ (4ag3)KK (1+ B²) + B²) + B² +K 21 K2 B
San
(t+
Summing the above equation from = O..... T, we have
.
(T + 2 P++1) S M(( .n)
N
+ (14472 + 41. /N) + 2/3/3
E - 1
100--11133
E
Σ
10
v (0)))
N
+ ( (1 + B²) + L²B3) + BK
+
K2 B²) + Sa 2) (T + 2)2//3
Dividing both sides by (T + 2)², we have
M( M( )o)
(T + 2)
+ (14472 + 4L / NN)on 2 El
2/3
111
E N
1
Σ 00 (0)))
(40030 K2 235² (1+ B²) L2 + 20 + 33222 21 R2 B + 8 2 + 2) 2/3
This completes the proof.
F.2 PROOF OF COROLLARY 4.15
If ao = 3o - K -- we have
M((0 R.41) 5 (
2 N2/3(T + 2/2/3
K2N5/3(T + 2/233 K2N2/3(T + 2/2/3
which is dominated by 0
N37
if72 > N.
37

---

Published as a conference paper at ICLR 2025
Table Parameter setting
Parameter
Description
Input size
4
Hidden size
128 x 128 x 128
Output size
2
Activation function
ReLu Iu
Number of episodes
500
Batch size
64
Discount factor
0.98
greedy parameter
0.01
Target update
30
Buffer size
10000
Minimal size
500
Learning rate
0.002, decays every 100 episodes
Feepet coor FRI vom PerDONis Fedleyn--iiat
x110
Retum
Retum
200
500
300
400
300
Episode
0.38
0.42
0.46
x110
Retum
Retum
1.00
200
300
500
100
400
500
200
300
Episode
Epicode
E Epbodde 0.62
0.5
0.58
X110f
2.20?
x110
Return
Return
1.00
200
300
400
500
100
200
300
400
500
1.00
200
300
400
500
Epicode
Epicode
0.66
0.7
0.74
Figure 8: Comparison control by DQN, FedDQN and PFEDDQN-REP in Cartpole Environments.
G ADDITIONAL EXPERIMENT DETAILS
Compute resources. The experiments are performed on a computer with Intel 14900k CPU with 48GB of RAM. No GPU involved.
PFED in the CartPole environment. We evaluate the performance PFEDDQN-REE REP in modified CartPole environment (Brockman al. (016). Similar to Jin et al. (2022), change the length of pole to create different environments. Specifically, we consider 10 agents with varying pole length from 0.38 to 0.74 with step size of 0.04. We compare PFEDDQN -REP with (i) conventional DQN that each agent learns its own environment independently; and (ii) federated version DQN (FedDQN) that allows all agents collaboratively learn single policy (without per- sonalization). We randomly choose one agent and present its performance in Figure 3 (a)) The results of the other agents are presented Figure 8 Again, we observe that our PFEDDQN-REP achieves the maximized return much faster than the conventional DQN due to leveraging shared representations among agents; and obtains larger reward than FedDQN, thanks to our personalized policy. We further evaluate the effectiveness of shared representation learned by PFEDDQN- -EP when generalizes to a new agent. As shown in Figure (b),,,b), our PFEDDQN-RES generalizes quickly to the environment. Detailed parameter settings can be found in Table
38

---

Published as a conference paper at ICLR 2025
PFedDMMRRy -ON
FedDOM
Fee0 FAL PeOOkAug FesteynG-mmaa
Retum
Retum
-10
10
100
80
60
pissode
picodd -0.12
-0.24
-0.18
x100"
x16""
x.10"
Return
Return
-10
-10
40
60
100
20
40
80
100
20
60
30
100
pisode
Episode
E pisode
0.06
x110
x100
x110>
Retu
Retum
-10
-10
-10
100
20
50
Episode
0.12
0.18
0.24
Figure 9: Comparison of control by DQN, FedDQN and PFEDDON- in Acrobot Environments.
PFEDDQN-RE: in Acrobot environment. We further evaluate FEDDQN-REP in modified
Acrobot environment (Brockman et al. (016). The pole length is adjusted with [-0.3, 0.3] with step size of 0.06, and the pole two benchmarks are compared as in Figure Bttop). The parameter setting remains the same except number of episodes decreases to 100. Similar observations can be made from Figure ((ootomm) and Figure 9 as those for the Cartpole enviroments.
with be adjusted accordingly (Jin et al. 2022). The same
G.1 MORE COMPLEX ENVIRONMENT: HOPPER
We consider another environment, Hopper from Gym, whose state and action space are both con- tinuous. To induce heterogeneity within between the agents' environments, we vary the length of legs be 0.02 + 0.001 i, where i is the i-th agent, while keeping other parameters (such healthy reward, forward reward, and control cost (the /2 cost function to penalize large actions), the same. We increase the number of agents 20, and plot the return with respect the number of frames. In addition to training, we also generate a new sampled transition to validate the algorithms' ability to generalize.
In order to fit the algorithm to the continuous setting, we modified the proposed algorithm to DDPG-based algorithm, similar to the DQN-related benchmarks. For FedQ-K, LFRL and FedAsynQ-ImAvg we discretize the state and action spaces. Similar to Cartpole and Acrobot envi- ronments, our proposed PFedDDPG-Rep achieves the best reward and generalizes to new environ- ments quickly, shown in Figure I0
G.2 VERIFYING THE LINEAR SPEEDUP RESULT
We now verify the main theoretical result empirically. the personalized setting, verifying this re- sult is not as straightforward as in the non-personalized setting because Theorem 4.13 Corollary 4.15 when parameters defined across environments (e.g.. Ts. C, and others) remain constant as the number of agents (environments) increases.
To properly address this issue, we design an experiment that duplicates 2 initial environments with pole lengths 0.36 and 0.42. We duplicate these two environments with 2, 3, 4, and 5 times, thereby obtaining situations with N = 2, 4, 6, 10. Because of this duplication, we know that the across- environment parameters (e.g., Ta. C, and others) remain constant.
39

---

Published as a conference paper at ICLR 2025
PFedDDP Re
DDPG
FedDDPG FedAsynO-lmAvg
FedQ-K
LFRL
PerDDPGAvg
x10
100
0.6
0.6
Return
0.4
Return
0.4
0.2
0.2
500
1000
1500
2000
500
1000
1500
2000
Frame
Frame
(a) Average return for 20 agents
(b) Generalization parameter 0.05
Figure 10: Hopper environment.
As shown in Figure [1 ship in the convergence time. course, note that in practice, there is certain amount of unavoidable overhead in the experiments, which means the speedup will not be exactly efficient predicted by theory.
increase the number of agents,
approximate linear relation-
140
FedDQN-Rep
Converge time(s)
120
100
80
60
40
2
4
8
10
Number of agents
Figure 11: Linear speedup for cartpole with duplicates of environments.
G.3 COMPUTATION AND WORST CASE PERSONALIZATIOO ERROR TRADE OFF
We show another experiment that quantifies the tradeoff between computation and personal- ization quality, which we define as the worst case personalization error across agents:
E IE[N]
(00 (())))(s)) Vi. (s)
(56)
The intuition behind this metric is that if all agents achieve good estimation error. then this metric is small (meaning we have personalized well), but if some agents perform poorly while others perform well, then this metric will be large (detecting that we did not personalize well).
We vary the number of agents from to 10 and examine naive DQN that runs independently on each environment, FedDQN (no personalization), and PFedDQN-Rep (our approach). In the left panel of Figure [2] we show the computational resources needed to run each algorithm, while right panel shows the personalization quality. We notice that for naive DQN, we can achieve no personalization error at the cost of high computation. At the other end of the spectrum, FedDQN leverages parallelization and reduces the computation, but has high personalization error. Finally, our algorithm. PFedDQN-Rey achieves the best of both worlds: low computation, while attaining low personalization emror.
40

---

Published as a conference paper at ICLR 2025
DQN
PFedDQNRRep
-FedDQN
4000
60
Time to convergence(s)
personalization error
3000
Worst case
40
2000
20
1000
2
4
6
8
10
10
Number of agents
Number of agents
Figure 12: Computation versus worst case personalization error trade off.
G.4 THE EFFECT OF ENVIRONMENT DISCREPANCY ON PERSONALIZATION ERROR
Recall that in the previous Hopper experiment, we vary the length of legs to be 0.02 0.001 where i is the i-th agent and 0.001 10 0.01 is the maximum pole length discrepancy between envirouments. In this section, we vary the maximum pole length discrepancy between (all 10 environments are identical) to 0.04 (the environments have substantial differences).
We compare the performance of the three algorithms that include personalization, PerDQNAvg. FedAsymQ-lmAAg, and PFedDQN-Rep (ours). The results are in Figure 13
We notice that as the discrepancy increases, all algorithm encounter degradations personalization quality (measured in terms of the worst case personalization error defined in (56), but our proposed algorithm achieves the least degradation.
60
PFadDON-Rop
personalization error
50
PerDQNAig
FodAsynQ-lmAAg
Worst case
40
30
0
0.01
0.02
0.03
0.04
Pole length discrepancy
Figure 13: Worst case personalization error with varying pole length discrepancy across environ- ments.
G.5 ANOTHER LOOK AT PERSONALIZATION
Here, we give another look at how personalization is achieved by PFedRL-Rep. We first compute the cosine similarity matrix of transition probabilities between pairs of agents. This represents the similarity of environments between any two agents. After the algorithm converges, we compute the cosine similarity matrix for the policy layer (last layer) of the neural network. This represents the similarity of the learned policy between any two agents.
In Figures 14 and 15 we observe that all agents reach their unique personalization. and the cosine similarity of personalization layer shows close distribution as the similarity of transition probability. This essentially means for environments of similar transition probability matrices. their personaliza- tion layer will reach similar stage. Since the agents share the representation layer, the heatmap stays identical.
41

---

Published as a conference paper at ICLR 2025
030 0,31
6:5 C20
334 335
E50
0:11
FMM PM
1.10
6.7
6.79
6.79
0.2
Agent
Agent
Agent
Transition probability heatmap
(b) Personalization layer heatmap.
(c) Representation layer heatmap.
Figure 14: Heatmap of Cartpole environment.
HULL
021
627
Agent
Agent
Agent
(a) Transition prohability heatmap (b) Personalization layer heatmap.
(c) Representation layer heatmap.
Figure 15: Heatmap of Acrobot environment.
G.6 MORE STATISTICS OF THE EMPIRICAL RESULTS
We report the return average, variance average, return median and total running time for environ- ments for Cartpole and Acrobot environments. Among all algorithms, our PFeDDQN-Rep achieves best return average and median, with top variance and running time, as summarized Tables and
We also provide zoom-in in shortened plot for both environments to show the quick adaptation speed when sharing representations as in Figure [6
Table 3: Statistics for Cartpole environment.
Algorithm
Return average
Variance average
Return median
Total running time(s)
PFedDQN-
143
43
154
466
DQN
135
54
127
3840
FedDON
101
67
88
387
FedQ-K
112
34
107
490
LFRL
117
47
99
434
PerDONAvg
127
48
131
520
FedAsynQ-Iman
119
51
117
501
42

---

Published as a conference paper at ICLR 2025
Table 4: Statistics for Acrobot environment.
Algorithm
Return average
Variance average
Return median
Total running time(s)
PFedDQN-Rep
-42
37
-29
618
DQN
-63
67
-57
5854
FedDQN
-714
162
-625
571
FedQ-K
-213
41
-202
621
LFRL
-207
58
-194
676
PerDQNAvg
-295
64
-277
602
FedAsynQ-ImAV
-191
36
-186
664
-FacDONN-R Rep
DOM
Tedt
TradDO
CaN DOA
PeOCRANg
PargOMoog
x10
Retm
Retum
10.s
50
200
Epsode
length
pole length
length
(b)
(a) Cartpole environment.
(b) Acrobot environment.
Figure 16: Shortened plot for cartpole and acrobot environment.
43


Reviews:
[{"review_id": "EtgAr3rwcV", "reviewer": "Reviewer_Um15", "summary": "This work propose a Personalized FedRL approach (similar to PFL but with RL) allowing local/per-agent learnable parameters for use in heterogeneous FedRL settings with convergence results. The authors find a linear relationship between the number of agent and the convergence timestep.", "questions": "Is the problem definition of PFedRL includes \"shared common structure\", how does it affects learning? or anything else in the theory? In theory they can be completely different problem without any shared structure and the results still stand?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- The theoretical analysis have a proof sketch and the assumptions are clearly stated.\n- I think the two timescale approximation result is novel\n- Interesting results from cliff-walking and cartpole", "weaknesses": "- Although the setup hold promise, evaluation is rather on the simple side. I consider it understandable for now since the main focus in on theory.\n- If I understand correctly, the linear speed up is not a particular exciting result, since sample collection in unit time grows linearly with N.\n- There is no comparison with other PFL methods, or parameter-sharing MARL methods.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "PzG3h427Y8", "reviewer": "Reviewer_LerK", "summary": "This paper introduces a personalized federated reinforcement learning (FedRL) framework, PFEDRL-REP, that incorporates shared representations to improve learning in heterogeneous environments. PFEDRL-REP collaboratively learns a shared feature representation among agents while maintaining agent-specific weight vectors for personalization. The authors analyze PFEDTD-REP, a variant using temporal difference learning, and prove it achieves a linear convergence speedup in terms of the number of agents, demonstrating scalability benefits. Experiments in both policy evaluation and control settings show that PFEDTD-REP enhances convergence and generalization in heterogeneous environments compared to non-personalized FedRL methods.", "questions": "1. Is PFEDRL-REP universally applicable across diverse heterogeneous federated RL problems? What if a shared feature representation could not represent the similarity between different environments and the heterogeneity could not be distinguished by the agent-specific weight vector?\n2. Does PFEDDQN-REP maintain strong performance on more complex RL tasks, such as those with sparse rewards or high-dimensional state spaces? Additional experimental results on more challenging environments would provide valuable insights.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "1. PFEDRL-REP is an innovative approach to FedRL, addressing a major challenge in heterogeneous environments by introducing a shared representation while allowing for agent-level personalization.\n2. The paper provides a rigorous theoretical foundation, including proofs of convergence speedup under Markovian noise using a two-timescale stochastic approximation framework. \n3. The paper is well-structured and clear, with detailed explanations of the problem formulation, the PFEDRL-REP framework, and the two-timescale convergence analysis.", "weaknesses": "1. The experimental evaluation could be extended to include more complex environments, such as those with sparse rewards or high-dimensional state spaces, to better assess the scalability of PFEDRL-REP. \n2. The applicability of PFEDRL-REP to all types of environmental heterogeneity is not fully guaranteed, as the combination of shared feature representations and personalized weight vectors may not capture all nuances of diverse environments.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "i2kujctX20", "reviewer": "Reviewer_mFXK", "summary": "The manuscript introduces PFEDRL, a framework for personalized federated reinforcement learning (FedRL) aimed at addressing heterogeneity across agent environments. The authors propose PFEDTD-REP, a specific instantiation of PFEDRL with temporal difference (TD) learning. Notably, they claim a linear speedup in convergence proportional to the number of agents, a desirable characteristic in large-scale federated RL systems. Experimental results in both value-based learning (CliffWalking) and control tasks (CartPole, Acrobot) validate the framework's advantages, with promising outcomes in personalization and convergence speedup.", "questions": "1. how does agent count N affects personalization? if we add more agents by increasing N, do we have better personlization or worse?\n2. is there a tradeoff between personalization and global performance or the speedup?\n3. how to determine the optimal number of agents? If we add more agents, will we get faster convergence? what about personalization? intuitively, more agents should increase the personalization complexity.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "+The paper is practically relevant and addresses real-world heterogeneity in federated RL environments\n\n+The manuscript made some innovations in FedRL theories:\n\n- First work to prove linear speedup in personalized federated RL with shared representations under Markovian noise\n\n- Rigorous analysis of convergence rates using two-timescale stochastic approximation theory", "weaknesses": "while there are merits in the paper’s theoretical contributions, I am concerned with a few critical points:\n\n1. regarding the motivation on personalization:\n\na. the paper lacks formal definition of what constitutes successful personalization. The authors should consider to design metrics to quantify personalization quality\n\nb. thus, no theoretical guarantees that learned personalization (via agent-specific parameters in the paper) captures meaningful environment-specific adaptations. What happens to personalization quality when environments are very different from each other?\n\nc. how does agent count N affects personalization? if we add more agents by increasing N, do we have better personlization or worse?\n\nd. is there a tradeoff between personalization and global performance or the speedup?\n\n2. regarding the problem formulation and approach:\n\na. does sharing this common representation breach privacy preservation?\n\nb. In section 2.2, why does the transition from (1) to (2) preserve the problem properties? \n\n3. regarding theoretical rigor\n\na. corollary 4.15 claims that linear speedup w.r.t agent count N is justified by “we can proportionally decrease T as N increases while keeping the same convergence rate”. However, this is not a precise claim as linear speedup generally implies that adding more agents N directly enhances the convergence rate, without requiring adjustments to T. Here, the convergence rate remains fixed only by reducing T, which does not reflect true linear acceleration in convergence. The paper could be misleading readers into believing that the convergence rate inherently improves with more agents, rather than simply adjusting the number of communication rounds to balance computational costs\n\nb. related to 3.a) above, if the authors claims linear speedup w.r.t agent count N, they should provide comprehensive experimental validation showing how convergence behavior scales with varying numbers of agents. Notable prior work making similar speedup claims, such as Fan et al. [1] which is one of the earliest FedRL works missing from the related work, included thorough ablation studies demonstrating the impact of agent count N on convergence. The absence of such analysis is particularly concerning given the centrality of the linear speedup claim to the paper's contributions.\n\nc. otherwise, how to determine the optimal number of agents? If we add more agents, will we get faster convergence? what about personalization? intuitively, more agents should increase the personalization complexity. \n\n4. regarding experimental evaluation.\n\na. Limited diversity in test environments (only classic control tasks) and no statistical significance is assessed \n\nb. how to empirically verify the personalization achieved?\n\nc. ablation on agent count N should be conducted.\n\n---\n[1] Fan, X., Ma, Y., Dai, Z., Jing, W., Tan, C., & Low, B. K. H. (2021). Fault-tolerant federated reinforcement learning with theoretical guarantee. *Advances in Neural Information Processing Systems*, *34*, 1007-1021.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)