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

Zhilin He; Gauri Joshi

cs.LGJul 22, 2026

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

Authors: Zhilin He, Gauri Joshi

Organizations: Carnegie Mellon University

Abstract

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 \textbf{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.

Figures & tables

Panel of Figure 1 — Figure 1: Method Visualization with 2 Clients Heterogeneous Settings. (a) Policy aggregation averages local optima $\mathbf{x}_{1}^{*},\mathbf{x}_{2}^{*}$ into global policy $\mathbf{x}^{(t,H)}$ , away from the global optimum $\mathbf{x}^{*}$ . (b) OT-MoE aggregates diffusion priors in distribution space, forming shared prior support $\bar{\pi}_{D}$ and personalized priors $\bar{\pi}_{D,1},\bar{\pi}_{D,2}$ (pink). The prior regularizer $\widehat{\mathcal{R}}$ (yellow) pulls local policies toward personalized prior supports. (c) FedGuide further uses the DICE value baseline (green) for policy improvement, shown as straighter policies updates and smaller endpoint dispersion ( $\mathrm{Variance}\downarrow$ ).

Figure 2

Variants	Diffusion Prior	DICE Value	Client Policy
FG-A	OT-MoE	Local	Avg.
FG-P	OT-MoE	N/A	Local
FG	OT-MoE	Local	Local

Table 1: Comparison of FedGuide (FG) variants: FedGuide keeps the DICE value baseline

\mathbf{V}_{i}

local, while FedGuide-P does not use

V_{\phi,i}

(

\beta=1

). FedGuide-A uses

V_{\phi,i}

but periodically averages client policies at the server. All variants use OT-MoE aggregation for diffusion prior.

	$\mathrm{CV}_{\sigma_{T}}$ $\downarrow$
Env	FG-A	FG-P	FG
Reacher	0.380	0.068	0.053
Hopper	0.402	0.164	0.155
Walker2D	0.208	0.179	0.127
HalfCheetah	0.628	0.235	0.241
MetaWorld10	0.364	0.272	0.260

Table 2: Normalized temporal volatility

\mathrm{CV}_{\sigma_{T}}

for FedGuide (FG) and variants, computed over the last 20 rounds as the across-seed mean of return standard deviation normalized by the absolute mean return. Smaller values indicate smoother late-stage returns and are consistent with, but do not directly measure, the variance-reduction effect in Lemma 2.

Panel of Figure 2 — Figure 2: Left: Demonstrations of different heterogeneous FRL environments: Reacher, Walker2D, Hopper, HalfCheetah, and MetaWorld10 with 10 heterogeneous manipulation tasks. Right: Visualization of toy-case Bandit2D with 4 isotropic Gaussian reward peaks ( $\sigma{=}0.2$ ) evenly spaced on the unit circle, each client sampled from an overlapping $120^{\circ}$ sector of the annulus $r\in[0.7,1.3]$ centered on one peak. All clients in FedAvg collapse to one peak, while FedGuide (including its two variants) captures each peak; its global diffusion prior recovers the ring-shaped union of client supports through OT-MoE aggregation.

Panel of Figure 3 — Figure 3: Client average return for 100 rounds at different heterogeneous FRL environments. We set client numbers to 10 in MetaWorld10 and keep 8 clients in other environments, evaluating on 5 random seeds with Behavior Cloning warm-up. Color regions indicate one standard deviation.

Figure 4: Client average return for FRL environments with stronger heterogeneity (“Hard”). We fixed client numbers at 8 and the same experimental settings with Fig. 3 .

Appendix figures & tables15 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 5: Client-level heterogeneity heatmaps for the main benchmark settings. Hopper, Walker2D, and HalfCheetah show normalized dynamics and reward metadata; Reacher shows normalized goal, action-noise, reward, and angle-shift metadata; MetaWorld10 shows task-semantic attributes.

Figure 6: Main-to-Hard metadata range expansion. For each metadata dimension, client values are normalized by the pooled Main and Hard range and plotted as side-by-side Main and Hard distributions.

Method		Heterogeneity Type	Aggregation Target	Alignment Space	Multi-Modal Support	Local Policy	Return-Aware Baseline
Baselines	FedAvg [ 23 ]	Data	Policy Params.	Parameter Space	✗	✗	✗
	FedKL [ 44 ]	Data	Policy Params.	Policy Output	✗	✗	P
	FedRL [ 13 ]	MDP / Dynamics	Actor + Critic	Parameter Space	✗	✗	P
	FedSVRPG-M [ 37 ]	MDP / Dynamics	Policy Grad.	Gradient Space	✗	✗	P
Variants	FedGuide-P	Data / MDP / Task	Diffusion Prior	Distribution Space	✓	✓	✗
	FedGuide-A	Data / MDP / Task	Prior + Policy	Distribution Space	✓	✗	✓

Appendix

Table 4: Comparison of the baselines and FedGuide variants used in our experiments. FedAvg aggregates policy parameters, FedKL regularizes local policies through KL penalties in the policy-output space, FedRL aggregates actor-critic models, and FedSVRPG-M aggregates variance-reduced policy-gradient updates. FedGuide instead aggregates diffusion behavior priors in distribution space while keeping each client’s policy local. FedGuide-P removes the DICE value baseline, and FedGuide-A additionally averages client policies. Here, ✓ denotes explicit support, ✗ denotes no explicit support, and P denotes partial or indirect support.

Method		Convergence	Communication Speedup	Local Updates Allowed	Distribution-Level Analysis
Baselines	FedAvg [ 23 ]	FL convergence	Empirical	✓	✗
	FedKL [ 44 ]	Asymptotic	Not claimed	P	Policy-output KL
	FedRL [ 13 ]	Finite to subopt.	Not claimed	✓	✗
	FedSVRPG-M [ 37 ]	Variance-reduced PG	Claimed	✓	Policy-gradient correction
Variants	FedGuide-P	Prior-only ablation	Not isolated	✓	OT-MoE diffusion prior
	FedGuide-A	Controlled avg. error	Not isolated	✓	OT-MoE prior + policy avg.

Appendix

Table 5: Theoretical and mechanism-level comparison of the baselines and FedGuide variants. FedAvg provides the standard federated policy averaging baseline, FedKL provides convergence analysis for KL-regularized policy-output alignment, FedRL analyzes federated RL under environment heterogeneity with convergence to a heterogeneity dependent suboptimal solution, and FedSVRPG-M uses momentum-based variance-reduced policy-gradient correction for efficient heterogeneous FRL. FedGuide keeps policies local and aggregates only diffusion-prior heads through OT-MoE; its analysis separates the effective stochastic variance, OT errors, prior or value surrogate errors, and local KL drift, while the DICE value baseline reduces the variance of local policy updates.

Parameter	Symbol	Bandit2D	Reacher	MuJoCo locomotion	MetaWorld10
Offline buffer size	$\|\mathcal{D}_{i}\|$	$10^{3}$	$2{\times}10^{4}$	$2{\times}10^{5}$ cap	$5{\times}10^{3}$
Offline source	–	local modes	client rollouts	D4RL	scripted policy
Diffusion steps	$K$	Gaussian	$1000$	$1000$	$1000$
Prior width / horizon	$d_{D},L_{D}$	Gaussian	$64,64$	$64,64$	$64,64$
Prior epochs / batch	$E_{D},B_{D}$	closed form	$40,512$	$40,512$	$40,512$
Prior learning rate	$\eta_{D}$	–	$10^{-4}$	$10^{-4}$	$10^{-4}$

Appendix

Table 6: Offline pretraining hyperparameters.

Parameter	Symbol	Bandit2D	Reacher	MuJoCo locomotion	MetaWorld10
Clients	$N$	$4$	$8$	$8$	$10$
Prior experts	$M$	$4$	$8$	$8$	$8$
Federated rounds	$H$	$60$	$100$	$100$	$100$
Rollout size / round	$T_{\mathrm{roll}}$	$200$	$2048$	$4096$	$4096$
Local PPO epochs	$U$	$4$	$4$	$4$	$4$
PPO mini-batch	$B$	$64$	$64$	$64$	$64$

Appendix

Table 7: Shared online hyperparameters for FedGuide variants.

Method	Prior aggregation	Value baseline	Policy aggregation	Key coefficients
FedGuide	OT-MoE on $\{\psi_{i}\}_{i=1}^{N}$	$\beta=0.5$	none	$\lambda_{1}=0.05$ , $\lambda_{2}=0.5$ , $\eta=0.05$
FedGuide-P	OT-MoE on $\{\psi_{i}\}_{i=1}^{N}$	$\beta=1$	none	$\lambda_{1}=0.05$ , $\lambda_{2}=0.5$ , $\eta=0.05$
FedGuide-A	OT-MoE on $\{\psi_{i}\}_{i=1}^{N}$	$\beta=0.5$	every 5 rounds	$\lambda_{1}=0.05$ , $\lambda_{2}=0.5$ , $\eta=0.05$

Appendix

Table 8: FedGuide variant hyperparameters.

Parameter	Symbol	FedAvg	FedKL	FedRL	FedSVRPG-M
Local optimizer	–	PPO	PPO	DDPG	PPO/SVRPG
Aggregated parameters	–	$\theta_{i}$	$\theta_{i}$	$\theta_{i}^{\mu},\theta_{i}^{Q}$	$\Delta\theta_{i}$
Global KL weight	$\lambda_{\mathrm{g}}$	$0$	$0.1$	–	–
Local KL weight	$\lambda_{\mathrm{l}}$	$0$	$0.05$	–	–
Local PPO epochs	$U$	$10$	$10$	–	$4$
Mini-batch size	$B$	$64$	$64$	$64$	$64$

Appendix

Table 9: Baseline hyperparameters.

Panel of Figure 7 — Figure 7: Bandit2D sensitivity to the expert count $M$ , the Sinkhorn temperature $\eta$ , and the prior weight $\lambda_{2}$ .

Panel of Figure 8 — Figure 8: Weak-overlap stress test. Client 5 has a true target mode at the origin, while the four experts contain only ring-shaped behavior support. OT-MoE routes Client 5 to a mixture of the existing experts, keeping its personalized prior on the ring rather than creating the unseen center mode.

Figure 9: FedGuide-family qualitative rollouts on Reacher, visualized with motion from a single final-round evaluation episode.

Figure 10: FedGuide-family qualitative rollouts on Hopper. Each cell overlays sampled frames from the final-round rollout to show the learned motion trace.

Figure 11: FedGuide-family qualitative rollouts on Walker2D, visualized with motion across heterogeneous clients.

Figure 12: FedGuide-family qualitative rollouts on HalfCheetah, visualized with motion overlays across the first eight clients.

Figure 13: FedGuide-family qualitative rollouts on MetaWorld10, visualized with across the ten task clients.

Explore similar work

Exploration-Driven Personalized Federated Reinforcement Learning via Intrinsic Motivation
Aug 11, 2026Md Rafid Islam, Rafsan Jany, Zahid Hasan +1Federated Learning Privacy
Department of Electrical and Computer Engineering North South University Dhaka, Bangladesh · Digital Health Research Division Korea Institute of Oriental Medicine Daejeon, South Korea · Department of Electrical and Computer Engineering The University of Alabama in Huntsville Huntsville, AL, USA
Efficient Federated RLHF via Zeroth-Order Policy Optimization
Apr 20, 2026Deyi Wang, Qining Zhang, Lei YingReinforcement Learning From Human Feedback Federated Learning
University of Michigan, Ann Arbor
FedQHD: Closed-Form Function-Space Federated Reinforcement Learning
May 27, 2026Yuchen Hou, Yongshan Chen, Zhuowen Zou +4Federated Averaging
Northeastern University · University of California, Irvine · The George Washington University

Aug 11, 2026, cs.LGJ/K move · Enter open · S save

Exploration-Driven Personalized Federated Reinforcement Learning via Intrinsic Motivation

Md Rafid Islam, Rafsan Jany, Zahid Hasan, Ratun Rahman

Department of Electrical and Computer Engineering North South University Dhaka, Bangladesh · Digital Health Research Division Korea Institute of Oriental Medicine Daejeon, South Korea · Department of Electrical and Computer Engineering The University of Alabama in Huntsville Huntsville, AL, USA

Federated Learning Privacy Intrinsic Motivation

Open paper AskPDF arXiv

Personalized Federated Reinforcement Learning (PFRL) takes a decentralized approach to storing and accessing information based on past experiences while keeping each client's data private during the learning of each client's policy. Many current methods for PFRL rely heavily on exploiting existing reinforcement learning reward signals to derive an optimal policy for each client, thereby neglecting exploration in non-stationary or sparse-reward environments. In this work, we introduce a new exploration-driven framework, Exploration-Driven Personalized Federated Reinforcement Learning via Intrinsic Motivation (EDPFRL-IM), that leverages an inherent curiosity-driven exploration at each client to promote local exploration and protect client privacy. Furthermore, to facilitate policy discovery via exploration in previously unexplored state spaces, clients add an intrinsic random network distillation (RND) signal to their extrinsic reward. Additionally, the server does not have access to clients' raw experiences or local gradient estimates; instead, the server sends global exploration priors and collects minimal novelty summaries from each client to enable both diverse and coordinated exploration among clients. Experiments in benchmark environments show that our framework outperforms average PFRL benchmarks in policy personalization and sample efficiency, primarily in delayed and sparse reward systems. Overall, EDPFRL-IM enables the integration of a flexible exploratory learning structure into federated reinforcement learning systems while preserving client privacy.

Abstract

Explore similar work

Exploration-Driven Personalized Federated Reinforcement Learning via Intrinsic Motivation

Efficient Federated RLHF via Zeroth-Order Policy Optimization

FedQHD: Closed-Form Function-Space Federated Reinforcement Learning