PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Zhilin He, Gauri Joshi

arXiv:2609.18964v1cs.LGcs.RO

TL;DR

Heterogeneous FRL must address mismatched client objectives and behavior distributions, which make naive model aggregation ineffective. FedGuide aggregates diffusion behavior priors through OT-MoE, uses a DICE value baseline for local policy improvement, and generally achieves higher returns and robust final- and worst-round performance.

  • Problem

    Heterogeneous robotic clients have mismatched local objectives and behavior distributions, making naive model aggregation ineffective.

    Source: §1 · p. 1

  • Method

    FedGuide shares diffusion behavior priors through OT-MoE instead of averaging local policies, and uses a state-only DICE value baseline to guide local policy improvement.

    Source: §1 · p. 2

  • Results

    FedGuide and its variants generally achieve higher returns over training and remain robust in final-round and worst-round performance across tested environments.

    Source: §6.1 · p. 7

  • Takeaways & Limitations

    Combining multimodal behavioral support with DICE-based local policy improvement is associated with more stable learning under heterogeneous clients.

    Source: §6.1 · p. 9

    FRL methods still struggle with heterogeneous task geometry beyond the environment-parameter heterogeneity mostly studied so far.

    Source: §6.1 · p. 8

Abstract

from arXiv · show

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.

1 Introduction

Heterogeneous FRL must align client behavior distributions, not only synchronize policy updates, because each client’s policy and MDP induce distinct trajectories. FedGuide addresses this mismatch with distribution-level diffusion priors and a state-only DICE baseline for local improvement.

1 Introduction

  • Client-specific policies interacting with local MDPs induce distinct trajectory distributions, which existing FRL methods do not explicitly model or align.This motivates distribution alignment rather than relying only on policy regularization, distance constraints, or momentum optimization.Source: §1 · p. 1
  • FedGuide communicates diffusion behavior priors in distribution space while keeping policy parameters local and avoiding raw-trajectory sharing.Personalized priors regularize local learning, and a state-only DICE baseline supports return-aware policy improvement.Source: §1 · p. 2
  • FedGuide couples local policy learning through OT-MoE personalized diffusion priors instead of directly averaging policies.The framework also introduces a prior-regularized PPO objective with a state-only DICE value baseline; the paper reports theoretical results on unbiasedness, variance reduction under positive residual correlation, and convergence to an approximate stationary point with an explicit residual error floor.Source: §1 · p. 2

2 Related Work

Prior FRL work addresses heterogeneity through policy-update optimization, shared-policy aggregation, or stabilization, while diffusion and DICE methods generally assume centralized data. FedGuide’s related-work gap is federated alignment of heterogeneous policy-induced distributions.

2 Related Work

  • Heterogeneous FRL methods use knowledge transfer, policy-divergence regularization, stabilized aggregation, and momentum-based or finite-time optimization analyses.These approaches primarily optimize policy updates or aggregate shared policies, leaving mismatched behavior modes at risk of collapse.Source: §2 · p. 2
  • Distribution-aware MoE and OT aggregation share structured distributions, but the cited approaches target static supervised data rather than policy-induced trajectories.The paper identifies this mismatch as a gap for heterogeneous FRL.Source: §2 · p. 2
  • Diffusion behavior models capture multimodal action and trajectory distributions, while DICE methods use distribution ratios to guide policies toward high-value in-sample actions.These methods generally assume one centralized behavior distribution and do not address communication, routing, or personalization across heterogeneous clients.Source: §2 · p. 2

3 Preliminaries

The setup has heterogeneous clients that share state and action spaces but differ in their MDP dynamics, rewards, and initial-state distributions. FedGuide uses diffusion models to represent data-supported actions and a state-only DICE estimate to center local policy advantages.

3.1 Problem Formulation

  • Clients share state and action spaces but have different transition, reward, and initial-state distributions, producing non-IID data and inconsistent gradients.The federated objective maximizes average expected return while keeping policies local and coupling clients through personalized distribution priors.Source: §3.1 · pp. 2–3

3.2 Diffusion Priors and DICE Value Functions

  • A state-conditioned diffusion model learns each client’s offline action distribution by corrupting data actions with noise and reversing the process toward data-supported actions.The resulting behavior policy models which actions occur in the dataset, regardless of return.Source: §3.2 · p. 3
  • DICE estimates discounted state-action distribution ratios through value functions, while the state-only offline value estimate centers advantages without directly scoring actions or reshaping the diffusion prior.The paper’s Figure 1 depicts the baseline as guidance for policy improvement and associates it with straighter updates and lower endpoint dispersion.Source: §3.2 · pp. 3–4, Figure 1 · p. 4

4 Methodology

FedGuide couples local policy learning through personalized diffusion-prior mixtures and uses a blended DICE value baseline for local advantage estimation. Its server routes client prior heads to experts, while policies and value estimators remain local.

4 Methodology

  • FedGuide combines OT-MoE alignment of personalized diffusion priors with low-variance guidance from a state-only DICE value baseline.The priors address distribution mismatch, while the baseline guides local policy improvement.Source: §4 · p. 4

4.1 OT-MoE Aggregation of the Diffusion Priors

  • Because averaging heterogeneous priors can collapse modes or assign mass to unsupported actions, OT soft-routes client prior heads to distributional experts.A Sinkhorn-regularized plan uses score-MSE costs, and column-normalized expert updates yield personalized mixtures for clients.Source: §4.1 · p. 4
  • Each local policy is regularized toward its personalized prior mixture, keeping policies client-specific while aligning them to offline-supported actions.The prior regularizer estimates a cross-entropy-style term using self-normalized importance weights over rollout samples.Source: §4.1 · p. 5

4.2 DICE Value Baseline Advantage Blending

  • The local advantage baseline blends the online critic with a pretrained state-only DICE value estimate before GAE and PPO updates.The DICE value enters through advantage estimation rather than directly reshaping the prior; the blend weight can be annealed toward the online critic as it becomes more reliable.Source: §4.2 · p. 5, §4.3 · p. 6

4.3 Putting Everything Together: FedGuide

  • Each round, the server routes client prior heads to experts, refreshes the expert distributions, and broadcasts personalized mixtures; clients perform local updates and upload prior heads.Policies, online critics, and DICE baselines remain local, and the trust-region term stabilizes local policy updates.Source: §4.3 · pp. 5–6

5 Theoretical Analysis

The analysis gives a convergence guarantee for FedGuide under stated smoothness, variance, drift, and approximation-error assumptions. It also shows that the state-only DICE baseline adds no direct score-function bias and can reduce stochastic-gradient variance under positive residual correlation.

5 Theoretical Analysis

  • The analysis assumes L-smoothness, bounded stochastic-gradient variance and local KL drift, plus bounded prior-surrogate, value-baseline, OT-cost, and Sinkhorn errors.These errors are denoted ϵSNIS, ϵV, ϵC, and ϵOT, respectively.Source: §5 · p. 6
  • The state-only DICE baseline does not add direct action-dependent score-function bias; its approximation error is included in ϵV.It centers the advantage estimate and does not directly reshape the diffusion prior or evaluate generated actions.Source: §5 · p. 6
  • Positive residual correlation lets the DICE value baseline act as a control variate and improve the stochastic-gradient constant.The bound groups approximation and optimization errors as Δ = ϵOT + ϵC + ϵSNIS + ϵV + δ.Source: §5 · p. 6
  • FedGuide satisfies a convergence guarantee under the stated assumptions when the policy-optimization stepsize is at most 1/(4L).The supplied passage states the guarantee's conditions but not its bound.Source: §5 · p. 6
  • Because policies and DICE value baselines remain client-local, the convergence bound has no policy-averaging heterogeneity term.Source: §5 · p. 6

6 Experimental Analysis

FedGuide and its variants generally achieve higher client-average returns across heterogeneous environments and remain robust in final-round and worst-round performance. The experiments also examine stronger heterogeneity, diffusion-prior mode preservation, and the value baseline’s association with smoother late-stage returns.

6 Experimental Analysis

  • The evaluation spans Bandit2D, Reacher, MuJoCo locomotion, and MetaWorld10, comparing FedGuide variants with federated policy-aggregation and policy-gradient baselines.FedGuide keeps the DICE value baseline local; FedGuide-P omits it, while FedGuide-A periodically averages client policies; all variants use OT-MoE for diffusion-prior aggregation.Source: §6 · pp. 6–7, Table 1 · p. 6

6.1 Experimental Results

  • FedGuide and its variants generally achieve higher client-average returns across tested environments and show robust final-round and worst-round performance.Figure 3 reports client-average returns over 100 rounds; Table 3 covers final-round and worst-round performance.Source: §6.1 · p. 7, Figure 3 · p. 8, Table 3 · p. 8
  • Under stronger heterogeneity, the FedGuide family attains the best final performance on Walker2D, HalfCheetah, and Reacher, while FedRL leads on Hopper-Hard.On Reacher-Hard, FedAvg and FedKL collapse to −103 returns and FedRL stalls at −100, while FedGuide variants converge to around 100.Source: §6.1 · p. 8, Figure 4 · p. 8
  • In Bandit2D, FedGuide’s variants capture all four reward peaks while FedAvg collapses to one, illustrating distributional mode preservation through OT-MoE.The toy setting has four isotropic Gaussian reward peaks evenly spaced on the unit circle.Source: Figure 2 · p. 7, §6.1 · p. 9
  • Using the value baseline usually lowers late-stage return volatility, while the caption cautions that this metric does not directly measure variance reduction.Table 2’s normalized temporal volatility is computed over the last 20 rounds; smaller values indicate smoother returns.Source: Table 2 · p. 7, §6.1 · p. 9

7 Conclusions

FedGuide combines OT-MoE aggregation of diffusion behavior priors with a DICE value baseline for local policy improvement. Its applicability is bounded by shared state-action spaces, meaningful overlap among client behavior priors, and OT-MoE and system-level constraints.

7 Conclusions

  • FedGuide aggregates diffusion behavior priors through OT-MoE and broadcasts personalized distribution-mixture priors to local clients.It uses a DICE value baseline for return-aware local policy improvement.Source: §7 · p. 9
  • The method assumes shared state-action spaces and meaningful overlap among client behavior priors; severe task-geometry or reward shifts may make aggregated priors misleading.Source: §7 · p. 9
  • OT-MoE scalability with many clients and unmodeled system-level heterogeneity—including adaptive client selection, communication delays, and compute constraints—remain limitations.Source: §7 · p. 9

A Theoretical Proofs

The appendix states regularity conditions behind FedGuide’s theoretical claims and analyzes its state-only DICE value baseline. The baseline does not add score-function bias and can reduce estimator variance when informative.

A.1 Assumptions

  • The analysis assumes smoothness, bounded gradient noise and local KL drift, an informative bounded DICE baseline, and bounded OT-MoE approximation errors.These are regularity conditions; the paper says A3 is enforced through local trust-region KL control, while A4 specifies when DICE is informative as a control variate.Source: §A.1 · pp. 12–13

A.2 Variance Reduction from the Value Baseline

  • A state-only baseline does not add action-dependent bias to the score-function estimator; finite-horizon GAE and imperfect value estimates remain approximation errors.The result concerns the score-function baseline term specifically.Source: §A.2 · p. 13
  • An informative DICE value baseline acts as a control variate and reduces variance for blending coefficients in the stated variance-reducing interval.The variance difference is strictly negative for some β in (0, 1), and the client-aggregated reduction has a positive coefficient.Source: §A.2 · pp. 13–14

A.3 Federated Convergence

The theoretical analysis gives convergence guarantees under stated assumptions and bounds when averaging prior-regularized local policies has limited impact. The environmental setup section describes the benchmarks and harder robustness settings.

A.3 Federated Convergence

  • Under A1–A5 and the stated stepsize condition, FedGuide has a nonconvex stationarity rate with controlled approximation terms.The approximation terms include OT and cost-proxy errors, prior and value surrogate errors, and local drift.Source: §A.3 · p. 14
  • Under the additional PL condition, FedGuide converges linearly to a neighborhood; the DICE baseline reduces the stochastic error floor, while other approximations contribute bounded terms.The neighborhood bound includes bounded contributions from OT-MoE, prior or value surrogates, and the local trust region.Source: §A.3 · p. 15

A.4 Policy Averaging under Prior-Regularized Alignment (FedGuide-A)

  • Policy averaging has limited impact when local policies remain distributionally aligned, and FedGuide-A is described as safe when its dispersion term Γπ is small.FedGuide’s OT-MoE prior aligns clients in distribution space while preserving behavior modes; the DICE baseline centers return-aware advantage estimates within that support.Source: §A.4 · pp. 16–17

B Environmental Setup

  • The environmental setup section introduces benchmarks, specifies client-level heterogeneity, and describes harder environments used to evaluate robustness under stronger heterogeneity.Its stated purposes include describing each benchmark’s relevance to heterogeneous FRL and how client heterogeneity is instantiated.Source: B Environmental Setup · p. 17

B.1 Environment Description

The benchmarks vary client MDPs through rewards, goals, dynamics, morphology, objects, or scenarios while sharing state and action spaces within each family. They range from a multimodal diagnostic bandit to goal-reaching, locomotion, and manipulation tasks.

B.1 Environment Description

  • Clients share state and action spaces within each benchmark family but differ in MDP components, producing distinct local distributions that test prior support preservation and local return-aware improvement.The configuration keeps each client’s policy, online critic, and frozen DICE baseline local; only diffusion-prior heads are uploaded for OT-MoE aggregation.Source: §B.1 · pp. 18–19
  • Bandit2D uses four separated Gaussian reward modes and client samples concentrated near one mode to illustrate distribution mismatch and mode preservation.This setup isolates how policy-space averaging can collapse modes or place mass in low-density regions.Source: §B.1 · p. 18
  • Reacher assigns clients different target positions and environment configurations, testing transfer without forcing a shared target-conditioned policy.Clients retain shared observation and action dimensions but have different reward landscapes and local state-action occupancy.Source: §B.1 · p. 18
  • Hopper, Walker2D, and HalfCheetah vary dynamics, morphology, contact, and gait, testing transferable locomotion priors alongside client-specific adaptation.Their client differences include environment parameters such as mass, damping, friction, actuation, target velocity, or reward coefficients.Source: §B.1 · pp. 18–19
  • MetaWorld10 assigns one manipulation task per client, creating differences in task geometry, object layouts, contact modes, and successful behavior distributions.The tasks have weaker overlap between local behavior supports; the setup evaluates whether distribution-space priors preserve multiple manipulation modes.Source: §B.1 · p. 19

B.2 Heterogeneity Setup

The main heterogeneity settings vary client metadata while retaining shared observation and action interfaces within benchmark families. Hard variants increase cross-client separation through stronger reward, structural, or dynamics perturbations.

B.2 Heterogeneity Setup

  • Main settings vary client goals and environment parameters in Reacher, locomotion dynamics and reward preferences, and task semantics in MetaWorld10.MetaWorld10’s ten clients represent distinct manipulation tasks; its heterogeneity includes object layout, contact sequence, success condition, and reward geometry.Source: §B.2 · pp. 19–20
  • Figure 5 encodes normalized dynamics and reward metadata for locomotion, goal, action-noise, reward, and angle-shift metadata for Reacher, and task-semantic attributes for MetaWorld10.The heatmap normalizes each continuous metadata column within its environment, emphasizing relative client variation rather than raw physical units.Source: §B.2 · p. 20, Figure 5 · p. 20

B.3 Heterogeneity Setup for Stronger Environments

  • Hard settings increase cross-client mismatch through categorical reward preferences, structural perturbations, or wider locomotion parameter ranges.Hopper-Hard and Walker2D-Hard use eight more separated dynamics regimes and more separated reward preferences; HalfCheetah-Hard widens perturbations and adds nonzero instability penalties.Source: §B.3 · p. 21

C Baseline Details

The baseline section contrasts parameter-, output-, actor–critic-, and gradient-based federation with FedGuide’s distribution-space prior aggregation and local value guidance. It then records the variants’ training settings and qualitative rollout visualizations.

C Baseline Details

  • The comparison uses matched client splits, rounds, rollout budgets, evaluation protocols, and seeds; Tables 4 and 5 organize methods by aggregation mechanism and analysis.The table descriptions contrast policy parameters, policy-output KL, actor–critic models, policy-gradient updates, and diffusion-prior heads.Source: C Baseline Details · pp. 21–22, Table 4 · p. 23, Table 5 · p. 24
  • FedAvg averages policy parameters, FedKL constrains policy-output drift with two KL penalties, FedRL uses a DDPGAvg-style actor–critic, and FedSVRPG-M aggregates variance-reduced gradient updates.FedKL’s penalties constrain drift from the server policy and between consecutive local updates; FedSVRPG-M uses momentum to reduce local gradient-estimation variance.Source: C Baseline Details · pp. 22–24
  • FedGuide keeps policies and value functions local, aggregating diffusion-prior heads through OT-MoE to return personalized mixture priors and using a blended DICE value baseline.The method aligns clients in distribution space rather than policy or gradient space, without sharing raw data, policies, or value functions.Source: Table 4 · p. 23, Table 5 · p. 24, C Baseline Details · p. 25

D Hyperparameters

  • Environment-specific hyperparameters are reported on the main page, while baseline implementations use PPO for FedAvg and FedKL, PPO/SVRPG for FedSVRPG-M, and DDPGAvg-style actor–critic for FedRL.Baseline hyperparameters are listed separately in Table 9.Source: D Hyperparameters · p. 25, §D.4 · p. 26, Table 9 · p. 27

D.1 Offline Pretraining

  • Clients pretrain behavior priors and DICE baselines from local offline buffers; the prior models local data support, and the frozen DICE baseline serves as a state-only baseline.Offline pretraining settings are listed in Table 6.Source: §D.1 · p. 25, Table 6 · p. 25

D.2 Online FedGuide Settings

  • Online PPO runs locally, with policies, critics, and DICE baselines kept local; only FedGuide and FedGuide-P communicate prior heads, while FedGuide-A also averages policies.Shared online settings are reported in Table 7.Source: §D.2 · p. 26, Table 7 · p. 26

D.3 FedGuide Variants

  • FedGuide variants share prior aggregation, client splits, rollout budgets, and PPO optimization; FedGuide-P removes DICE, while FedGuide-A additionally averages client policies.These variants isolate the value-baseline choice and policy averaging after personalized prior alignment.Source: C Baseline Details · pp. 24–25, §D.3 · p. 26

E Visualizations

  • Figures 7–11 show single final-round evaluation episodes across heterogeneous clients, with rows ordered FedGuide-A, FedGuide-P, then FedGuide and the best of five seeds selected for display.The visualizations cover Reacher, Hopper, Walker2D, HalfCheetah, and MetaWorld10; quantitative results instead use seed-averaged metrics.Source: E Visualizations · p. 26, Figure 7 · p. 28, Figure 8 · p. 29

Related papers

Loading 2609.18964v1…