Venue: ICLR 2021 Conference

Reviews:

final_score: rating: 7: Good paper, accept
review_content: This paper studies the asymptotic convergence properties of (population-level) policy gradient methods with two-layer neural networks, softmax parametrization, and entropic regularization, in the mean-field regime. By modelling the hidden layer as a probability distribution over the parameter space, the training dynamics of policy gradient methods can be written as a partial differential equation. Under certain regularity conditions, the paper shows that if the training dynamics converge to a stationary point, this limiting point is a globally optimal policy. The paper also presents results for finite-time convergence of the training dynamics for neural networks to the mean-field limit.

The optimization landscape and convergence properties of policy gradient methods have drawn attention in RL theory for a long time, and it is nice to see a work that studies this problem from the perspectives of mean-field limit of neural networks, albeit being completely asymptotic. Overall I think this makes an interesting contribution, and I appreciate the sketch of proof ideas in the simpler bandit case. Technically, it seems that the main results are built upon existing frameworks of Mei et al., (2018), Chizat and Bach et al., (2018), etc. But the author also pointed out an interesting technical novelty, which is the use of density arguments when the problem structure is in lack of the hidden convexity used in other works.

On the other hand, it appears to me that one major weakness of the result is that the theorem holds true only when the dynamics converges to a stationary point. Can the authors provide conditions under which this can happen? For example, would it be possible to establish some compactness under additional regularity conditions and use it to show the convergence of a subsequence? If the convergence does fail to happen in certain regimes, how will the dynamics behave? Will it convergence to a limiting cycle or diverge? Are there some natural counter-examples? It would be helpful if the authors could provide more discussions on this condition.

Additionally, it seems to me that the paper actually shows that (due to entropic regularization) the limiting point is the Boltzman policy induced by the optimal Q function (at temparature $\tau$), instead of the optimal Q function iteslf. If that is the case, this needs to be stated clearly in the theorem.
review_title: Interesting asymptotic results on policy gradient methods

final_score: rating: 7: Good paper, accept
review_content: This paper provides a mean-field characterization of entropy-regularized policy gradient dynamics for wide single hidden layer neural networks. The evolution of neural network parameters is described by a transport partial differential equation. And the convergence properties of the dynamics are established.

Overall, I vote for accepting. The paper is well-written. The technical contents seem sound and a comprehensive literature review is provided. And authors also conduct numerical experiments to validate the theory. My two minor comments are as follows:

- The remarks in Section 4.1 are mainly explaining why authors need the assumptions to establish the theories. It would be nicer to provide examples when these assumptions hold or discuss the generality of the assumptions.

- It would be nicer to spend more space to discuss the major differences of the theoretical analysis in this paper compared to earlier results on the mean-field limit in the supervised learning setting. It would be helpful to discuss different aspects of the technical analysis more explicitly.
review_title: Review for "Global optimality of softmax policy gradient with single hidden layer neural networks in the mean-field regime"

final_score: rating: 7: Good paper, accept
review_content: Overall, I vote for accepting. This paper extends previous work in the parameter dynamics of simple neural networks to reinforcement learning framework in continuous state and action spaces with nonlinear function approximation and overcomes the challenge of lack of convexity. The main concern of mine would be that the theorems in the paper are not powerful enough to help us fully understand the experiments.

Pros:

1. This paper introduces the mean-field formulation into the reinforcement learning framework. The technical proof seems highly challenging.
2. Under mild conditions, it demonstrates interesting convergence properties of the particle dynamics to the mean-field counterpart and further, the mean-field dynamics to the global optima. This provides new insight into theoretical understanding of this problem.

Cons:
1. There still exists gap between the theoretical results and numerical experiments to be filled. The experiments shown in Figure 1 are conducted with finite number of neutrons and constant step-size, but the theorems are stated under the adiabatic limit. Hence there is lack of quantitative results under the setting of finite number of samples and finite gradient step size.
2. The paper only shows convergence, but fails to give more detailed properties like convergence rate, etc. 

review_title: This paper provides a mean-field formulation of policy gradient dynamics in parameter space in the reinforcement learning framework. It proves the convergence of the particle dynamics to their mean-field counterpart and the convergence of the mean-field dynamics.

final_score: rating: 7: Good paper, accept
review_content: This paper studies the mean-field limit of the policy gradient method (with entropy regularized) and proves that any stationary point under this setting is a global minimizer. I am not able to verify the entire proof as it involves a lot of standard steps to bridging the finite parameter case and the mean-field limit. The result seems promising and well complements several theory results in RL in the past year, e.g. the optimality of policy gradient under NTK regime and the TD algorithm in the mean-field regime. 

Although the paper does not provide the convergence guarantee of the mean-field density flow to a stationary point (please correct me if this is wrong), the characterization of the optimality is still a good contribution. It well explains why a neural network policy is globally optimal given (1)it is stationary under the training via the first-order method (2)its parameterization has strong expressive power, e.g. it has infinite parameters or it is essentially a nonparametric model. 

The convergence (section 4.2) to the many-particle limit. i.e. mean-field limit, seems standard, as the authors claim it is very similar to the case of supervised learning. I still would like to ask whether the authors found any key differences between the supervised learning case and RL objective, i.e., maximizing the total reward. In particular, does the absence of a strongly convex loss function cause any difficulty in the proof?
review_title: Solid contribution to RL theory

Meta-review: This paper takes a step towards understanding the role of nonlinear function approximation--- more specifically, function approximation via (two-layer) neural nets---in some variants of the policy-gradient algorithms. The authors borrow the mean field analysis idea recently popularized in studying shallow neural nets, and investigate the mean-field limits of the training dynamics in the current RL settings. The results and analyses are interesting as they nicely complement another line of linearization-based analyses (i.e., the one based on neural tangent kernels) towards understanding non-linear function approximation. As suggested by a reviewer, it would be nice to add discussions in the revised paper regarding when the dynamics can be guaranteed to converge to a stationary point. 

Decision: Accept (Poster)