Personalized Federated Reinforcement Learning via Model-Agnostic Meta-Learning: Convergence of Exact and Hessian-Free Meta-Policy Gradients
Authors: Ali Beikmohammadi, Sarit Khirirat, Sindri Magnússon
Organizations: Department of Computer and Systems Sciences, Stockholm University, Stockholm, Sweden. · Faculty of Information and Communication Technology, Mahidol University, Nakhon Pathom, Thailand.
We study personalized federated reinforcement learning, in which n agents, each acting in its own Markov decision process, collaborate through a server to learn a shared MAML-style policy initialization that becomes effective for an individual agent once that agent adapts it with a single local policy-gradient step. We propose Per-FedAvg-PG, in which agents take τ local stochastic meta-policy-gradient steps between communication rounds, and prove that it reaches an ε-approximate first-order stationary point of the personalized objective in K=O(ε−3/2) rounds with τ=Θ(ε−1/2) local steps. The analysis rests on a structural feature of the reinforcement learning setting: under standard policy-class regularity, the per-agent objectives have uniformly bounded gradients and Hessians with explicit constants, so the bounded-gradient and bounded-heterogeneity conditions imposed by the supervised theory hold automatically and no separate heterogeneity assumption is needed. The exact meta-gradient requires the inner-loop policy Hessian, which our experiments identify as the practical bottleneck. We therefore analyze the Hessian-free variant, bound its bias, and exhibit fixed points at which the meta-gradient is nonzero and of order α, showing that the resulting stationarity floor is a property of the method rather than of the bound. Experiments on tabular and neural navigation confirm the predicted behavior and show transfer to unseen agents at an order of magnitude lower sample cost than independent training. Together these results identify the adaptation step size as a tunable personalization knob and the curvature estimate as the quantity that governs whether exact meta-gradients are affordable.
Figures & tables
federated
RL
personalization mechanism
convergence guarantee
Per-FedAvg [ 7 ]
✓
—
MAML init. (1 step)
ε -FOSP, K=O(ε−3/2)
SG-MRL [ 5 ]
—
✓
MAML init. ( ≥1 steps)
ε -FOSP, finite-sample
MAML-LQR [ 25 ]
—
LQR
MAML init. (PG)
stabilizing, up to heterogeneity bias
FedNPG [ 29 ]
✓
✓
none (shared policy)
global opt. (tabular, entropy-reg.)
FedSARSA / FedPG [ 30 , 15 , 31 ]
✓
✓
none (shared policy)
finite-time, up to heterogeneity gap
Shared representations [ 28 ]
✓
✓
shared repr. + local heads
finite-time, linear speedup
Table 1: Position of this paper among the most closely related works. “ ε -FOSP” abbreviates convergence to an ε -approximate first-order stationary point. Unlike Per-FedAvg, our guarantees need no bounded-gradient or bounded-heterogeneity assumption: both are theorems in the policy-gradient setting (Lemma 1 , Remark 1 ).
Figure 1: Tabular gridworld, exact evaluation. Left: post-adaptation value F vs. per-agent trajectories; dashed is the exact variant before adaptation, dotted a random policy. Right: final F vs. α ; at α=0 all three coincide, as Proposition 3 requires.
Figure 2: Neural navigation, Monte Carlo evaluation. Left: curvature bottleneck: return of the exact variant vs. curvature batch mh (per-seed points, median line); destabilized until mh≈1000 , consistent with the α2LJ2/mh term of Lemma 3 . Right: few-shot return on held-out agents vs. adaptation trajectories; success rates in Figure 7 of Appendix C .
Appendix figures & tables7 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 3: One communication round of Per-FedAvg-PG (Algorithm 1 ). Only parameters cross the link; trajectories stay local.
Figure 4: Tabular family: the eight goal cells. Agents differ only in the goal.
Figure 5: Neural arc family: shared-structure goals with a held-out split.
Experiment
Env.
K
τ
α
β
min
mh
mout
Personalization; learning curves
tabular
80
5
2
0.3
10
10
10
Adaptation step size α
tabular
80
5
{0,41,21,1,2,3}
0.3
10
10
10
Curvature bottleneck
neural
150
5
1
0.2
10
{10,50,100,200,500,1000}
10
Few-shot personalization
neural
150
5
1
0.2
10
—
10
Appendix
Table 2: Hyperparameters. All runs use full participation and 10 seeds, with n=8 agents (tabular) or 6 training agents (neural). The first two rows run the exact variant, the first-order variant and non-personalized FedAvg-PG ; the third runs the exact variant while sweeping the Hessian batch mh ; the fourth runs the first-order variant. Here mh is used by the exact variant and, within its single policy-gradient estimate, by FedAvg-PG , so the first-order variant spends 20 trajectories per local step against 30 for the other two.
Figure 6: Learning curves (tabular, exact evaluation). Post-adaptation return against communication round k . All methods improve monotonically and the personalized variants stay ahead of non-personalized FedAvg-PG at every round, so the advantage reported in Figure 1 is not an artifact of where training was stopped. This is the same data as Figure 1 (left) with the horizontal axis expressed in rounds rather than per-agent trajectories, the two being proportional at fixed τ and batch sizes.
Figure 7: Success rate on held-out agents (neural, Monte Carlo evaluation): fraction of held-out episodes that enter the goal region, against the number of trajectories used to form the single adaptation step. Companion to Figure 2 (right), which reports return on the same runs.
Figure 8: Held-out rollouts (neural). Trajectories on an unseen goal (green star, shaded goal region; black dots are episode starts), each drawn until it first reaches the goal region (dotted paths never reach; reach counts in the panel titles). Before adaptation the first-order initialization already steers toward the region and reaches it in most rollouts; after a single 200 -trajectory adaptation step every rollout reaches the goal; from-scratch training at the same budget reaches it in fewer than half.
This paper considers reinforcement learning from human feedback in a federated learning setting with resource-constrained agents, such as edge devices. We propose an efficient federated RLHF algorithm, named Partitioned, Sign-based Stochastic Zeroth-order Policy Optimization (Par-S2ZPO). The algorithm is built on zeroth-order optimization with binary perturbation, resulting in low communication, computation, and memory complexity by design. Our theoretical analysis establishes an upper bound on the convergence rate of Par-S2ZPO, revealing that it is as efficient as its centralized counterpart in terms of sample complexity but converges faster in terms of policy update iterations. Our experimental results show that it outperforms a FedAvg-based RLHF on four MuJoCo RL tasks.
Despite the popularity of the actor-critic method and the practical needs of collaborative policy training, existing works typically either overlook environmental heterogeneity or give up personalization altogether by training a single shared policy across all agents. We consider a federated actor-critic framework in which agents share a common linear subspace representation while maintaining personalized local policy components, and agents iteratively estimate the common subspace, local critic heads, and local policies (i.e., actors). Under canonical single-timescale updates with Markovian sampling, we establish finite-time convergence via a novel joint linear approximation framework. Specifically, we show that the critic error converges to zero at the rate of O~(1/((1−γ)4TK)), and the policy gradient norm converges to zero at the rate of O~(1/((1−γ)6TK)), where T is the number of rounds, K is the number of agents, and γ∈(0,1) is the discount factor. These results demonstrate linear speedup with respect to the number of agents K, despite heterogeneous Markovian trajectories under distinct transition kernels and coupled learning dynamics. To address these challenges, we develop a new perturbation analysis for the projected subspace updates and QR decomposition steps, together with conditional mixing arguments for heterogeneous Markovian noise. Furthermore, to handle the additional complications induced by policy updates and temporal dependence, we establish fine-grained characterizations of the discrepancies between function evaluations under Markovian sampling and under temporally frozen policies. Experiments instantiate the framework within PPO on federated Hopper-v5 action-map heterogeneity, showing gains over Single PPO and FedAvg PPO and downstream transfer from the learned shared trunk.
Federated reinforcement learning (FedRL) enables multiple agents to collaboratively train a global policy without sharing raw data, making it ideal for privacy-sensitive applications. However, FedRL faces challenges in heterogeneous environments where differing state-transition dynamics lead to non-identical input distributions and imbalanced parameter updates during aggregation. Therefore, this paper develops a personalized observation normalization (PON) method, allowing each agent to locally normalize raw state inputs using a continuously updated running mean and variance. This design ensures consistent scaling of local feature without overshadowing across agents during aggregation. Furthermore, we demonstrate that sharing normalization parameters across agents is ineffective due to the diverse local input distributions, which highlights the necessity of personalized statistics. Experiments on heterogeneous MuJoCo tasks show that our developed PON accelerates training and achieves superior performance compared to baseline methods.
Yiran Pang, Zhen Ni, Xiangnan Zhong
Department of Electrical Engineering & Computer Science, Florida Atlantic University, Boca Raton, FL, USA