Paper: Finite-Time Analysis of Federated Temporal Difference Learning with Linear Function Approximation under Environment and Computation Heterogeneity (ICLR-2024)

Reviewer comment (Experiments):
The simulations were conducted on the platform: GridWorld. Expanding simulations to more complex or diverse environments could provide a more comprehensive understanding of the algorithm's performance in practice.

Author response:
Based on your comment, in the revised paper, we have added simulation results in the setting of MountainCar-v1 from OpenAI Gym in the supplementary material. For the mountain car experiment, we first discretize the continuous state spaces with a uniform grid by tile coding. Each tiling has $10 \times 10$ grids. The action can be selected from {push left, no push, and push right}. The simulation results show that the running error is slightly larger than that of I.I.D. sampling in Gridworld with same heterogeneity level ${\kappa ^2} = 1$ . This is because in this setting, there exists discretization error. However, other results coincide with those in Gridworld.

Reviewer comment (Evaluation):
A deeper exploration of the algorithm's limitations or challenges, especially in real-world scenarios, might add depth to the research.

Author response:
The assumptions used in this paper are common in existing theoretical studies on RL (including FRL). Indeed, some of these assumptions would not hold in practice for real-world applications. For example, linear function approximation (also a common assumption in many existing works) assumes that the true value function can be well approximated by a linear function, using a fixed feature vector $\phi (s)$  for state $s$. However, in real-world settings, it is hard to find a good approximate feature matrix for all states. Also, in this paper, we assume that each agent participates in FRL in each round in a synchronous manner. However, in real-world settings, it can be more efficient for agents to participate in some but not all rounds of FRL in an asynchronous manner. For future work, we will explore the problem of this paper in more general settings by relaxing some of the assumptions used in this paper, such as using non-linear function approximation, and considering partial and asynchronous participation of agents.

Reviewer comment (Theory):
Would a real-world setting disrupt the assumptions required for analysis and negatively impact performance?

Author response:
The assumptions used in this paper are common in existing theoretical studies on RL (including FRL). Indeed, some of these assumptions would not hold in practice for real-world applications. For example, linear function approximation (also a common assumption in many existing works) assumes that the true value function can be well approximated by a linear function, using a fixed feature vector $\phi (s)$  for state $s$. However, in real-world settings, it is hard to find a good approximate feature matrix for all states. Also, in this paper, we assume that each agent participates in FRL in each round in a synchronous manner. However, in real-world settings, it can be more efficient for agents to participate in some but not all rounds of FRL in an asynchronous manner. For future work, we will explore the problem of this paper in more general settings by relaxing some of the assumptions used in this paper, such as using non-linear function approximation, and considering partial and asynchronous participation of agents.

Reviewer comment (Theory):
The average mean gradient $\bar{g}(\theta)$ doesn't equal to the gradient of Eq(6). This discrepancy arises because $\bar{A}$ is nonsymmetric, whereas a true gradient should be such that the hessian is symmetric. Although Eq. (6) defines the objective of the HFTD algorithm as minimizing the MSPBE, the algorithm actually converges to $\theta^*$, which satisfies the average mean gradient $\bar{g}\left(\theta^*\right)=0$. Hence, $\theta^*$ does not denote the MSPBE's minimum. As such, HFTD failed in finding the optimal value of the MSPBE.

Author response:
For each agent $i$, the function ${{\bar g_i}}(\theta )$ is similar to a "gradient" but is actually *not* the true gradient of any fixed objective function. This is because the "gradient descent" step, $g_i(\theta_t ) = (r_t + \gamma \phi {({s_{t + 1}})^{\rm{T}}}\theta  - \phi {(s_t)^{\rm{T}}}\theta )\phi (s_t)$, taken at time $t$ makes the estimated value ${V_\theta^i }(s_t)$ closer to the target $r_t^i + \gamma {V_\theta^i }({s_{t + 1}})$, but $r_t^i + \gamma {V_\theta^i }({s_{t + 1}})$ itself depends on $V_\theta^i $. When there is no confusion, we omit the subscript $i$ in the discussion below. It has been proven in [3] that the projected Bellman operator is a contraction operator as ${\left\| {{\Pi_D}{T_\mu }{V_{\theta_1} }} - {{\Pi_D}{T_\mu }{V_{\theta_1} }} \right\|_D}$ $\le$
 
$\gamma {\left\| {{V_{\theta_1} }-{V_{\theta_2} }} \right\|_D}$.

By repeatedly applying the contraction, one can obtain the approximate optimal value function $\Phi {\theta ^*}$ as

${\left\| {{V_{{\theta ^*}}} - {V_\mu }} \right\|_D}$ $\le$

 $\frac{1}{{\sqrt {1 - {\gamma ^2}} }}{\left\| {{\Pi_D}{V_\mu } - {V_\mu }} \right\|_D}$

It measures the root-mean squared deviation between the value predictions of the limiting TD value function and the true value function ${{V_\theta }}$. This means that, although ${{\bar g_i}}(\theta )$ is not the true gradient of the MSPBE ${F_i}(\theta ) $,
the value $\theta _i^*$ that satisfies  $\bar g_i(\theta _i^*) = 0$ is still the optimal solution that minimizes $F_i(\theta)$. Moreover, a powerful ODE method was provided in [3] using the stochastic approximation $\dot \theta  = {{\bar g}_i}(\theta ) = {{\bar A}_i}\theta  + {{\bar b}_i}$ where ${{\bar A}_i}$ is negative definite and symmetric. Thus, this ODE is exponentially stable and $\theta$ will converge to a globally and asymptotically stable equilibrium point ${\theta_i ^*}$ where $\theta _i^* =  - \bar A_i^{ - 1}{{\bar b}_i}$. Therefore, although $\bar g(\theta )$ is not the true gradient of (6), $\bar g({\theta ^*}) = 0$ is still the optimal solution to (6) (Eq.(5) in the revised paper).

Reviewer comment (Writing):
The paper's main body contains multiple conflicting statements. The abstract mentions a mixed environment as the average of $N$ heterogeneous environments. However, in Sec 4, it's conveyed that this environment is randomly drawn from the heterogeneous environments. These descriptions are contradictory.

Author response:
Based on your comment, in the revised paper, we have revised the abstract to clarify that the "mixture of environments" is the environment *randomly drawn* from agents' heterogeneous environments (marked in red in the revised paper). This *mixture* environment is different from the *single virtual* environment $\( S,A,\bar P,\bar R,\gamma\)$ defined in [5-6], which is obtained from averaging the transition kernels and reward functions of agents' environment as $\overline P  = \frac{1}{N}\sum\nolimits_{i = 1}^N {{P_i}} $ and $\bar R = \frac{1}{N}\sum\nolimits_{i = 1}^N {{R_i}} $. The goal of this paper is to find the optimal value function model for the mixture environment of agents. This goal is similar in spirit to that of federated supervised learning, as the latter aims to find the optimal model that minimizes the average training loss for all data samples of all clients.

Reviewer comment (Experiments):
The authors didn't provide a fair comparison between the results of their paper and those of existing works. In table (1), the objective of Wang (2023) was to find the optimal value function of $i$-th agent, which more focused on the personalization. However, in this paper, the objective function is to find the optimal value function of all $N$ environments, which was in the average sense. If the authors changed the optimality to $i$-th agent's optimal value function, their methods still cannot converge exactly without the gradient heterogeneity in Assumption 2.

Author response:
We compare the results of this paper with those of related existing works as follows. The goal of Wang (2023) was to find the optimal value function for the *single virtual* environment $\( S,A,\bar P,\bar R,\gamma\)$  where $\overline P  = \frac{1}{N}\sum\nolimits_{i = 1}^N {{P_i}} $ and $\bar R = \frac{1}{N}\sum\nolimits_{i = 1}^N {{R_i}} $. They obtained the personalization result $\mathbb{E}\left\| {{V_{\theta_T} } - {V_{{\theta_i} ^*}}} \right\|_{\bar D}^2$. It is bounded by $T_1$ and $T_2$, where $T_1$ denotes 

$\mathbb{E}\left\| {{V_{\theta_T} } - {V_{\theta ^*}}} \right\|_{\bar D}^2$, and $T_2$ denotes

$\mathbb{E}\left\| {{V_{{\theta_i} ^*} } - {V_{{\theta ^*}}}} \right\|_{\bar D}^2$,


where $\bar D$ is the stationary distribution of the averaged MDP $\( S,A,\bar P,\bar R,\gamma\)$; ${\theta ^*}$ is the optimal solution to the value evaluation problem of the averaged MDP. Nevertheless, $T_1$ and $T_2$ both involve some non-vanishing terms. Specifically, [6] focused on the term $T_1$. In handling per round progress by using the update rule, the inner product includes $\left\| {\bar g(\theta ) - \frac{1}{N}\sum\limits_{i = 1}^N {{{\bar g}_i}(\theta )} } \right\|$ where ${\bar g({{\bar \theta^t }})}$ is the expected gradient of the *virtual* MDP. It is non-zero, seeing Lemma 2 in [6]. That is the reason why $T_1$ has non-vanishing terms. For $T_2$, it is non-vanishing due to environment heterogeneity as $\theta _i^* \ne {\theta ^*}$. 

Compared to [6], in this paper, we consider the \textit{mixture} environment which is the environment \textit{randomly drawn} from agents' heterogeneous environments, rather than the single virtual environment [6] mentioned above. Based on this setting, we propose the HFTD algorithm which aims to iteratively find the optimal value function model that minimizes the average MSPBE of agents' heterogeneous environments. Therefore, the objective function fo the FRL problem in this paper is different from that in [6]. We show that the HFTD algorithm can asymptotically achieve the optimal model in the settings of this paper. We do not claim that the HFTD algorithm can asymptotically achieve the optimal model for the learning problem in [6]. Also, the HFTD algorithm would not asymptotically achieve the local optimal model for any agent due to environment heterogeneity.

Reviewer comment (Writing):
The presentation of the paper can be improved with a lot of grammar issues and inaccurate statements (see Questions below for details). For example, there shall be citations for the statement "Hence, inspired by FL, federated reinforcement learning (FRL) has been proposed as a promising approach which allows agents to collectively learn the policy without sharing raw data samples." In the paper, only one citation on federated RL is provided (Wang et al. arXiv 2023). It is difficult to claim that FRL has been a "promising" approach. Overall, I am not convinced by merging the FL setup with RL. If every agent faces heterogeneous MDPs, what is the goal of the FRL?

Author response:
A1: RL faces severe sample inefficiency in real systems. For example, when RL is applied to train a walking robot, the walking ability of the robot is limited by the number of samples used by the robot. A natural idea to deal with the sample-efficiency issue of RL is federated RL (FRL) which allows agents to share the information so that agents can collectively learn a better policy. This policy may perform slightly worse than the optimal policy for each local agent, but it can possess a broad applicability to effectively address all local objectives. With federation, agents can accelerate the process of learning based on other's experience. 

We have presented thorough investigation on FRL with theoretical guarantees [1-7] in the part of federated learning as an independent subsection in Related Works. Besides, we have added citations for the statement "Hence, inspired by FL, federated reinforcement learning (FRL) has been proposed as a promising approach which allows agents to collectively learn the policy without sharing raw data samples, to address the sample inefficiency of RL in real systems" [1-7]. We marked the revisions in red.

Reviewer comment (Novelty):
The algorithm is a combination of the FL and TD learning algorithm; and the analysis also follows mostly from existing literature.

Author response:
A2: The HFTD algorithm proposed in this paper is much more than a simple combination of FL and TD learning. There are three major differences compared to the existing works on FRL [1-7], discussed as follows.

1) The FRL settings considered in this paper have non-trivial differences from those in prior works on FRL. Compared to parallel RL where agents interact with identical environments [1, 2, 6], FRL where agents have heterogeneous environments remains largely under-explored. In a few recent works on FRL [3-4] for agents with heterogeneous environments, the objective is to find the optimal value function model for a *single virtual* environment $\( S,A,\bar P,\bar R,\gamma\)$ obtained from averaging the transition kernels and reward functions of agents' environment as $\overline P  = \frac{1}{N}\sum\nolimits_{i = 1}^N {{P_i}} $ and $\bar R = \frac{1}{N}\sum\nolimits_{i = 1}^N {{R_i}} $. However, such an ``averaged'' virtual environment may not coincide with the individual environment of any of the agents for FRL. In this paper, we consider a *mixture* environment defined as the environment *randomly drawn* from agents’ heterogeneous environments. Note that this mixture environment is different from the virtual environment [3-4] mentioned above. Thus, our goal is to find the optimal value function model for this mixture environment of agents. This goal is similar in spirit to that of federated supervised learning, as the latter aims to find the optimal model that minimizes the average training loss for all data samples of all clients.

2) The HFTD algorithm can asymptotically achieve the optimal value function model in the FRL settings considered in this paper. In existing works on FRL for heterogeneous environments [3-4], the bound of the convergence error is non-vanishing. In this paper, to our best knowledge, we provide the first result in existing works on FRL with heterogeneous environments that the convergence error can diminish to zero asymptotically. This result is because that the goal of the HFTD algorithm is to find the optimal value function model for the mixture environment that is randomly drawn from agents’ heterogeneous environments. To this end, the algorithm aims to minimize the average of agents’ mean squared projected Bellman errors (MSPBEs) for their respective environments. In particular, a key property of the global gradient of the average MSPBE allows us to remove a non-vanishing bias in the convergence error, so that only vanishing error terms remain. We also show that the HFTD algorithm achieves a sample complexity of ${\cal{O}}\left( {\frac{1}{{N{\varepsilon ^2}}}} \right)$ and a linear convergence speedup, which match the results of existing TD learning algorithms [6, 8].


3) The HFTD algorithm allows agents to use different numbers of local iterations based on their capabilities of collecting samples. The existing works on FRL [1-4] assume that agents use the same number of local iterations in each round of the algorithm. However, in practice, agent's capabilities of sample collection and local computation can be highly diverse, so that stragglers can significantly slow down the training process. Therefore, it is more efficient for agents to use different numbers of local iterations based on their respective capabilities.

Reviewer comment (Writing):
There are quite a lot of typos/grammar issues; See e.g., "each agents in FRL collects" "none of these works have shown" “design a federated TD algorithm that asymptotically converge to” “how much is the sample complexity of this algorithm? ” "We provide numeral results"

Author response:
A3: In the revised paper, we have corrected all grammatical errors and typos (marked in red in the revised paper). For example, "each agents in FRL collects" is changed to "each agent in FRL collects", "none of these works have shown" is changed to "none of these works has shown", "design a federated TD algorithm that asymptotically converge to" is changed to "design a federated TD algorithm that asymptotically converges to", "we provide numeral results" is changed to "we provide numerical results".

Reviewer comment (Experiments):
The paper claims that using federated multiple agents to collaboratively evaluate the policy can accelerate the convergence achieving linear speedup. It would be great if numerical results can be provided to demonstrate this linear speedup under heterogenous settings.

Author response:
A4: Indeed, these numerical results were all conducted under the heterogeneous setting as Fig 1. Furthermore, Fig 1(c) demonstrates this linear speedup under heterogeneous settings. Following your suggestion, we have added more details of heterogeneous settings. Assumption 2 defines the environment heterogeneity shown as bounded gradient heterogeneity, $\frac{1}{N}\sum\nolimits_i {\left\| {{{\bar g}_i}(\theta )} \right\|_2^2}  \le {\beta ^2}\left\| {\frac{1}{N}\sum\nolimits_i {{{\bar g}_i}(\theta )} } \right\|_2^2 + {\kappa ^2}$ where ${\kappa ^2}$ denotes the heterogeneity level. Here we conduct the simulations with ${\kappa ^2}=1$. Based on the fact that ${{\bar g}_i}(\theta ) = {\Phi ^{\rm{T}}}{D_i}(\gamma {P_i}\Phi  - \Phi )\theta  + {\Phi ^{\rm{T}}}{D_i}{R_i}$, a series of heterogeneous transition probability kernel and reward function are generated $\( {P_i},{R_i}\) _{i = 1}^N$ to satisfy Assumption 2. After that, these environments are executed on GridWorld. Fig 1(c) shows that one achieves a linear speedup with respect to the number of workers which matches our theoretical result ${\mathbb{E}\left\| {{{\bar \theta }^T} - {\theta ^*}} \right\|_2^2 \le {\cal{O}}\left( {{e^{ - \sqrt {NKT} }}} \right) + {\cal{O}}\left( {\frac{{N(K - 1)}}{T}} \right) + {\cal{O}} \left( {\frac{1}{{\sqrt {NKT} }}} \right)}$. Typically, when $T$ is sufficiently large, the total sample complexity for achieving an $\epsilon$-accurate optimal solution ${\mathbb{E}}\left\| {{{\bar \theta }^T} - {\theta ^*}} \right\|_2^2 \le \varepsilon $ is $KT = {\cal{O}}\left( {\frac{1}{{N{\varepsilon ^2}}}} \right)$.