Contribution
Proposes FedHSA, the first federated stochastic approximation algorithm that simultaneously handles agent heterogeneity (different local operator roots), Markovian data, and sparse communication, while achieving an M-fold linear speedup in sample complexity with no heterogeneity-induced bias term. Closes a fundamental open problem in federated reinforcement learning theory.
Key insight
Standard local-update federated SA algorithms converge to the wrong point under heterogeneity because multiple local steps using different local operators accumulate drift; injecting a global correction signal at each local step (the FedHSA update) cancels this drift exactly, enabling unbiased convergence even with Markovian noise and no projection step.
Assumptions
- SA operators G_i are contractive (strongly monotone) and Lipschitz- continuous; the global operator G-bar is also strongly monotone.
- Data at each agent is generated by an ergodic, aperiodic Markov chain with a unique stationary distribution; mixing time tau_mix is finite.
- Agents share a common finite state space; cross-agent Markov chains are independent.
- Agent heterogeneity is bounded: the discrepancy between local roots theta_i* and the global root theta* is finite.
- Communication occurs synchronously every H steps; agents are homogeneous in speed.
- Step sizes satisfy a two-timescale condition: local step eta decays faster than global step alpha_g.
Key results
- Proposition 1 (FedAvg-style SA bias under heterogeneity). Vanilla local-update federated SA with constant step size converges to a point that differs from the global root theta* by a non-vanishing bias proportional to the discrepancy between local operator roots and the global root. Holds when: Noiseless case (sigma=0); constant step size; heterogeneous agents with distinct local roots.
- Theorem 2 (FedHSA convergence). FedHSA converges exponentially fast to a ball around theta* with radius proportional to the step size; no heterogeneity-induced bias term appears in the bound. Holds when: Strongly monotone, Lipschitz SA operators; Markovian data; H local steps; two-timescale step sizes.
- Corollary 1 (Sample complexity with decaying step size). With linearly decaying global step size, FedHSA achieves sample complexity O-tilde(1/(MHT)) to reach epsilon-accuracy, giving M-fold linear speedup over single-agent SA. Holds when: Linearly decaying alpha_g; same operator and sampling assumptions as Theorem 2; complexity hides log factors and mixing time tau_mix.
Claims
| id | claim | strength | support |
|---|---|---|---|
| C1 | Vanilla local-update federated SA (FedAvg-style) fails to converge to the correct root under heterogeneity with a constant step size, even in the noiseless case. | strong | Proposition 1 gives an exact closed-form expression for the non-vanishing error as a function of the discrepancy between local roots and the global root. |
| C2 | FedHSA converges exponentially fast to a ball around the correct global root theta-star with no heterogeneity-induced bias, achieving M-fold linear speedup. | strong | Theorem 2 provides non-asymptotic bound; Corollary 1 shows sample complexity O-tilde(1/(MHT)) under linearly decaying step size. |
| C3 | FedHSA does not require a projection step to bound iterates, unlike competing heterogeneous federated SA algorithms. | strong | Algorithm description and proof in appendix explicitly avoid projection; a more involved analysis handles unbounded iterates directly. |
Method
FedHSA (Federated Heterogeneous Stochastic Approximation).
At each communication round, the server broadcasts the global parameter to all agents. Each agent performs H local SA steps using a corrected update: the noisy local operator is augmented by adding the current global average operator and subtracting the agent’s local operator evaluated at the round’s starting point. This correction term cancels the drift that would otherwise accumulate from heterogeneous local operators. After H steps, agents send their parameter deltas to the server, which averages them with a global step size to update the global parameter. Additionally, each agent sends its local operator evaluated at the new global parameter to enable the next round’s correction.
- Drift-correction term: G(theta-bar) - G_i(theta-bar, o_{i,0}) added to each local update
- Two-timescale step sizes: local step size eta and global step size alpha_g
- H local steps between communication rounds
- Server aggregation of agent deltas (not raw parameters)
Concepts
- Stochastic Approximation (SA) — Iterative algorithm for solving root- finding problems G-bar(theta)=0 using only noisy evaluations of G, covering SGD, TD learning, and Q-learning as special cases.
- Markovian sampling — Data at each agent is generated by an ergodic Markov chain rather than drawn i.i.d., introducing temporal correlations that slow convergence by a mixing-time factor.
- Heterogeneity-induced bias — Non-vanishing error term in convergence bounds of naive federated SA methods arising from differences between agents’ local operator roots and the global root.
- M-fold linear speedup — Sample complexity scales as 1/(MHT), meaning M agents collectively achieve the same accuracy as one agent using M times fewer samples per agent.
- Mixing time — Number of steps for a Markov chain to become epsilon- close in total variation to its stationary distribution; governs the penalty for Markovian versus i.i.d. data.
Connections
Builds on.
- Communication-Efficient Learning of Deep Networks from Decentralized Data (FedAvg) (LIT-307) — FedHSA extends the intermittent-communication federated learning paradigm to general nonlinear SA with heterogeneity correction.
- SCAFFOLD and related drift-correction methods in federated optimization — FedHSA adapts drift-mitigation ideas from federated optimization to the broader, nonlinear SA setting.
Recommendations
- R1 — Use FedHSA instead of FedAvg-style local SA when agents have heterogeneous environments (e.g., different MDPs in federated RL) to avoid convergence to the wrong fixed point. Topic: federated reinforcement learning · Strength: strong · When: Contractive/strongly monotone nonlinear SA operators; Markovian data at each agent; intermittent communication with a server.
- R2 — Use a linearly decaying global step size alpha_t to achieve the optimal O-tilde(1/(MHT)) sample complexity; constant step size leaves a residual bias proportional to alpha. Topic: step size scheduling · Strength: strong · When: When exact convergence to theta-star (rather than a neighborhood) is desired.
Bearing on the record
Naive local averaging converges to the wrong fixed point when agents have different environments. That is a correctness failure rather than a rate penalty, which makes it the more important half of the federated results here.
Limitations
- Requires agents to transmit their local operator evaluated at the new global parameter each round, adding one extra communication per round compared to vanilla FedAvg.
- Convergence guarantees require Lipschitz and strong-monotonicity assumptions that may not hold for general nonlinear RL operators.
- Agents’ Markov chains must share a common finite state space, limiting applicability to some continuous-state RL settings.
- Analysis assumes independent data across agents; correlated inter-agent data is not handled.
Open questions
- Can FedHSA-style bias correction be extended to settings where agents’ state spaces also differ?
- What is the minimum communication frequency H needed to maintain linear speedup as M grows?
- Does the mixing-time inflation factor tau/(1-rho) in the Markov bound represent the true lower bound, or can it be improved?