🤖 machine learning

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

FedGuide is a novel Federated Reinforcement Learning framework that addresses distributional mismatch in heterogeneous environments by aggregating diffusion priors via Optimal-Transport Mixture-of-Experts and employing a Distribution Correction Estimation value baseline to achieve robust, high-performance collaborative policy learning.

Original authors: Zhilin He, Gauri Joshi

Published 2026-09-17

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Zhilin He, Gauri Joshi

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

In the world of robotics, learning is often a solitary endeavor. A robot learns to walk or grasp an object by trial and error within its own specific environment, collecting data that is unique to its sensors, its motors, and the floor beneath its feet. This works well for a single machine, but it becomes a bottleneck when we want a fleet of robots to learn together. If every robot stays isolated, they cannot benefit from the successes or failures of their peers. The solution lies in a collaborative approach where machines share what they have learned without sending their raw, private data to a central server. This is the promise of federated learning: a way for distributed agents to build a shared intelligence while keeping their local experiences private. However, a major hurdle remains. When robots operate in different settings—perhaps one on a smooth factory floor and another on a rocky construction site—their experiences are fundamentally mismatched. Simply averaging their learned behaviors often results in a confused, ineffective strategy that fits neither environment well. The challenge is to find a way to combine these diverse experiences without forcing them into a single, compromised shape.

Researchers at Carnegie Mellon University have tackled this problem with a new framework called FedGuide, designed specifically for situations where robots face different physical realities. The core idea is to stop trying to average the robots' final decision-making rules, which often leads to a loss of important details. Instead, the team focuses on the underlying patterns of movement and action that each robot has discovered. They use a type of artificial intelligence model known as a diffusion model, which acts like a sophisticated memory of all the actions a robot has successfully taken in the past. Rather than sharing the robot's final policy, the system shares these "behavior memories." On a central server, these memories from different robots are carefully blended using a mathematical technique that respects the unique modes of behavior each robot has developed. This process ensures that a robot learning to walk on ice retains its specific caution, while a robot on dry pavement keeps its confidence, even as they learn from one another.

To make this collaboration even more effective, the researchers added a second layer of guidance. While the shared memories tell the robots what actions are possible, they do not necessarily say which actions are the most rewarding. To solve this, the team introduced a value estimator that acts as a local compass. This tool helps each robot understand how good a specific action is within its own unique environment, providing a stable signal that prevents the learning process from becoming erratic. By combining these shared behavioral memories with local, reward-aware guidance, the system allows each robot to improve its own performance without being dragged down by the differences in its peers' environments.

The researchers tested this approach across a variety of simulated worlds, ranging from simple tasks like reaching for a target to complex locomotion challenges involving hopping, walking, and running. In these tests, they compared their new method against standard techniques that simply average the robots' policies. The results showed a clear advantage for the new framework. In environments where the robots faced significant differences in their tasks or physical setups, the standard methods often struggled, sometimes failing to learn anything useful or collapsing into a single, poor strategy. In contrast, the new approach maintained high performance across all clients. It not only achieved higher final scores but also demonstrated greater stability, avoiding the wild swings in performance that often plague collaborative learning. Even in the most difficult scenarios, where the differences between robots were extreme, the system managed to keep every robot on a path toward improvement.

A key finding from the study is the importance of keeping the robots' specific behaviors distinct while still allowing them to learn together. When the researchers visualized the learning process, they saw that standard methods tended to force all robots into a single, average behavior pattern, effectively erasing the unique solutions each robot had found. The new method, however, preserved these diverse patterns. It allowed the system to recognize that there are multiple valid ways to solve a problem depending on the context. This preservation of diversity was crucial for robustness. The system did not just perform better on average; it was more reliable in the worst-case scenarios, ensuring that no single robot was left behind or forced to adopt a strategy that did not fit its reality.

The study also highlighted the specific roles played by the two main components of the system. The shared behavioral memories were essential for providing a foundation of safe, data-supported actions, preventing the robots from wandering into dangerous or impossible moves. The local value estimator was equally important, acting as a stabilizer that reduced the noise in the learning process. When the researchers removed either of these components, the system's performance dropped, confirming that both the shared memory of diverse behaviors and the local guidance on what is valuable are necessary for success. The results suggest that for robots to learn effectively in a distributed world, they need a way to share their experiences that honors their differences rather than smoothing them away.

This work represents a significant step forward in making collaborative robot learning practical for real-world applications. By moving away from the idea of a single, universal policy and toward a system that respects and integrates diverse local experiences, the researchers have shown that machines can learn together without losing their individual effectiveness. The framework provides a blueprint for how future fleets of robots, from warehouse workers to autonomous vehicles, might learn from each other in a way that is both efficient and robust. While the current tests were conducted in simulation, the principles demonstrated offer a clear path toward solving the fundamental problem of how to teach machines to work together when they live in different worlds. The success of this approach lies not in forcing uniformity, but in finding a way to align diverse perspectives into a coherent, collective intelligence.

Technical Summary: FedGuide

1. Problem Statement

Federated Reinforcement Learning (FRL) aims to enable collaborative policy learning across distributed agents operating in heterogeneous environments. However, existing FRL methods face significant challenges due to distributional mismatch. In standard FRL, clients generate local trajectories based on their specific Markov Decision Processes (MDPs), leading to policy-induced state-action distributions that differ across clients.

Current approaches typically address heterogeneity through variance reduction, divergence penalization, or momentum optimization. While these methods improve stability, they primarily rely on directly averaging policy or value-network parameters. This approach fails to explicitly model or align the underlying heterogeneous behavior distributions. Consequently, aggregating policies with mismatched trajectory distributions often leads to "mode collapse," where the global policy averages out distinct, valid local behaviors, resulting in ineffective or unsupported actions for specific clients.

2. Methodology: FedGuide

The authors propose FedGuide, a framework that shifts the aggregation interface from policy parameters to diffusion behavior priors. Instead of averaging policies, FedGuide aligns clients in the distribution space while keeping policy parameters local.

Core Components

Diffusion Priors as Behavior Models:
- Each client trains a state-conditioned diffusion model on its offline dataset to capture the multimodal distribution of actions (the "behavior prior").
- Unlike standard policies, these priors model what actions occur in the data, regardless of return, providing a data-supported action space.
Optimal-Transport Mixture-of-Experts (OT-MoE) Aggregation:
- The server maintains $M$ global diffusion prior experts.
- Instead of hard clustering or simple averaging, FedGuide uses an Optimal Transport (OT) plan to soft-route clients to experts based on the score-matching distance between client priors and expert priors.
- This results in a personalized distribution-mixture prior ( $\bar{\pi}_{D,i}$ ) for each client, which is a weighted combination of expert distributions. This preserves heterogeneous behavior modes in the distribution space without collapsing them into a single mode.
Distribution Correction Estimation (DICE) Value Baseline:
- To enable return-aware policy improvement within the support of the diffusion prior, FedGuide employs a state-only DICE value baseline ( $V_{\phi,i}$ ).
- This baseline is pre-trained on offline data to estimate the ratio between the improved policy distribution and the data distribution.
- Crucially, the DICE baseline is used only to center advantage estimates in the policy gradient, not to reshape the diffusion prior or act as a direct policy-gradient term. This avoids evaluating out-of-distribution actions and provides a stable, low-variance signal.
Local Optimization:
- Clients perform local updates using a Proximal Policy Optimization (PPO) objective.
- The loss function includes a prior regularizer that keeps the local policy close to the personalized diffusion mixture prior ( $\bar{\pi}_{D,i}$ ).
- The advantage estimator blends the online critic with the DICE value baseline to reduce variance.

3. Key Contributions

The paper outlines two primary contributions:

C1 (Formulation & Aggregation): The authors formulate heterogeneous FRL as local policy learning coupled by distribution-level diffusion priors. They develop an OT-MoE aggregation mechanism that routes client prior heads to server experts and broadcasts personalized distribution-mixture priors, avoiding the direct averaging of local policies.
C2 (Algorithm & Theory): They introduce a prior-regularized local PPO objective with a state-only DICE value baseline. Theoretically, they prove:
- Unbiasedness: The state-only baseline does not introduce score-function bias.
- Variance Reduction: The baseline reduces stochastic gradient variance when positively correlated with the online-baseline residual.
- Convergence: The method converges to an approximate stationary point of the federated surrogate with an explicit residual error floor determined by approximation errors in the OT cost, DICE estimation, and prior surrogates.

4. Experimental Results

FedGuide was evaluated across diverse heterogeneous FRL benchmarks, including Reacher, MuJoCo locomotion (Hopper, Walker2D, HalfCheetah), MetaWorld10, and a toy Bandit2D environment.

Performance: FedGuide consistently outperformed representative baselines (FedAvg, FedKL, FedRL, FedSVRPG-M) in client-average returns, final-round performance, and worst-round robustness.
Robustness: In "Hard" settings with increased heterogeneity (e.g., Reacher-Hard, Hopper-Hard), FedGuide maintained stability and convergence, whereas baselines like FedAvg and FedKL often collapsed or stagnated.
Ablation Studies:
- Diffusion Prior: Visualization on Bandit2D showed that FedGuide (using OT-MoE) successfully captured all multimodal reward peaks, whereas FedAvg collapsed to a single peak.
- Value Baseline: The inclusion of the DICE baseline significantly reduced late-stage return volatility (temporal variance), confirming its role in stabilizing local updates.
Limitations Observed: In MetaWorld10 with 10 clients, all methods showed only slight improvements with high variance, suggesting that FRL still struggles with complex heterogeneous task geometries beyond simple parameter shifts.

5. Significance and Claims

The paper claims that FedGuide addresses a fundamental gap in FRL: the lack of explicit distributional alignment among heterogeneous clients. By decoupling the aggregation of behavior distributions (via diffusion priors) from policy parameters, the framework allows clients to share knowledge about valid action modes without forcing them into a single, potentially invalid global policy.

The significance lies in:

Preserving Heterogeneity: Unlike parameter averaging which collapses modes, OT-MoE aggregation preserves distinct local behavior modes in the distribution space.
Stable Learning: The combination of a data-supported prior and a low-variance DICE baseline enables stable, return-aware learning even under strong heterogeneity.
Theoretical Grounding: The work provides theoretical guarantees on convergence and variance reduction, distinguishing it from heuristic approaches.

The authors modestly note that the method assumes shared state-action spaces and meaningful overlap among client priors. It may struggle under severe task-geometry shifts where aggregated priors become misleading, and it faces scalability challenges with OT-MoE as the number of clients grows.