English

Federated Distributional Reinforcement Learning with Distributional Critic Regularization

Machine Learning 2026-03-19 v1

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.

Keywords

Cite

@article{arxiv.2603.17820,
  title  = {Federated Distributional Reinforcement Learning with Distributional Critic Regularization},
  author = {David Millard and Cecilia Alm and Rashid Ali and Pengcheng Shi and Ali Baheri},
  journal= {arXiv preprint arXiv:2603.17820},
  year   = {2026}
}

Comments

9 pages, 4 Figures, conference

R2 v1 2026-07-01T11:26:21.883Z