On the Linear Speedup of Personalized Federated Reinforcement Learning with Shared Representations

Authors: GUOJUN XIONG, Shufan Wang, Daniel Jiang, Jian Li
Primary area: reinforcement learning
Keywords: personalized federated reinforcement learning, shared representations, stochastic approximation

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.

Decision: accept (poster)
Reviewer scores: 6, 8, 6

Reviews and discussion:
<thread>
creation date:
1730532973806

thread_type:
official_review

<thread_object>
creator: reviewer_mfxk

<content>
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.

strengths:
+The paper is practically relevant and addresses real-world heterogeneity in federated RL environments

+The manuscript made some innovations in FedRL theories:

- First work to prove linear speedup in personalized federated RL with shared representations under Markovian noise

- 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:

1. regarding the motivation on personalization:

a. the paper lacks formal definition of what constitutes successful personalization. The authors should consider to design metrics to quantify personalization quality

b. 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?

c. how does agent count N affects personalization? if we add more agents by increasing N, do we have better personlization or worse?

d. is there a tradeoff between personalization and global performance or the speedup?

2. regarding the problem formulation and approach:

a. does sharing this common representation breach privacy preservation?

b. In section 2.2, why does the transition from (1) to (2) preserve the problem properties? 

3. regarding theoretical rigor

a. 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

b. 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.

c. 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. 

4. regarding experimental evaluation.

a. Limited diversity in test environments (only classic control tasks) and no statistical significance is assessed 

b. how to empirically verify the personalization achieved?

c. ablation on agent count N should be conducted.

---
[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.

questions:
1. how does agent count N affects personalization? if we add more agents by increasing N, do we have better personlization or worse?
2. is there a tradeoff between personalization and global performance or the speedup?
3. 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.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730714493442

thread_type:
official_review

<thread_object>
creator: reviewer_lerk

<content>
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.

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.
2. The paper provides a rigorous theoretical foundation, including proofs of convergence speedup under Markovian noise using a two-timescale stochastic approximation framework. 
3. 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. 
2. 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.

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?
2. 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.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730838356201

thread_type:
official_review

<thread_object>
creator: reviewer_um15

<content>
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.

strengths:
- The theoretical analysis have a proof sketch and the assumptions are clearly stated.
- I think the two timescale approximation result is novel
- 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.
- If I understand correctly, the linear speed up is not a particular exciting result, since sample collection in unit time grows linearly with N.
- There is no comparison with other PFL methods, or parameter-sharing MARL methods.

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?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732495937846

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Revision posted

comment:
Dear reviewers,

Thanks for the valuable feedback. We wanted to make sure you are aware that we've uploaded a revision of the paper before the end of the discussion phase. The new results are in blue and located in Appendix H:

* **More complex environment.** We've added a more complex environment (Hopper), under the recommendation of Reviewer LerK. Hopper has continuous state and action spaces, so we adapted our personalized FedRL framework to DDPG, resulting in a new instantiation of the algorithm PFedDDPG, which performs well compared to baselines (See Appendix H1 in the revision). In addition, this result illustrates that the PFedRL framework can be used quite broadly (TD, DQN, DDPG).

* **Linear speedup wrt to agent count N.** As suggested by Reviewer mFXK, we've conducted an experiment varying the number of agents from 2 through 10 to verify our theoretical results. Indeed, we see that the convergence time decreases nearly linearly as the number of agents increases. Please see Appendix H2.

Thanks again!
</content>
</thread_object>

</thread>
<thread>
creation date:
1732728324086

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Another revision posted

comment:
Dear reviewers,

Given our discussion with Reviewer mFXK, we have posted another revision of the paper. Please refer to the latest revision because we have formatted the new results to be easier to read (on separate pages and different sections). Our revision includes several new experiment results:
* A more complex environment (Hopper) in **Appendix H1**
* Verifying the linear speedup theoretical result in **Appendix H2**
* Examining the tradeoff between computation and personalization quality in **Appendix H3**
* Examining the effect of environment discrepancy on personalization error in **Appendix H4**

As always, thanks for your comments and please follow up if we can answer or clarify additional points as the discussion period winds down soon.
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper studies personalized federated reinforcement learning (FedRL) with shared representations, presenting theoretical convergence results and demonstrating linear speedup in the proposed framework, PFedRL-Rep, under Markovian noise. The results are supported by rigorous theoretical analysis and experimental validation, including performance improvements in heterogeneous environments. The reviewers highlighted the novelty and soundness of the theoretical contributions, particularly in establishing linear speedup and addressing personalization in FedRL. During the rebuttal phase, the authors effectively addressed concerns raised by the reviewers, providing additional experimental ablations and clarified theoretical guarantees. Overall, the reviewers reached a positive consensus, and the paper makes a substantial theoretical contribution to the field, making it acceptable for publication.