GAdapHeadFRL: Grouped Adaptive Head Mixing for Personalized Multi-Task Federated Reinforcement Learning

+24.5%

Relative gain over the strongest personalized MT-FRL baseline on MiniGrid MT7

50.0% → 72.9%

Converged-task ratio on MiniGrid versus the personalized baseline

>50%

Lower average performance standard deviation than centralized MT-RL references

12 / 50

Up to 12 MiniGrid tasks and 50 MetaWorld tasks across heterogeneous benchmarks

Abstract

Multi-task reinforcement learning (MT-RL) trains a single agent to solve multiple tasks by leveraging shared knowledge across task domains. However, most MT-RL methods assume centralized access to task data, making them impractical for privacy-sensitive or large-scale distributed settings. Multi-task federated reinforcement learning (MT-FRL) mitigates this issue by enabling agents to collaborate through shared models rather than raw trajectories, but often suffers from unstable performance under task heterogeneity, negative transfer, and imbalanced learning across tasks. A fundamental question in MT-FRL is how each agent should decide with whom to share skills and how strongly to rely on shared knowledge.

To answer this question, we propose GAdapHeadFRL, a personalized MT-FRL framework based on geometry-aware grouping and adaptive decision-layer task-head mixing. First, we decompose the local Bellman residual into representation and task-head estimation errors, and derive a closed-form optimal sharing strength governed by task mismatch and estimation uncertainty. Second, we instantiate this insight by identifying compatible client groups from local update geometry, constructing group-specific task-head prototypes, and dynamically mixing them with local heads using lightweight practical proxies.

Experiments on heterogeneous benchmarks, including MiniGrid (up to 12 tasks) and MetaWorld (up to 50 tasks), show that GAdapHeadFRL consistently outperforms state-of-the-art personalized MT-FRL baselines. It achieves up to 24.5% relative improvement on MiniGrid and increases the converged-task ratio from 50.0% to 72.9% over the personalized baseline. In addition, the proposed method reduces the average performance standard deviation by over 50% compared with centralized MT-RL references. These results demonstrate that geometry-aware grouping and adaptive task-head mixing provide an effective and scalable principle for heterogeneous MT-FRL.

Method

Whom to share with, and how much

A prevalent strategy for heterogeneous federated RL decouples the policy or value network into a shared representation backbone and task-specific decision heads. Existing methods then pick one of two extremes for the heads: fully shared heads maximize collaboration but risk negative transfer, while fully local heads avoid conflicts but forfeit decision-level knowledge transfer. GAdapHeadFRL moves beyond this binary choice along two axes.

Whom to share with (server). Clients send task signatures extracted from their local updates. The server clusters them with spherical K-means, automatically selects the number of groups using a cosine-silhouette criterion, and maintains one task-head prototype per group.

How much to share (client). Each client mixes its local head $\hat\phi_{i,k}$ with its group prototype $\hat\phi_{g,k}$ as $\tilde\phi_{i,k}(\alpha)=(1-\alpha)\hat\phi_{i,k}+\alpha\hat\phi_{g,k}$, where $\alpha_{i,k}\in[0,1]$ is computed locally from lightweight proxies of task mismatch (Mahalanobis distance) and estimation instability.

Theory: the optimal sharing strength is not a free hyperparameter

For a fixed backbone $\theta$, the local Bellman residual splits exactly into a representation error, which does not depend on $\alpha$, and a task-head estimation error that does.

Lemma (Orthogonal Decomposition). For any $\alpha\in[0,1]$ and fixed $\theta$, $$\varepsilon_{i,k}(\alpha)=\varepsilon_{\mathrm{rep},i,k}(\theta)+\big\|\tilde\phi_{i,k}(\alpha)-\phi_{i,k}^\star(\theta)\big\|_{\Sigma_i(\theta)}^2,$$ where $\varepsilon_{\mathrm{rep},i,k}(\theta)$ is independent of $\alpha$ and $\Sigma_i(\theta)$ is the feature covariance under the local distribution.

Modeling the local and shared heads as stochastic estimators, $\hat\phi_i=\phi_i^\star+m_i+\xi_i$ and $\hat\phi_g=\phi_g^\star+m_g+\xi_g$, with oracle mismatch $b_{i,g}=\phi_g^\star-\phi_i^\star$, noise energies $V_i,V_g$, and noise coupling $C_{ig}$, the expected head error becomes a strictly convex quadratic in $\alpha$.

Theorem (Optimal Adaptive Mixing). Under imperfect noise coupling ($|C_{ig}|\le\rho\sqrt{V_iV_g}$ for some $\rho<1$), the unconstrained minimizer is $$\alpha_{\mathrm{unc},i,k}=\frac{V_i-C_{ig}-\langle m_i,\;m_g+b_{i,g}-m_i\rangle_{\Sigma_i}}{\|m_g+b_{i,g}-m_i\|_{\Sigma_i}^2+V_i+V_g-2C_{ig}},$$ and the constrained optimum is $\alpha^\star_{i,k}=\Pi_{[0,1]}(\alpha_{\mathrm{unc},i,k})$.

Fully shared heads (FedAvg-style, $\alpha\to1$) and fully local heads (personalized, $\alpha\to0$) are the two boundary cases of this single adaptive principle: sharing is favored when task mismatch is small and local noise is large, and rejected when the mismatch dominates. Because the quality of the shared prior bounds what any mixing rule can achieve, grouping compatible clients first makes the prototype worth borrowing from.

Main Results

Success rate / normalized return (%) on MetaWorld and MiniGrid, reported as mean ± standard deviation. MT10, MT5 and MT7 average over the original split and two additional randomly sampled splits. The upper block is centralized multi-task RL; the lower block is federated RL.

Method	Venue	MetaWorld		MiniGrid
Method	Venue	MT10 (3Set)	MT50	MT5 (3Set)	MT7 (3Set)	MT12
Centralized multi-task RL
MOORE	ICLR	79.2 ± 5.2	72.3 ± 3.3	73.7 ± 8.0	74.3 ± 7.3	77.8 ± 6.0
McLean et al. (2025)	RLJ	80.9 ± 5.4	72.5 ± 3.5	68.0 ± 9.4	66.4 ± 6.8	64.6 ± 6.0
Federated RL
FedRL-SAC/PPO	DCS	78.1 ± 7.4	65.6 ± 3.1	57.6 ± 5.9	65.3 ± 4.5	68.7 ± 3.2
PFEDRL-REP	ICLR	78.1 ± 9.5	67.3 ± 1.0	66.2 ± 1.1	66.1 ± 8.2	80.0 ± 2.6
AdapHeadFRL (Ours)	–	80.7 ± 2.2	68.9 ± 0.9	79.1 ± 8.3	81.8 ± 2.7	85.9 ± 0.8
GAdapHeadFRL (Ours)	–	84.9 ± 3.6	72.0 ± 0.6	80.8 ± 6.1	82.3 ± 3.2	86.0 ± 0.9

MiniGrid learning curves

Average return on MiniGrid under two benchmark scales. Top: MT5 with the original split and two random splits (MT5-S1/S2). Bottom: MT7 with the original split and two random splits (MT7-S1/S2). Mean and 95% confidence interval over 30 independent runs. The dashed line is the centralized baseline MOORE. GAdapHeadFRL improves over PFEDRL-REP on MT7 from 66.1% to 82.3%, a 24.5% relative improvement.

Performance and stability

Mean performance versus standard deviation on MiniGrid MT12

Each point shows the mean success rate and standard deviation of a method. The bottom-right is the desirable regime of higher mean and lower variability. On MT12, GAdapHeadFRL reaches the best mean return with a standard deviation of 0.9, 85.0% lower than the strongest centralized MT-RL reference.

Task-level outcome distribution

Distribution of task-level outcomes on MiniGrid

Left: task instances grouped by peak normalized return. GAdapHeadFRL produces fewer failed cases and more converged cases than the baselines. Right: violin plots of peak returns across all task instances. Our methods shift the distribution upward and shrink the low-return tail.

Adaptive vs. fixed mixing weights

The adaptive coefficient outperforms every fixed $\alpha$. A small fixed ratio ($\alpha=0.2$) helps, but larger fixed values degrade performance, confirming that the optimal sharing strength must track task mismatch and estimation uncertainty during training.

When does selective sharing help?

Gain decomposition over PFEDRL-REP across heterogeneity levels

Gain over PFEDRL-REP on three MetaWorld MT10 sets ordered by update heterogeneity: $\Delta_\alpha$ from adaptive mixing, $\Delta_g$ additionally from grouping. Grouping matters most at medium heterogeneity (+6.5 points); adaptive mixing dominates at high heterogeneity (+6.0 points).

Practical coefficient vs. empirical oracle

On MiniGrid MT5 we freeze the backbone, local head, group prototype and Bellman targets, sweep $\alpha\in[0,1]$, and compare the practical coefficient with the critic-loss minimizer on an independent selection split.

Round	Convex (%)	Mean $\alpha_{\mathrm{oracle}}$	Mean $\alpha_{\mathrm{practical}}$	MAE
250	90	0.236	0.219	0.074
500	94	0.287	0.271	0.068

Deployment robustness

New-client assignment (MetaWorld MT10, held-out drawer-open/close). Joining the most compatible refreshed centroid reaches 0.8 return in 175k steps, 41.7% fewer than local-only adaptation.

Integration strategy	Early AUC ↑	Steps to 0.8 ↓	Final return ↑
Local only	0.5050	300k	0.9101
Global prototype	0.6041	275k	0.8761
Refreshed centroid	0.6272	175k	0.9968

Partial participation (MiniGrid MT12). Retaining the previous prototype whenever a group is absent keeps 98.5% of full-participation performance at 75% participation and 91.9% at 50%.

Participation	Final performance	Retained	RMST to 90% baseline
50%	79.0 ± 0.7	91.9%	160
75%	84.7 ± 0.5	98.5%	156
100%	86.0 ± 0.9	100%	66

Conclusion

We studied MT-FRL under structured task heterogeneity and proposed GAdapHeadFRL, which combines geometry-aware grouping with adaptive head mixing. By interpolating between local and shared task heads according to task mismatch and estimation uncertainty, it builds more compatible shared priors while avoiding negative transfer. On MiniGrid and MetaWorld, it consistently improves over federated baselines and remains competitive with strong centralized MT-RL methods, highlighting adaptive intermediate sharing as an effective principle for heterogeneous federated RL.

Two directions follow from this work. The analysis is exact for linear task heads and extends to nonlinear heads only locally, so a global characterization under nonlinear heads remains open. Group assignments are also fixed after warmup, whereas task relationships may drift as the shared backbone evolves, and periodic regrouping would extend the framework to non-stationary deployments.

BibTeX

@article{pang2026gadapheadfrl,
  title   = {Grouped Adaptive Head Mixing for Personalized Multi-Task Federated Reinforcement Learning},
  author  = {Pang, Yiran and Ni, Zhen and Pados, Dimitris and Zhong, Xiangnan},
  year    = {2026},
  url     = {https://yiranpang.github.io/GAdapHeadFRL/}
}