Technical Summary: Collaborative Yet Personalized Policy Training via Single-Timescale Federated Actor-Critic

Problem Formulation

The paper addresses the challenge of collaborative yet personalized policy training in Federated Reinforcement Learning (FRL) under environmental heterogeneity. In many modern applications, such as Embodied AI, agents operate in distinct environments characterized by different transition kernels ($P^k$) and reward functions ($R^k$).

Existing approaches typically fall into two categories:

1. Ignoring Heterogeneity: Training a single shared policy across all agents, which often leads to performance degradation due to misaligned learning signals.
2. Abandoning Collaboration: Training entirely independent local policies, forfeiting the sample efficiency gains of collaboration.

While Personalized Federated Learning (PFL) has seen success in supervised learning, its application to FRL remains underexplored. Specifically, there is a lack of theoretical frameworks that handle personalized actor-critic (AC) methods under Markovian sampling (where data is temporally correlated) and heterogeneous transition kernels. Previous works often focus on value-based methods, common policies, or assume independent and identically distributed (i.i.d.) data, failing to capture the coupled dynamics of policy updates and heterogeneous Markov noise in a single-timescale setting.

Methodology: pFedAC

The authors propose pFedAC, a single-timescale personalized federated actor-critic algorithm. The core architectural innovation is a shared linear subspace representation combined with personalized local heads.

Model Structure

The value function for agent $k$ is approximated as:
$$V^k_\theta(s) \approx \phi(s)^\top B \omega^k_\theta$$
where:

• $\phi(s) \in \mathbb{R}^d$ is a known feature map.
• $B \in \mathbb{R}^{d \times r}$ is a common orthonormal subspace (shared across all $K$ agents) representing the underlying structure of the value functions.
• $\omega^k_\theta \in \mathbb{R}^r$ are local head parameters specific to each agent $k$, allowing for personalization.

Algorithm Mechanics

The algorithm operates in rounds $t = 0, \dots, T-1$ with three coupled update steps per agent, followed by a server aggregation step:

1. Critic Update (Local Head & Subspace):

   • Agents perform a local TD($L$) update to estimate the advantage function.
   • Using the TD error $\delta^k_{t,L}$, agents update their local head $\omega^k_t$ and the common subspace estimate $B_t$.
   • The update for $B_t$ involves a projected stochastic gradient step followed by a QR decomposition to maintain orthonormality. This projection is crucial for controlling perturbation amplification in the presence of heterogeneous noise.

2. Actor Update (Policy):

   • To mitigate the discrepancy between the limiting state distribution (used in critic updates) and the discounted state visitation distribution (used in policy gradients), the authors employ a companion MDP with a transition kernel $\tilde{P}^k = \gamma P^k + (1-\gamma)\eta$.
   • Agents sample trajectories from this auxiliary chain to compute policy gradients, ensuring more accurate gradient estimation under Markovian sampling.

3. Aggregation:

   • Agents send their updated subspace estimates $\bar{B}_{t+1}$ to the server.
   • The server computes the average $\bar{B}_{t+1}$ and performs a QR decomposition to broadcast the updated orthonormal basis $B_{t+1}$ back to all agents.

Key Contributions

1. First Personalized AC in Federated RL

The paper presents the first algorithm enabling personalized actor-critic learning in a federated setting. Unlike prior works that focus on shared policies or value functions, pFedAC explicitly learns a shared representation while maintaining agent-specific policy components.

2. Finite-Time Convergence Analysis

The authors establish finite-time convergence guarantees for the algorithm under Markovian sampling and heterogeneous transition kernels. The analysis introduces a novel joint linear approximation framework that simultaneously tracks three tightly coupled error dynamics:

• The error in the local head estimates ($\omega^k$).
• The error in the shared subspace estimate ($B$).
• The error in the policy gradient (Actor).

Convergence Rates:
The paper proves that:

• The critic error converges to zero at a rate of $\tilde{O}\left(\frac{1}{(1-\gamma)^4 \sqrt{TK}}\right)$.
• The policy gradient norm converges to zero at a rate of $\tilde{O}\left(\frac{1}{(1-\gamma)^6 \sqrt{TK}}\right)$.

Here, $T$ is the number of rounds, $K$ is the number of agents, and $\gamma$ is the discount factor. The notation $\tilde{O}$ hides polylogarithmic factors and problem-dependent constants.

3. Linear Speedup

The results demonstrate a linear speedup with respect to the number of agents $K$. This implies that adding more agents to the federated system proportionally reduces the sample complexity required to reach a target accuracy, even in the presence of heterogeneous Markovian trajectories.

4. Technical Innovations in Analysis

To achieve these results, the authors developed several new analytical tools:

• Perturbation Analysis for QR Decomposition: A new analysis bounding the distortion caused by QR decomposition steps in the presence of heterogeneous Markovian noise.
• Conditional Mixing Arguments: Techniques to handle the coupling between policy updates and Markovian sampling, specifically characterizing discrepancies between function evaluations under Markovian sampling versus temporally frozen policies.
• Geometric Structure Exploitation: Utilizing the geometric properties of the shared subspace and principal angle distances to control the interaction between local head errors and subspace estimation errors.

Experimental Results

The authors instantiated the framework within Proximal Policy Optimization (PPO), creating FedPer PPO, and evaluated it on a federated Hopper-v5 benchmark with action-map heterogeneity (where agents share dynamics but have different action mappings).

• Performance: FedPer PPO significantly outperformed both Single PPO (trained independently) and FedAvg PPO (standard federated averaging). In the "6-UNIQUE" setting (where every agent has a distinct action map), FedPer achieved a 2.63$\times$ improvement over Single PPO and 3.78$\times$ over FedAvg PPO.
• Transfer Learning: The learned shared "trunk" (the subspace $B$) was shown to be a transferable foundation. When the trunk was frozen and only new heads were trained on downstream tasks (including out-of-distribution variants), the model achieved 5.42$\times$ the performance of training from scratch on vanilla tasks and 3.50$\times$ on OOD tasks.

Significance and Claims

The paper claims that pFedAC successfully bridges the gap between collaboration and personalization in heterogeneous FRL environments. By leveraging a shared linear subspace, agents can benefit from collective learning without sacrificing adaptability to local environment characteristics.

The authors explicitly note that while the convergence rates demonstrate linear speedup, the dependence on the discount factor $(1-\gamma)$ may not be optimal, particularly in the near-undiscounted regime ($\gamma \to 1$). They identify improving this dependence as a key direction for future work, suggesting that sharper analyses of Markovian structure or variance reduction techniques could yield further improvements.

In summary, this work provides a rigorous theoretical foundation and empirical validation for personalized federated actor-critic methods, addressing the critical need for scalable, collaborative, yet adaptable reinforcement learning in complex, heterogeneous systems.