Personalized Observation Normalization for Federated Reinforcement Learning in Simulation Environments with Heterogeneity

Yiran Pang; Zhen Ni; Xiangnan Zhong

cs.LGApr 10, 2026

Personalized Observation Normalization for Federated Reinforcement Learning in Simulation Environments with Heterogeneity

Authors: Yiran Pang, Zhen Ni, Xiangnan Zhong

Organizations: Department of Electrical Engineering & Computer Science, Florida Atlantic University, Boca Raton, FL, USA

Abstract

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.

Figures & tables

Panel of Figure 1 — Fig. 1 : (a) Illustration of the FedRL system. The global model is distributed to and aggregated from local models, with each local model interacting with its distinct environment. (b) Observation distributions from three agents in heterogeneous HalfCheetah environments. The upper plot demonstrates the density of observation values across agents, highlighting the mean and variation in input distributions caused by differing state-transition dynamics. The lower plot shows the effect of applying personalized observation normalization (PON), which aligns the distributions across agents.

Fig. 2 : Illustration of the issue where a model with a smaller norm is overshadowed by one with a larger norm during aggregation, as discussed in FL scenario [ 29 ] .

Fig. 3 : The train observation variation over environment steps for agents with and without PON in heterogeneous HalfCheetah. The figure illustrates that the input variation fluctuates due to evolving policies and environment heterogeneity, while PON effectively reduces variation.

Fig. 4 : The mean observation trajectories over training steps for three agents using ON with normalization parameters shearing (shared ON). The shared parameters lead to inconsistent observation scaling across heterogeneous environments, causing significant fluctuations in the observation mean.

Fig. 5 : Performance comparison of independent learning PPO, FedRL-PPO (PPO with FedAvg for Actor and Critic), and FedRL-PPO with PON in three Mujoco environments. Each environment includes three heterogeneous settings and the average performance of agents. The shaded area represents the standard deviation. Objective value is plotted against the number of environment steps. All experiments were repeated with five different random seeds. FedRL with PON achieves faster convergence and higher final performance across all settings.

Method	HalfCheetah	Ant	Walker2d
PPO	$1923.08\pm 828.42$	$2168.26\pm 529.07$	$2304.43\pm 571.84$
FedRL-PPO	$2239.04\pm 1098.35$	$3370.11\pm 386.52$	$2931.00\pm 646.13$
FedRL-PPO w/ PON	$\mathbf{3938.03\pm 288.82}$	$\mathbf{4354.00\pm 517.79}$	$\mathbf{3532.11\pm 403.20}$

TABLE I : Performance comparison of different methods across several heterogeneous environments. The value represents the average and standard deviation of the best performance of three agents in five repeated experiments.

Fig. 6 : Illustrates the FedRL training curves of PON and its shared-parameter variant, shared ON, on the heterogeneous HalfCheetah-v3 task. The solid line and shaded area represent the average and standard deviation of the objective value of three agents.

Federated RL

Federated reinforcement learning lets distributed agents learn policies collaboratively while keeping their local data and experience private.

See all 34 papers

Personalized Federated Reinforcement Learning via Model-Agnostic Meta-Learning: Convergence of Exact and Hessian-Free Meta-Policy Gradients
Sep 19, 2026Ali Beikmohammadi, Sarit Khirirat, Sindri MagnússonFederated RL Meta-Learning
Department of Computer and Systems Sciences, Stockholm University, Stockholm, Sweden. · Faculty of Information and Communication Technology, Mahidol University, Nakhon Pathom, Thailand.
FedGuide: Diffusion Prior Alignment and Value Baseline Guidance for Heterogeneous Federated Reinforcement Learning
Jul 22, 2026Zhilin He, Gauri JoshiFederated RL Federated Learning
Carnegie Mellon University
Efficient Federated RLHF via Zeroth-Order Policy Optimization
Apr 20, 2026Deyi Wang, Qining Zhang, Lei YingFederated RL Communication-Efficient Distributed Training
University of Michigan, Ann Arbor

Sep 19, 2026J/K move · Enter open · S save

Personalized Federated Reinforcement Learning via Model-Agnostic Meta-Learning: Convergence of Exact and Hessian-Free Meta-Policy Gradients

Ali Beikmohammadi, Sarit Khirirat, Sindri Magnússon

Department of Computer and Systems Sciences, Stockholm University, Stockholm, Sweden. · Faculty of Information and Communication Technology, Mahidol University, Nakhon Pathom, Thailand.

Federated RL Meta-Learning Meta-Reinforcement Learning Policy Gradient Methods Second-Order Optimization

Open paper AskPDF arXiv

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

\varepsilon

-approximate first-order stationary point of the personalized objective in

K=\mathcal O(\varepsilon^{-3/2})

rounds with

τ=Θ(\varepsilon^{-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.

Abstract

Federated RL

Explore similar work

Personalized Federated Reinforcement Learning via Model-Agnostic Meta-Learning: Convergence of Exact and Hessian-Free Meta-Policy Gradients

FedGuide: Diffusion Prior Alignment and Value Baseline Guidance for Heterogeneous Federated Reinforcement Learning

Efficient Federated RLHF via Zeroth-Order Policy Optimization