Mean Field Langevin Actor-Critic: Faster Convergence and Global Optimality beyond Lazy Learning

Thank you for your helpful comments. We address the technical comments below.

Regarding Finite MDPs

The global optimality should be analyzed in finite MDPs. Is the global optimality guaranteed in finite MDPs?

Finite MDPs represent a simpler subset within our MDP framework and are encompassed within our settings. Notably, our analysis stands out in its capacity to handle continuous state and action spaces. Additionally, in cases with a finite number of action spaces, the policies in MFLPG can be characterized by following a softmax function to probabilistically select actions. Research in reinforcement learning on MDPs, particularly policy gradient methods, has been actively pursued recently, and analyzing global optimality within such settings remains a significant topic (Leathy et al. (2022), Zhang et al. 2021, 2022).

If some states are not reached, how can we guarantee global optimality?

In this study, we consistently regard the expected value on the visit distribution as what to get easily. Specifically, we derive results based solely on population updates, not relying on stochastic gradient descent, which aligns with recent studies (Leathy et al. (2022), Zhang et al. 2021, 2022). Moreover, the completeness of sampling subsets of states crucial for estimating the optimal policy is aggregated within the moment regularity condition, outlined in Assumption 4, related to the visitation distribution based on the optimal policy and the distribution at each time step.

Regarding the Motivation and Superiority of Experimentation in Section 5

The experiment is only on CartPole.

The computational experiments presented in Section 5 serve the explicit purpose of addressing two key objectives: (1) To experimentally demonstrate the necessity of analyzing network representation learning in the mean-field regime surpassing existing lazy-scaling theories like NTK concerning optimization and generalization capabilities, and (2) To showcase the practical utility of our newly proposed method, the Double-Loop MFLTD, by comparing it practically with the widely used, more basic TD(1). It's essential to note that our study's focal point lies in conducting convergence analyses of TD learning and PG methods in the mean-field regime, capturing feature learning without relying on lazy-scaling, and deriving rapid convergence rates under mild assumptions. In general, analytical studies of deep learning theory, especially within reinforcement learning, lack computational experiments or are confined to simplified reward structures, table transition settings, or classic toy models like Cart Pole. While showcasing the effectiveness of our algorithms in larger models or complex environments with vast state spaces would offer significant insights, considering the current stage with theoretical convergence analyses, we find our experimental settings adequate. But we certainly acknowledge that expanding the models and settings remains a future research direction.

There should be an experiment on over-parameterized cases.

The positioning of the computational experiments in our study aligns with the explanation provided above and holds meaningful implications even for experiments on MF networks that are not entirely over-parameterized. Particularly, differences in performance stemming from the representational power between MF networks and NTK networks are evident and verifiable even within finite-width MF networks. Further details regarding this point are directed to the response for reviewer GTpe (Benefits of feature learning in the context of this RL).

Why are the networks over-parameterized?

The use of over-parameterized neural networks, such as those in mean-field theory and NTK, does not derive from the motivation that increasing the number of network parameters inherently improves performance. Over-parameterization is an approach aimed at providing theoretical guarantees for neural network optimization and is driven by demands from the theoretical side. Therefore, while achieving theoretical guarantees with non-over-parameterized finite-width MF networks would be an ideal goal, the analysis of neural networks remains incomplete even on supervised learning, particularly in the context of reinforcement learning, where analyses often remain limited to linear function approximation like NTK or more powerful assumptions. Hence, the deep learning theory community is actively exploring statistical and optimization properties beyond NTK. Our research contributes as a significant step within this array of research directions in reinforcement learning.

(continued below)

Thank you for your helpful comments. We address the technical comments below.

Regarding Assumption 5; some practical example

Thank you for your insightful remarks regarding Assumption 5 in our manuscript. Indeed, Assumption 5 plays a crucial role in bridging the gap within the proof. Sufficient increase of the parameters $R$ and $M$ ensures the richness of $\mathcal{F}\_{R, M}$. On the other hand, it's worthing to note that Assumption 5 guarantees that when one has a policy, one can always approximate the advantage function in the gradient direction of the policy gradient within the finite KL-divergence ball. Indeed, for example, Assumption 5 is satisfied when $A_{\pi}/\|A_{\pi}\|\_{\sigma_{\pi}}\in \mathcal{F}\_{R, M}$. Now that $Q_{\pi}\in\mathcal{F}\_{R, M}$ is assumed by Assumption 2, the geometric regularity of Assumption 5, coupled with the richness of the function class $\mathcal{F}\_{R, M}$, is a further generalization of the above example $A_{\pi}/\|A_{\pi}\|\_{\sigma_{\pi}}\in \mathcal{F}\_{R, M}$, with the assumption relaxed by an amount corresponding to $B>0$.

Addressing questions:

I appreciate your invaluable feedback, which prompted several clarifications and adjustments in our manuscript:

1. Please see Theorem 1 in Sutton et al. (1999). Regarding the definition of $d^{\pi}$ in Sutton et al. (1999)'s Theorem 1, it is indeed associated with the visitation distribution in their paper. While they initially define $d^\pi$ in the context of the stationary distribution, they later redefine it as the visitation distribution just before Theorem 1. This transition might cause confusion, but it is crucial to note that within Theorem 1, $d^{\pi}$ refers to the visitation distribution. This alignment is evident from the proof. Consequently, for our Proposition 1, similar consideration of expectations concerning the visitation distribution is necessary, leading to the current statement. Moreover, you are correct in stating that when the initial distribution is the stationary distribution, the visitation distribution coincides with the stationary distribution, naturally fulfilling one of the conditions in Assumption 4.
2. I have rectified the confusion regarding dimensions by modifying it to $\mathcal{S}\times\mathcal{A}\subset \mathbb{R}^{D}$.
3. The correction has been made to reflect that particles range from 1 to $m$: $[1,m]$.
4. Regarding the inner optimization objective function $\mathcal{L}\_l$ in MFLTD, it has been emphasized that this function is convex over the probability distribution space and that the values of the inner objective function at its stationary point is larger than a constant multiple of the true objective function, the expected squared error $\frac{1-\gamma}{2}\mathbb{E}{\varsigma_{\pi}}[(Q_q - Q_{\pi})^2]$, with an error of at most $\mathcal{O}(\lambda_{TD})$. This implies a continuous reduction of $\frac{1-\gamma}{2}\mathbb{E}_{\varsigma_{\pi}}[(Q_q - Q_{\pi})^2]$ with each inner update. To enhance clarity, I have revised the respective section for better comprehension.
5. As per your observation, I have rectified notation inconsistencies and modified the definition of the second $L_l[q]$.
6. The erroneous reference to $q_*^{(l+1)}$ instead of $Q_*^{(l+1)}$ has been corrected.
7. In our study's setup, although we initially maintained $\lambda_{TD}$ as a constant across all steps, your insight on adapting $\lambda_{TD}$ as a function of the inner optimization time $s\in[0, S]$, such as $\lambda_{TD} = 1/\log(s)$, offers a means to progressively decrease biases in the convergence values of the objective function. This adjustment aims to converge asymptotically towards the true optimal value.
8. Your suggestion regarding additional references is duly noted. I will include the relevant citation in the specified section.
9. Apologies for any confusion caused by overloaded notations. I have revised the notation system to a clearer and unambiguous format to mitigate any potential misunderstandings.
10. In MFLPG (Algorithm 1), in each policy evaluation step, we uniformly sample $l\in[T_{TD}]$ and adopt $Q^{(l)}$ as $Q_t$ from the estimated Q-functions $\{Q^{(l)}\}\_{l\in[T_{TD}]}$ obtained by MFLTD (Algorithm 1). Then the error in the value function using $Q_t$ in the sense of expected value is upper-bounded by the error on the right-hand side of Theorem 1. Supplementary explanation is added in the paper.

For assumption 5, do you mean the choice of f could depend on $\pi$? If this is the case, then please clarify. What I understand from Assumption 5 was that the choice of f is independent of $\pi$. Question 1 was my mistake, thank you for the clarification.

Thank you for your helpful comments. We address the technical comments below.

Hard to verify assumptions, which weaken the aim of supporting empirical success by theoretical footing.

It is unclear how practical those proposed actor-critic methods are. The authors argued that those assumptions are moderate, but the work also does not show the proposed methods are still close enough to what have been used in practice.

Supplementary Information on Assumptions' Validity and Limitations

Each assumption in our study aligns closely with the commonly used assumptions in theoretical analyses of neural reinforcement learning. They are quite familiar within this domain.

Assumptions 1 and 3 effectively dictate the richness of the neural network function class induced by determining the activation functions within neural networks. Our coverage encompasses a wide range of neural networks employing typical activations, such as sigmoid or hyperbolic tangent functions applied component-wise. These assumptions are prevalent in other analytical studies, like Agazzi & Lu (2021), Leathy et al. (2022), and Zhang et al. (2021, 2022). Additionally, Assumption 1 is validated in cases where kernels are smooth and light-tailed, like the RBF kernel as mentioned in Suzuki et al. (2023) for applications such as MMD and KSD estimation. Assumption 3 naturally arises from the richness of the neural network class introduced by Assumption 1. As detailed in Appendix B.1, this is an inherent deduction from the benefits of the Barron class' universal function approximation capability, among other factors.

Assumption 5 ensures that when a policy is present, approximating the advantage function within a finite KL-divergence ball in the gradient direction of the policy gradient is always possible. Specifically, for instance, when $A_{\pi}/\|A_{\pi}\|\_{\sigma_{\pi}}\in \mathcal{F}\_{R, M}$. Given $Q_{\pi}\in\mathcal{F}\_{R, M}$ by Assumption 3, the moderate geometric regularity of this assumption, coupled with the richness of the function class $\mathcal{F}\_{R, M}$, is evident. This assumption is a relaxed version accommodating $B>0$, as elaborated in the mentioned example.

We've supplemented and complemented the main text with additional clarifications, realistic examples of meeting these assumptions, and more. If there are further questions or doubts regarding specific assumptions' limitations or validity, please feel free to inquire.

Insights on the Practicality of Proposed Algorithms

Section 5 of our study offers valuable insights into the practicality of our proposed algorithms. For instance, our experiments demonstrate that while MFLTD constitutes a double-loop algorithm, it performs comparably to a general TD(1), displaying even higher performance accuracy. This emphasizes the practical adequacy of our algorithm in obtaining more accurate Q-functions.

How do we justify widely used actor-critic methods in practice are learning representations in the same way of how the proposed methods learn?

"Representation learning" is essentially the advantage of using neural networks as learning machines themselves. In contrast to kernel methods, multi-layered neural networks (having two or more trainable layers) achieve more efficient learning by aligning features based on the data used for learning. In our study, we capture the learning dynamics of widely used neural network representations through the Mean-Field (MF) regime. This understanding aligns with a firmly established "common knowledge" within the community, especially within the context of general supervised learning, where learning dynamics of neural networks in the MF regime partially mirror those utilized in practical scenarios.

For instance, considering MFLPG as policy gradient where Gaussian noise is added each gradient update in the discrete-time algorithm, this noise inclusion leads to an optimization under Langevin dynamics, ensuring analytical proof. That is, this addition does not fundamentally alter the learning of network representations from a viewpoint of representation learning. Consequently, real-world neural networks inherently possess a learning mechanism quite similar to those handled in our study. This fact contributes to supporting the assertion that widely used actor-critic methods in practice learn representations in a manner akin to the proposed methods.

The learning of "representations" in the implementation of individual reinforcement learning algorithms, while crucial, is not straightforward to elucidate. There remains ample scope for further, more rigorous research into the specific implications of implicit "representation learning" through neural networks within the realm of reinforcement learning. For more insight into neural network representation capabilities, please refer to the response to reviewer GTpe.

We sincerely appreciate your time and effort in reviewing our manuscript and providing valuable feedback. We have taken them into serious consideration for improving the clarity and coherence of our work.

Benefits of feature learning in the context of this RL

As you rightfully pointed out, the Mean-field (MF) scaling approach utilized in our study, in comparison to the Neural Tangent Kernel (NTK) scaling prevalent in existing research, covers a more diverse function class by aligning features along the data used for learning. In Section 5, we performed a comparison between NTK scaling and MF scaling using finite-width three-layer neural networks. Admittedly, extrapolation from the analysis results of our study focused on infinite-width neural networks may require further consideration in this section. By employing finite-width MF networks, we aim to elucidate the benefits of feature learning, specifically highlighting three aspects: (1) the performance difference stemming from the representational capabilities of MFLTD and NTK-TD, (2) Sample complexity with respect to the dimension, (3) Reinforcement learning unique feature learning trends:

(1) The performance difference due to representational capabilities.

• Primarily, NTK scaling employs fixed feature bases throughout training, only learning the second layer based on these bases, and its expressive power relies on the number of available bases determined by the network's width. In contrast, in MF scaling, the bases corresponding to the first layer's features can adequately adapt to the data used for learning.
• As a specific example, to represent the polynomials of few relevant dimensions in a high (, say $d$-,) dimensional space, superpolynomial width is required in the NTK scaling because with high probability the first layer weight of each neuron does not align with such relevant directions. On the other hand, for the MF scaling, the required width does not depend on the ambient dimension $d$ (although here we ignored an optimization aspect). Detailed discussion can be found in Damian et al. (2022).Like this, in terms of representation ability, the function class realized by finite-width NTK networks is sometimes clearly distinguished from that by the NTK limit, and less able to approximate functions common in learning theory than those by MF networks with finite width (although the function class of finite-width NTK networks are not completely included in that of finite-width MF networks). Therefore, the class of finite width networks in Cai et al. (2019) is not as flexible as that of finite width MF networks.
• We wanted to emphasize not only the order itself but also the definition is important to discuss the convergence rates. The measurement of convergence depends on with which the MSE is measured (, which we called global minima). For a fixed finite width, the global minima in MF can be significantly closer to the true function than that in NTK as we argued above. Although our network is infinite width, this is mainly not for representation but for optimization, and it would be possible to convert our results to the finite width setting.

(2) Sample complexity with respect to the dimension.

The MF network with infinite width also has advantage over the NTK network with infinite width, in terms of sample complexity in high-dimensional problems.

• If you look at the learnt network in MF, the first layer parameters typically align with the important directions. But not in NTK, because NTK weight does not allowed to travel so much.
• As a result, the MF learnt network typically has low-dimensionality, and therefore requires fewer samples. Specific examples are parity function, polynomials of few relevant dimensions, low dimensionality, hierarchical functions, and single index models. For example, in order to learn the $2$-parity in a $d$-dimensional space, $\mathcal{O}(d^2/\varepsilon)$ sample is required for NTK but $\mathcal{O}(d/\varepsilon)$ for MF (Telgarsky (2023)) to achieve the accuracy. However, the dimension dependency will be absorbed by just looking $T$ or $T_{TD}$.
• This is not because MF has less (or more) representation ability and the whole hypothesis class is smaller (or larger), but because the parameters can gain the aligned structure a result of training (=feature learning).

(continued below)

I appreciate your thorough evaluation of our work and have made specific adjustments based on your invaluable suggestions:

authors should try to make an enumerated list to summarize all definitions…

I am sincerely grateful for your feedback. In line with your recommendation, I intend to enhance the readability of the manuscript by compiling a supplementary document summarizing essential symbols used in proofs and theorems.

The paper does not show how the actor and critic updates are combined together.

We would like to emphasize that our main result Theorem 2 evaluate the errors induced by both algorithms by combining the error analysis for each algorithm. Indeed, you can find that the MFLAC in Algorithm 1 entails running the policy evaluation (MFLTD, Algorithm 2) at each update of policy improvement. Then, the error given by Theorem 2 yields the statement under the condition of $T_{\mathrm{TD}} = O(1/\lambda_{\mathrm{TD}})$ for the number of iterations in the inner loop.

The discrete time analysis of MFLPG, unlike MFLTD, is missing.

Our research analyzed the continuous-time dynamics of MFLPG. Expanding this analysis to discrete-time dynamics can naturally extend from the continuous-time framework. This extension, dependent on the learning rate $\eta>0$, converges toward the optimal policy reward within an error margin, as observed in recent studies on time and particle discretization (Suzuki et al. (2023)).

The technique used to show linear convergence of inner loop MFLD does not seem to be new compared to previous established works.

Impact on the paradigm of analysis of TD learning of NNs in MFLTD

Acknowledging your point, the uniqueness of stationary points in the inner loop of MFLTD and its global optimality leverage previously established convergence proofs. However, the contribution of MFLTD in our study lies in proposing this dual-loop algorithm and deriving non-asymptotic results regarding the convergence of its outer loop. As described on page 5, introducing a naive semi-gradient method using gradient descent on the optimal problem in the mean-field regime, i.e., as an optimal problem on the probability distribution space, cannot be achieved. Specifically, the Fokker-Planck equation $\partial_t \rho_t = \lambda\cdot\Delta \rho_t + \nabla\cdot\left(\rho_t \cdot v_t\right)$, where $v_t := \nabla\mathbb{E}{\varsigma_{\pi}}[(Q_q-\mathcal{T}^{\pi}Q_q)h_{\theta}]$, does not converge. This is due to the fact that the naive semi-gradient on the probability subspace $\mathbb{E}{\varsigma_{\pi}}[(Q_q-\mathcal{T}^{\pi}Q_q)h_{\theta}]$ does not consistently align with the true oracle gradient direction $\mathbb{E}{\varsigma_{\pi}}[(Q_q-Q_{\pi})h_{\theta}]$. This discrepancy opposes the intuition derived from the demonstration in Cai et al. (2019), where, in naive vector space TD learning optimization with linear function approximation of the Q-function, the semi-gradient $\mathbb{E}{\varsigma_{\pi}}[(Q_{\Theta}-\mathcal{T}^{\pi}Q_{\Theta})\cdot\nabla_{\Theta} Q_{\Theta}]$ aligns with the oracle gradient direction $\mathbb{E}{\varsigma_{\pi}}[(Q_{\Theta}-Q_{\pi})\cdot\nabla_{\Theta} Q_{\Theta}]$. Hence, to estimate the true function $Q_{\pi}$, we propose a novel TD learning algorithm using the direction of the TD error $\mathbb{E}{\varsigma_{\pi}}[(Q-\mathcal{T}^{\pi}Q)h_{\theta}]$. This proximal dual-loop algorithm goes beyond the analysis of existing NTK scaling, achieving probabilistic distribution optimization and demonstrating convergence in MF networks, contrary to the aforementioned intuitive discrepancy. Our study provides the first convergence analysis of TD learning in MF networks without great regularization (e.x. using sufficiently large regularization parameter $\lambda>$const. in Leahy et al. 2022). Additionally, the practical validation of this newly proposed algorithm in Section 5 is a noteworthy contribution as well.