Pith. sign in

REVIEW 4 major objections 5 minor 28 references

Heterogeneous Federated Reinforcement Learning Using Wasserstein Barycenters

T0 review · 4 major / 5 minor · reviewed 2026-08-06 · deepseek-v4-flash

Challenge this review Re-run · record.json Download PDF Read on arXiv ↗

Pith's one-line read Fusing neural networks with Wasserstein barycenters produces a single DQN that balances poles across environments of different lengths.

desk verdict The core aggregation step in Algorithm 1 has a sign error that makes the algorithm unimplementable as written, despite a genuinely new HFRL application. read the letter →

arxiv 2506.15825 v1 pith:WAOCKBLD submitted 2025-06-18 cs.LG

Luiz Pereira , M. Hadi Amini This is my paper
classification cs.LG
keywords FederatedlearningHeterogeneousreinforcementWassersteinbarycenterModelfusionOptimaltransportDeepQ-networkCartPole
verification ladder T0 review T1 audit T2 compute T3 formal

The pith

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

The reading

This paper tries to establish that Wasserstein barycenters, the geometric averages of probability distributions, can serve as a model-fusion rule in federated learning and that this rule extends to heterogeneous federated reinforcement learning. It presents FedWB, which flattens each agent's weight matrices, shifts them by their minima, scales them to unit sum, computes the Wasserstein barycenter of the resulting probability vectors, and rescales back to form a global network. In a distributed MNIST experiment, the fused global model reaches high accuracy with fewer epochs than FedAvg in early training. In CartPole with different pole lengths, the global DQN produced by FedWB controls all environments and sustains a longer balance time early on, while FedAvg catches up later. The paper claims this geometric averaging preserves environment structure that plain arithmetic averaging destroys, and it recommends using FedWB early and FedAvg later.

What carries the argument

The machine is the Wasserstein barycenter in one dimension, computed after flattening each layer's weight matrix into a vector, shifting by its minimum, and normalizing to unit sum. This conversion turns each network's weights into a probability distribution so that the barycenter, the distribution minimizing the sum of squared Wasserstein distances to all agents' distributions, acts as the aggregation operator. The un-normalization step uses the arithmetic means of the per-agent scaling sums and minima. The argument relies on the barycenter's tendency to place mass along a geodesic between the input distributions, which the paper reads as preserving geometric structure rather than splitting mass bimodally as arithmetic averaging does.

What would settle it

Take a single trained network, randomly permute the neurons of a hidden layer, and run Algorithm 1 on the original and permuted versions as if they were two agents. If the two fused outputs differ materially, the method is sensitive to unit ordering that real networks possess. Alternatively, run FedWB against FedAvg and against a neuron-aligned optimal-transport fusion across several heterogeneous RL tasks with larger environment differences; if FedWB consistently underperforms either baseline after equal compute, the claimed early advantage is refuted.

Watch

Extended reading notes

Core claim

The central discovery is that a Wasserstein barycenter of flattened, shifted, and normalized weight vectors is a workable global model in both federated supervised learning and federated reinforcement learning. The algorithm treats each layer's weight matrix as a one-dimensional distribution by flattening, adding the matrix's minimum to make entries nonnegative, and dividing by the sum; the barycenter of these distributions is then un-normalized using the averaged minimum and sum across agents. In MNIST, this global model exceeds 90% test accuracy, and it improves faster in early epochs than FedAvg. In CartPole with heterogeneous pole lengths, the fused global DQN balances poles across all environments, and by the end of 600 episodes both FedWB and FedAvg reach the 500-action maximum, with FedWB holding the early advantage. The paper frames the result as a speed-versus-accuracy tradeoff and suggests a hybrid schedule.

Load-bearing premise

The load-bearing premise is that after flattening each weight matrix, shifting by its own minimum, and scaling to unit sum, the Wasserstein barycenter of the resulting probability vectors is a meaningful average of the networks, and that un-normalizing with the global average of the minima and sums recovers valid network weights without any neuron or permutation alignment across agents.

Editorial extensions

If this is right

  • If FedWB works as claimed, federated training can produce a global model by geometric averaging of weight vectors alone, without aligning neurons across agents, and the global model matches arithmetic averaging in the long run.
  • The reported early advantage implies a practical hybrid: run FedWB for the first part of training to get a strong starting point, then switch to FedAvg to finish cheaply; the paper suggests this directly.
  • Because the aggregation applies to Q-networks, the same procedure should be applicable to any DQN-based multi-agent system facing environment heterogeneity.
  • The one-dimensional reduction keeps the barycenter computation tractable, but the paper reports that it is slower than averaging, so the gain in early accuracy must be weighed against aggregation time.

Reading between the lines

Editorial extensions of the paper, not claims the author makes directly.

  • The paper's flatten-then-average construction is sensitive to the ordering of neurons: permuting the units of a hidden layer changes the flattened vector and therefore the barycenter, so the method is not permutation invariant even though the underlying network function would be unchanged. This is an implicit property the paper does not examine.
  • A natural extension would be to compare FedWB against neuron-aligned optimal-transport fusion on the same CartPole suite; the paper's design does not include that comparison, so whether the barycenter's benefit survives alignment is untested.
  • The claimed benefit could be stress-tested by increasing environment heterogeneity beyond pole length, for example by varying observation or reward structure; the paper only varies one scalar parameter.
Share X Bluesky LinkedIn Reddit HN

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, and a circularity audit.

Referee Report

4 major / 5 minor

Summary. The paper proposes FedWB, a federated aggregation method that replaces arithmetic averaging with a Wasserstein barycenter of flattened, min-shifted, sum-normalized weight matrices. It applies FedWB to distributed MNIST classification and to heterogeneous federated deep Q-learning on CartPole environments with different pole lengths, claiming faster early convergence than FedAvg and a global Q-network that works across environments.

Significance. If correct, FedWB would be a useful extension of Wasserstein-barycenter model fusion to heterogeneous federated reinforcement learning. The problem is well motivated, and the abstract claim of faster early convergence is falsifiable and is compared against FedAvg. However, the central aggregation step contains an algebraic sign error, is not derived or compared with the aligned OT-fusion methods it builds on, and the experiments confound aggregation rule with network architecture and omit key details. The reported results therefore do not currently support the paper's claims.

major comments (4)
  1. [Section 4.1, Algorithm 1] Algorithm 1, lines 3–6 and line 9: the normalization and un-normalization steps are internally inconsistent. Line 4 computes W := W + min(W), which does not prevent negative numbers; it shifts entries downward (e.g., for weights [−1, 100] the result is [−2, 99]) and can make the scaling sum zero or negative. The correct nonnegativity transform would be W − min(W), and the inverse in line 9 would then add the averaged minimum, not subtract it. As written, line 6 does not produce a probability vector for general signed weights and line 9 does not invert lines 4–6. Since Algorithm 1 is the core of both the MNIST and CartPole pipelines, the experiments cannot be reproduced from the pseudocode.
  2. [Section 4.1, Algorithm 1] The paper never justifies treating each flattened weight matrix as a 1D probability distribution and taking its Wasserstein barycenter as the global model. No neuron or permutation alignment is performed across agents, and no comparison is made with the aligned model-fusion method of [8] or the layerwise barycenter method of [9]. The entire FedWB result therefore rests on an unvalidated heuristic that is load-bearing for every experimental claim.
  3. [Section 5.1, Figure 5] The comparison between FedWB and FedAvg is confounded by architecture. Section 4.1 states that FedAvg uses two hidden layers of 200 nodes each, while FedWB uses a single hidden layer of 256 nodes. Thus Figure 5 compares different networks, not different aggregation rules, and the claimed early-convergence advantage cannot be attributed to the Wasserstein barycenter.
  4. [Section 5.2, Algorithm 3, Figures 6–7] The HFRL experiment omits essential details: number of agents, pole-length values, aggregation frequency C, replay buffer size, learning rate, and the epsilon-greedy annealing schedule are not reported. No multiple-seed statistics or error bars are shown, and the text does not report per-environment performance of the global DQN. The abstract's claim that a single global DQN functions across all environments is therefore not demonstrated by the presented evidence.
minor comments (5)
  1. [Section 3.2.1, Eq. (7)] The matrix P is called a permutation matrix, but equation (6) defines a transport polytope or coupling, which is not a permutation matrix.
  2. [Section 2.2, Eq. (5)] The update rule theta = wT * theta + (1 - wT) * theta is a tautology; a term involving the mature model's parameters appears to be missing.
  3. [Figure 5 caption] The caption lists both normal training and FedAvg as the 'red line'; one of the curves is likely misdescribed.
  4. [Figure 7 caption] The caption uses 'FHRL' instead of 'HFRL'.
  5. [Algorithms 1–3] The loop structure of Algorithm 1 is underspecified: lines 2–6 use W(i)_d with no enclosing loop over layers, while line 8 later loops over i. Please clarify the intended order of operations.

Circularity Check

0 steps flagged · score 0.0 of 10

No significant circularity: FedWB defines an explicit aggregation rule and its claims are tested against FedAvg and standard training, not derived from the method's own output.

full rationale

The paper's derivation chain is self-contained and not circular. FedWB's global model is obtained from a concrete, stated recipe: flatten each weight matrix, apply a scaling transformation, compute a Wasserstein barycenter, and invert the scaling (Algorithm 1, lines 1-11). The method is benchmarked against FedAvg and ordinary single-agent training on MNIST and CartPole, so the headline comparisons are external rather than quantities built from the method itself. No fitted parameter is renamed as a prediction; Table 1 is an explicitly labeled extrapolation from measured epochs and single-agent wall-clock time, not a claim derived from FedWB. The cited model-fusion works [8] and [9] are prior external work by other authors, and no load-bearing self-citation or imported uniqueness theorem appears. The paper's references do not include the present authors' prior work, so self-citation is absent. Separately, Algorithm 1's normalization step appears mathematically inconsistent with its stated goal of preventing negative numbers (adding the minimum rather than subtracting it), and this is a substantive correctness risk; however, a sign error or ill-defined normalization is not a circularity, because the algorithm does not redefine its output as a restatement of its inputs. No circular step can therefore be exhibited.

Assumptions & free parameters 6 free parameters · 4 assumptions · 0 invented entities

The central claim rests on a heuristic normalization and barycenter computation that is not derived, on unstated hyperparameters, and on the assumption that flattened weight distributions can be averaged without permutation alignment. No new physical or mathematical entities are introduced.

free parameters (6)
  • Per-layer minimum shift and scaling sum = Data-dependent per agent; not learned
    Algorithm 1 uses each agent's min and sum to convert weights to probability vectors; the global un-normalization uses their averages, a choice not derived.
  • Hidden layer width = 256 for FedWB, 200 for FedAvg
    The compared architectures differ in size, confounding the MNIST speed and accuracy comparison.
  • Epsilon-greedy annealing schedule = Not reported
    The DQN exploration schedule strongly affects CartPole training but is never specified.
  • Target network update frequency C = Not reported
    Algorithm 3 requires a value for C but none is given.
  • CartPole pole lengths = Not reported
    Heterogeneity is created by varying pole lengths, but the actual lengths are never stated.
  • Wasserstein barycenter solver = Not reported
    The core aggregation step depends on an OT solver, but no solver, discretization, or implementation detail is described.
assumptions (4)
  • ad hoc to paper Flattened, shifted, and scaled weight matrices can be treated as probability distributions whose 1D Wasserstein barycenter is a valid fusion of neural networks
    Algorithm 1 normalizes weights to histograms and applies WB; no proof or reference supports this exact procedure.
  • ad hoc to paper No permutation or neuron alignment is required across agents before averaging weights
    Algorithm 1 flattens weight matrices in place and never aligns neurons, despite known permutation symmetries in neural networks (contrast with [8]).
  • standard math The Bellman target and Huber loss are valid for DQN training
    Section 3.3 invokes standard Q-learning and DQN updates.
  • ad hoc to paper Averaging Q-network parameters across heterogeneous MDPs yields a policy that works across all environments
    The HFRL design assumes a single global Q-network can handle all pole lengths; this is tested only on a toy problem without theoretical support.

how reviews work

0 comments
Cite this review

Pith. "Pith review of Heterogeneous Federated Reinforcement Learning Using Wasserstein Barycenters." pith.science (2026). https://pith.science/paper/WAOCKBLD

@misc{pith2026250615825,
  author       = {Pith},
  title        = {Pith review of: Heterogeneous Federated Reinforcement Learning Using Wasserstein Barycenters},
  year         = {2026},
  howpublished = {\url{https://pith.science/paper/WAOCKBLD}},
  note         = {Machine review of arXiv:2506.15825}
}
read the original abstract

In this paper, we first propose a novel algorithm for model fusion that leverages Wasserstein barycenters in training a global Deep Neural Network (DNN) in a distributed architecture. To this end, we divide the dataset into equal parts that are fed to "agents" who have identical deep neural networks and train only over the dataset fed to them (known as the local dataset). After some training iterations, we perform an aggregation step where we combine the weight parameters of all neural networks using Wasserstein barycenters. These steps form the proposed algorithm referred to as FedWB. Moreover, we leverage the processes created in the first part of the paper to develop an algorithm to tackle Heterogeneous Federated Reinforcement Learning (HFRL). Our test experiment is the CartPole toy problem, where we vary the lengths of the poles to create heterogeneous environments. We train a deep Q-Network (DQN) in each environment to learn to control each cart, while occasionally performing a global aggregation step to generalize the local models; the end outcome is a global DQN that functions across all environments.

Figures

Figures reproduced from arXiv: 2506.15825 by the authors.

Figure 1
Figure 1. Heterogeneous Federated Reinforcement Learning (HFRL) - Model Architecture. Model architecture for [PITH_FULL_IMAGE:figures/full_fig_p003_1.png] view at source ↗
Figure 2
Figure 2. Compact NN design. 3.2 Preliminaries of Optimal Transport 3.2.1 General Theory Optimal transport (OT) deals with the idea of how to move mass from one location to another in a transportation-cost￾minimizing fashion. While the original work was proposed by Monge [20], the more well-known formulation was proposed by Kantorovich [21], where the idea of mass-splitting allowed for broadening of the applications of OT. Th… view at source ↗
Figure 3
Figure 3. The number of epochs needed to reach at least 90% accuracy given a number of agents in the distributed [PITH_FULL_IMAGE:figures/full_fig_p009_3.png] view at source ↗
Figures from the paper (4 more)
Figure 4
Figure 4. Figure 4: Accuracy per number of epochs. Results yielded from an architecture using 10 agents. The black line is the [PITH_FULL_IMAGE:figures/full_fig_p010_4.png]
Figure 5
Figure 5. Figure 5: Model comparison between normally training with one agent (red line), FedWB (black line), and FedAvg (red [PITH_FULL_IMAGE:figures/full_fig_p011_5.png]
Figure 6
Figure 6. Figure 6: Result comparison between FedWB and FedAvg in training a global DQN for the CartPole problem. [PITH_FULL_IMAGE:figures/full_fig_p011_6.png]
Figure 7
Figure 7. Figure 7: Difference in the "duration" per epoch for FedWB based and FedAvg based FHRL. [PITH_FULL_IMAGE:figures/full_fig_p012_7.png]

Discussion (0). Sign in to comment.

Reference graph

Works this paper leans on

28 extracted references · 18 canonical work pages

  1. [9]

    Wasserstein barycenter-based model fusion and linear mode connectivity of neural networks

    Aditya Kumar Akash, Sixu Li, and Nicolás García Trillos. Wasserstein barycenter-based model fusion and linear mode connectivity of neural networks. arXiv preprint arXiv:2210.06671, 2022

    arXiv 2022

  2. [8]

    Model fusion via optimal transport

    Sidak Pal Singh and Martin Jaggi. Model fusion via optimal transport. Advances in Neural Information Processing Systems, 33:22045–22055, 2020

    2020

  3. [1]

    Deep reinforcement learning from human preferences

    Paul F Christiano, Jan Leike, Tom Brown, Miljan Martic, Shane Legg, and Dario Amodei. Deep reinforcement learning from human preferences. Advances in neural information processing systems, 30, 2017

    2017

  4. [2]

    Reinforcement learning from simulated environments: An encoder decoder framework

    Benjamin Choo, Graham Crannel, Stephen Adams, Faraz Dadgostari, Peter A Beling, Ann Bolcavage, and Roy McIntyre. Reinforcement learning from simulated environments: An encoder decoder framework. In 2020 Spring Simulation Conference (SpringSim), pages 1–12. IEEE, 2020. 12

    work page 2020

  5. [3]

    Federated reinforcement learning: Techniques, applications, and open challenges

    Jiaju Qi, Qihao Zhou, Lei Lei, and Kan Zheng. Federated reinforcement learning: Techniques, applications, and open challenges. arXiv preprint arXiv:2108.11887, 2021

    arXiv 2021

  6. [4]

    Release strategies and the social impacts of language models

    Irene Solaiman, Miles Brundage, Jack Clark, Amanda Askell, Ariel Herbert-V oss, Jeff Wu, Alec Radford, Gretchen Krueger, Jong Wook Kim, Sarah Kreps, et al. Release strategies and the social impacts of language models. arXiv preprint arXiv:1908.09203, 2019

    arXiv 1908

  7. [5]

    Mastering the game of go with deep neural networks and tree search

    David Silver, Aja Huang, Chris J Maddison, Arthur Guez, Laurent Sifre, George Van Den Driessche, Julian Schrittwieser, Ioannis Antonoglou, Veda Panneershelvam, Marc Lanctot, et al. Mastering the game of go with deep neural networks and tree search. nature, 529(7587):484–489, 2016

    2016

  8. [6]

    Multi-task learning with attention for end-to-end autonomous driving

    Keishi Ishihara, Anssi Kanervisto, Jun Miura, and Ville Hautamaki. Multi-task learning with attention for end-to-end autonomous driving. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 2902–2911, 2021

    work page 2021

Show all 28 references
  1. [7]

    Communication- efficient learning of deep networks from decentralized data

    Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas. Communication- efficient learning of deep networks from decentralized data. In Artificial intelligence and statistics , pages 1273–1282. PMLR, 2017

    2017

  2. [10]

    Federated machine learning: Concept and applications

    Qiang Yang, Yang Liu, Tianjian Chen, and Yongxin Tong. Federated machine learning: Concept and applications. ACM Transactions on Intelligent Systems and Technology (TIST), 10(2):1–19, 2019

    2019

  3. [11]

    Federated learning with differential privacy: Algorithms and performance analysis

    Kang Wei, Jun Li, Ming Ding, Chuan Ma, Howard H Yang, Farhad Farokhi, Shi Jin, Tony QS Quek, and H Vincent Poor. Federated learning with differential privacy: Algorithms and performance analysis. IEEE Transactions on Information Forensics and Security, 15:3454–3469, 2020

    2020

  4. [12]

    Federated reinforcement learning with environment heterogeneity

    Hao Jin, Yang Peng, Wenhao Yang, Shusen Wang, and Zhihua Zhang. Federated reinforcement learning with environment heterogeneity. In Gustau Camps-Valls, Francisco J. R. Ruiz, and Isabel Valera, editors,Proceedings of The 25th International Conference on Artificial Intelligence ...

    2022

  5. [13]

    Federated reinforcement learning: Linear speedup under markovian sampling

    Sajad Khodadadian, Pranay Sharma, Gauri Joshi, and Siva Theja Maguluri. Federated reinforcement learning: Linear speedup under markovian sampling. InInternational Conference on Machine Learning, pages 10997–11057. PMLR, 2022

    2022

  6. [14]

    Federated reinforcement learning acceleration method for precise control of multiple devices

    Hyun-Kyo Lim, Ju-Bong Kim, Ihsan Ullah, Joo-Seong Heo, and Youn-Hee Han. Federated reinforcement learning acceleration method for precise control of multiple devices. IEEE Access, 9:76296–76306, 2021

    2021

  7. [15]

    Semi-distributed resource management in uav-aided mec systems: A multi-agent federated reinforcement learning approach

    Yiwen Nie, Junhui Zhao, Feifei Gao, and F Richard Yu. Semi-distributed resource management in uav-aided mec systems: A multi-agent federated reinforcement learning approach. IEEE Transactions on Vehicular Technology, 70(12):13162–13173, 2021

    2021

  8. [16]

    Federated reinforcement learning for energy management of multiple smart homes with distributed energy resources

    Sangyoon Lee and Dae-Hyun Choi. Federated reinforcement learning for energy management of multiple smart homes with distributed energy resources. IEEE Transactions on Industrial Informatics, 18(1):488–497, 2020

    2020

  9. [17]

    Deep learning

    Yann LeCun, Yoshua Bengio, and Geoffrey Hinton. Deep learning. nature, 521(7553):436–444, 2015

    2015

  10. [18]

    Mini-batch gradient descent: Faster con- vergence under data sparsity

    Sarit Khirirat, Hamid Reza Feyzmahdavian, and Mikael Johansson. Mini-batch gradient descent: Faster con- vergence under data sparsity. In 2017 IEEE 56th Annual Conference on Decision and Control (CDC) , pages 2880–2887. IEEE, 2017

    2017

  11. [19]

    Introduction to multi-layer feed-forward neural networks

    Daniel Svozil, Vladimir Kvasnicka, and Jiri Pospichal. Introduction to multi-layer feed-forward neural networks. Chemometrics and intelligent laboratory systems, 39(1):43–62, 1997

    1997

  12. [20]

    Mémoire sur la théorie des déblais et des remblais

    Gaspard Monge. Mémoire sur la théorie des déblais et des remblais. Mem. Math. Phys. Acad. Royale Sci., pages 666–704, 1781

  13. [21]

    On the transfer of masses: Doklady akademii nauk ussr

    L Kantorovich. On the transfer of masses: Doklady akademii nauk ussr. 1942

    1942

  14. [22]

    Computational optimal transport: With applications to data science

    Gabriel Peyré, Marco Cuturi, et al. Computational optimal transport: With applications to data science. Founda- tions and Trends® in Machine Learning, 11(5-6):355–607, 2019

    2019

  15. [23]

    Genesis of bimodal distributions

    Isidore Eisenberger. Genesis of bimodal distributions. Technometrics, 6(4):357–363, 1964

    1964

  16. [24]

    Reinforcement learning: A survey.Journal of Machine Learning Research, 18(153):1–90, 2017

    Lihong Li, Wei Zhang, Li Zhang, Yinyu Zeng, Qi Yu, and Shuo Yang. Reinforcement learning: A survey.Journal of Machine Learning Research, 18(153):1–90, 2017. 13

    2017

  17. [25]

    Optimal control theory: an introduction

    Donald E Kirk. Optimal control theory: an introduction. Courier Corporation, 2004

    2004

  18. [26]

    Robust estimation of a location parameter

    Peter J Huber. Robust estimation of a location parameter. Breakthroughs in statistics: Methodology and distribution, pages 492–518, 1992

    1992

  19. [27]

    Yann LeCun, Corinna Cortes, and Christopher J.C. Burges. Gradient-based learning applied to document recognition. Proceedings of the IEEE, November 1998

    1998

  20. [28]

    Human-level control through deep reinforcement learning

    V olodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Marc G Bellemare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, et al. Human-level control through deep reinforcement learning. nature, 518(7540):529–533, 2015. 14

    2015

Pith tools

Reviewed August 6, 2026 · model on record in the stance chip above.