🪨 Asynchronous Federated Policy Gradient

Literature Review

Lan et al. [1] proposed AFedPG, the first asynchronous federated reinforcement learning framework with policy gradient updates and provable convergence guarantees.

• The framework addresses a fundamental limitation of synchronous federated RL:
  • in heterogeneous multi-agent systems, synchronous aggregation forces the server to wait for the slowest agent at every global step, making the time complexity dependent on the worst-case agent $t_{\max}$.
• AFedPG eliminates this bottleneck by allowing the server to update the global policy as soon as any agent completes its local computation, while other agents continue sampling and computing in parallel.
• The key technical challenge this introduces is that agents compute gradients based on stale policies:
  • the global model has been updated $\delta_k$ times since agent $i$ last received it,
• and in RL, unlike supervised FL, stale policies mean agents are collecting trajectories from a different distribution than the current global policy.

To handle this staleness, Lan et al. design a delay-adaptive lookahead technique specific to federated RL.

• At the $k$-th global iteration with delay $\delta_k$, instead of using the received model parameters directly, the agent computes a lookahead estimate $\tilde{\theta}_k \leftarrow \theta_k + \frac{1 - \alpha_{k-\delta_k}}{\alpha_{k-\delta_k}}(\theta_k - \theta_{k-1})$, where $\alpha_{k-\delta_k}$ is a delay-dependent learning rate.
• The agent then collects trajectories according to the policy $\pi_{\tilde{\theta}_k}$ and computes the updating direction.
  • This mechanism is designed to cancel out second-order Hessian correction terms $(1 - \alpha_{k-\delta_k})\nabla^2 J(\theta_k)(\theta_{k-1} - \theta_k) + \alpha_{k-\delta_k}\nabla^2 J(\theta_k)(\tilde{\theta}_k - \theta_k) = 0$ that arise uniquely in the RL setting due to policy-dependent data collection.
• On the server side, updates are normalised as $\theta_{k+1} \leftarrow \theta_k + \eta_k \frac{d_{k-\delta_k}}{|d_{k-\delta_k}|}$ to ensure controllable step sizes regardless of delay magnitude.

The convergence analysis establishes two results.

• For global convergence, AFedPG achieves $\mathcal{O}(\epsilon^{-2.5}/N)$ sample complexity at each agent, matching the single-agent rate of $\mathcal{O}(\epsilon^{-2.5})$ with a linear speedup in the number of agents $N$.
  • The time complexity improves from synchronous FedPG's $\mathcal{O}(t_{\max}/N)$ to $\mathcal{O}((\sum_{i=1}^{N} 1/t_i)^{-1})$
    • ... the harmonic mean of agent computation times, which is always less than or equal to $t_{\max}/N$ and significantly smaller when $t_{\max} \gg t_{\min}$.
  • The average delay $\bar{\delta}$ is shown to be naturally bounded by the average concurrency $\bar{\omega} \leq N$ without requiring any assumption on the maximum delay.

However, AFedPG assumes a centralised server, homogeneous agents running identical policy gradient algorithms, and identical environments across agents. Extending the delay-adaptive convergence analysis to decentralised aggregation with heterogeneous RL algorithms (e.g., SAC and PPO within the same federation), non-IID data across clusters, and off-policy learning remains an open problem .

References

[1] G. Lan, D.-J. Han, A. Hashemi, V. Aggarwal, and C. G. Brinton, "Asynchronous federated reinforcement learning with policy gradient updates: Algorithm design and convergence analysis," in Proc. Int. Conf. Learn. Representations (ICLR), 2025, arXiv:2404.08003.