Federated Distributional Reinforcement Learning with Distributional Critic Regularization

David Millard; Cecilia Alm; Rashid Ali; Pengcheng Shi; Ali Baheri

cs.LGMar 18, 2026

Federated Distributional Reinforcement Learning with Distributional Critic Regularization

Authors: David Millard, Cecilia Alm, Rashid Ali, Pengcheng Shi, Ali Baheri

Organizations: Dept. of Mechanical Engineering, Rochester Institute of Technology, Rochester, NY 14623, USA · Dept. of Psychology and School of Information, Rochester Institute of Technology, Rochester, NY 14623, USA · Department of Engineering Science, University West, Trollhättan 461 30, Sweden · Golisano College of Computing and Information Sciences, Rochester Institute of Technology, Rochester, NY 14623, USA

Abstract

Federated reinforcement learning typically aggregates value functions or policies by parameter averaging, which emphasizes expected return and can obscure statistical multimodality and tail behavior that matter in safety-critical settings. We formalize federated distributional reinforcement learning (FedDistRL), where clients parametrize quantile value function critics and federate these networks only. We also propose TR-FedDistRL, which builds a per client, risk-aware Wasserstein barycenter over a temporal buffer. This local barycenter provides a reference region to constrain the parameter averaged critic, ensuring necessary distributional information is not averaged out during the federation process. The distributional trust region is implemented as a shrink-squash step around this reference. Under fixed-policy evaluation, the feasibility map is nonexpansive and the update is contractive in a probe-set Wasserstein metric under evaluation. Experiments on a bandit, multi-agent gridworld, and continuous highway environment show reduced mean-smearing, improved safety proxies (catastrophe/accident rate), and lower critic/policy drift versus mean-oriented and non-federated baselines.

Figures & tables

Panel of Figure 1 — Fig. 1: Each panel shows how aggregation operators combine scalar, multimodal client reward distributions in the bandit setting (i.e., $Z^{\pi}=r_{t}$ , $\gamma=0$ ). Each client is parameterized with QR-DQN models and averaging is computed via Equation 9 . Left: critic parameter averaging ( FedAvg ) induces an output-space arithmetic averaging effect that can smear modes and attenuate tail mass (mean-smearing). Middle: an unweighted Wasserstein-1 barycenter aggregates in distribution space and better preserves geometric structure (modes/tails) compared with arithmetic averaging. Right: a CVaR-risk-weighted Wasserstein-1 barycenter biases the barycenter toward lower-tail behavior by upweighting clients with larger lower-tail risk (CVaR at level $\alpha=0.1$ ), helping retain tail-relevant mass while still preserving multimodality.

Fig. 2: Case Study 2 client environments for Clients 1–3. Each panel shows one client’s environment with obstacles, hazard zone, agent start/goal locations, and object start/goal locations. The three client environments are arranged for direct side-by-side comparison of environmental heterogeneity. All instantiations use a $10\times 10$ grid with 3 agents and 3 objects.

Method	Return $\uparrow$	Output Drift $\downarrow$	Catastrophe $\downarrow$
Local	483.09 $\pm$ 128.40	113.37 $\pm$ 27.94	0.0323 $\pm$ 0.0076
FedAvg	554.76 $\pm$ 95.51	159.71 $\pm$ 18.56	0.0240 $\pm$ 0.0054
FedAvg (CVaR)	544.56 $\pm$ 74.24	151.51 $\pm$ 17.09	0.0203 $\pm$ 0.0037
FedAvg (CVaR + TR)	436.21 $\pm$ 54.13	98.07 $\pm$ 11.77	0.0172 $\pm$ 0.0032

TABLE I: Results are mean

\pm

standard deviation over 30 random seeds. Return is the average episodic return (higher is better). Output Drift quantifies across-round instability of the distributional critic outputs on a fixed evaluation/probe set (lower means the critic’s predicted return laws change less from round to round and are less perturbed by federation). Catastrophe is the evaluation catastrophe rate (fraction of episodes ending in a catastrophic event; lower is better).

Fig. 3: Bars report the percent change in catastrophe rate for each heterogeneous client relative to the Local baseline (baseline corresponds to 0% by definition), aggregated over 30 random seeds; negative values indicate fewer catastrophes (safer) than Local . Catastrophe rate is the fraction of evaluation episodes that terminate in a designated catastrophic event (environment-defined unsafe terminal condition). Error bars denote variability across seeds (standard deviation).

Fig. 4: Accident rate across training steps for each heterogeneous client in the highway environment (Case Study 3), comparing Local , FedAvg , FedAvg (CVaR) , and FedAvg (CVaR + TR) over 30 random seeds. Shaded regions denote standard deviation across seeds. Lower accident rate indicates safer behavior. FedAvg (CVaR + TR) consistently achieves the lowest accident rate across clients, demonstrating that the proposed risk-aware trust-region regularization improves safety in a continuous-control setting under environment heterogeneity.

Method	Return $\uparrow$	Accident $\downarrow$
Local	23.197 $\pm$ 4.830	0.811 $\pm$ 0.063
FedAvg	27.797 $\pm$ 4.219	0.736 $\pm$ 0.068
FedAvg (CVaR)	30.168 $\pm$ 3.996	0.678 $\pm$ 0.066
FedAvg (CVaR + TR)	31.231 $\pm$ 3.943	0.653 $\pm$ 0.066

TABLE II: Results are mean

\pm

standard deviation over 30 random seeds. Return is the average episodic return (higher is better). Accident is the evaluation accident rate (fraction of episodes with at least one collision/accident event; lower is better). Methods mirror Case Study 2.

Federated RL

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

See all 34 papers

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
Constraint-Aware Aggregation for Federated Reinforcement Learning in Microgrid Energy Coordination
Jul 14, 2026Usman Haider, Karl MasonFederated Learning Aggregation Federated RL
School of Computer Science, University of Galway, Ireland.
Learning from Unreliable Trajectories: Adversarially-Robust Federated Q-Learning
Oct 2, 2026Sreejeet Maity, Aritra MitraFederated RL Byzantine-Robust Federated Learning
Department of Electrical and Computer Engineering North Carolina State University, Raleigh

Jul 22, 2026J/K move · Enter open · S save

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

Zhilin He, Gauri Joshi

Carnegie Mellon University

Federated RL Federated Learning Heterogeneous Federated Learning

Open paper AskPDF arXiv

Federated Reinforcement Learning (FRL) enables collaborative policy learning across distributed agents with heterogeneous environments. While recent methods based on variance reduction, divergence penalization, and momentum optimization improve FRL under heterogeneous settings, they still primarily synchronize policy or value-network parameters and do not explicitly address distributional mismatch among heterogeneous clients. Therefore, we propose FedGuide, a FRL framework that uses diffusion priors as behavior models to provide personalized data supported distributions for heterogeneous local policy learning. Instead of directly averaging local policies, FedGuide aggregates those diffusion priors through Optimal-Transport Mixture-of-Experts (OT-MoE), preserving heterogeneous behavior modes in distribution space. It further develops a Distribution Correction Estimation (DICE) value baseline to provide low-variance, return-aware guidance for local policy improvement. Experiments across heterogeneous environments show that FedGuide outperforms representative FRL methods in client-average returns, final-round performance, and worst-round robustness, while maintaining stable learning under stronger heterogeneity.

Abstract

Federated RL

Explore similar work

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

Constraint-Aware Aggregation for Federated Reinforcement Learning in Microgrid Energy Coordination

Learning from Unreliable Trajectories: Adversarially-Robust Federated Q-Learning