Global optimality of softmax policy gradient with single hidden layer neural networks in the mean-field regime
Authors: Andrea Agazzi
Jianfeng Lu
Venue: ICLR 2021 Poster (2021)
Keywords: policy gradient
entropy regularization
mean-field dynamics
neural networks

Abstract: We study the problem of policy optimization for infinite-horizon discounted Markov Decision Processes with softmax policy and nonlinear function approximation trained with policy gradient algorithms. We concentrate on the training dynamics in the mean-field regime, modeling e.g. the behavior of wide single hidden layer neural networks, when exploration is encouraged through entropy regularization. The dynamics of these models is established as a Wasserstein gradient flow of distributions in parameter space.  We further prove global optimality of the fixed points of this dynamics  under mild conditions on their initialization.

Published a conference paper at ICLR 2021
GLOBAL OPTIMALITY OF SOFTMAX POLICY GRADIENT WITH SINGLE HIDDEN LAYER NEURAL NETWORKS IN THE MEAN-FIELD REGIME
Andrea Agazzi,
Jianfeng Lu
Department of Mathematics, Duke University,
Department of Mathematics,
Department of Physics and Department of Chemistry, Duke University
agazzi @math. duke. edu
I an fenggggmmtth duke.
ABSTRACT
We study the problem of policy optimization for infinite-horizon discounted Markov Decision Processes with softmax policy and nonlinear function approximation trained with policy gradient algorithms. We concentrate on the training dynamics in the mean-field regime, modeling e.g., the behavior of wide single hidden layer neural networks, when exploration is encouraged through entropy regularization. The dynamics of these models is established Wasserstein gradient flow of distributions in parameter space. We further prove global optimality of the fixed points of this dynamics under mild conditions on their initialization.
INTRODUCTION
In recent years, deep reinforcement learning has revolutionized the world of Artificial Intelligence by outperform- ing humans in a multitude of highly complex tasks and achieving breakthroughs that were deemed unthinkable at least for the next decade. Spectacular examples of such revolutionary potential have appeared over the last few years, with reinforcement learning algorithms mastering games and tasks of increasing complexity, from learning to walk to the games of Go and Starcraft (Mnih et al., 2013; 2015; Silver et al., 2016; 2017; 2018; Haarnoja et al., 201 8a; Vinyals et al., 2019).
In most cases, the main workhorse allowing artificial intelligence pass such unprecedented milestones was a variation of a fundamental method to train reinforcement learning models: policy gradient (PG) algorithms (Sutton et al., 2000). This algorithm has disarmingly simple approach to the optimization problem hand: given parametrization of the policy, it updates the parameters in the direction of steepest ascent of the associated integrated value function. Impressive progress has been made recently the understanding of the convergence and optimization properties of this class of algorithms in the tabular setting (Agarwal et al., 2019; Cen et al., 2020; Bhandari & Russo, 2019), in particular leveraging the natural tradeoff between exploration and exploitation offered for entropy- regulareee rewards by softmax policies (Haarnoja al., 2018b; Mei et al., 2020). However, this simple algorithm alone is not sufficient to explain the multitude of recent breakthroughs in this field: in application domains such Stareraft, roboties or movement planning, the space of possible states and actions are exceedingly large- or even continuous and can therefore not be represented efficiently by tabular policies (Haarnoja et al., 201 8a). Consequently, the recent impressive successes of artificial intelligence would be impossible without the natural choice of neural networks to approximate value functions and or policy functions in reinforcement learning algorithms (Mnih et al., 2015; Sutton et al., 2000).
While neural networks, in particular deep neural networks, provide a powerful and versatile tool to approximate high dimensional functions on continuous spaces (Cybenko, 1989; Hornik, 1991; Barron, 1993), their intrinsic nonlinearity poses significant obstacles to the theoretical understanding of their training and optimization properties. For instance, it is known that the optimization landscape of these models is highly nonconvex, preventing the use of most theoretical tools from classical optimization theory. For this reason, the unprecedented success of neural networks in artificial intelligence stands in contrast with the poor understanding of these methods from a theoretical perspective. Indeed, even in the supervised setting, which can be viewed as a special case of reinforcement learning, deep neural networks are still far from being understood despite having been an important and fashionable research focus in recent years. Only recently, a theory of neural network learning has started to emerge, including recent works on mean-field point of view of training dynamics (Mei al.,

---

Published a conference paper at ICLR 2021
2018; Rotskoff & Vanden-Eijnden, 2018; Rotskoff et al., 2019; Wei et al., 2018; Chizat & Bach, 2018) and on linearized dynamics in the over- wver-parametrized regime (Jacot et al., 2018; Allen-Zhu et al., 2018: Du et al., 2018; 2019; Zou et al., 2018: Allen-Zhr et al., 2019; Chizat et al., 2019; Oymak & Soltanolkotabi, 2020; Ghorbani et al., 2019; Lee et al., 2019). More specifically to the context of reinforcement learning, some works focusing on value-based learning (Agazzi & Lu, 2019; Cai et al., 2019; Zhang et al., 2020), and others exploring the dynamics of policy gradient algorithms (Zhang et al., 2019) have recently appeared. Despite this progress, the theoretical understanding of deep reinforcement learning still poses a significant challenge to the theoretical machine learning community, and it is of crucial importance to understand the convergence and optimization properties of such algorithms to bridge the gap between theory and practice.
CONTRIBUTIONS.
The main goal of this work is to investigate entropy regulaizee policy gradient dynamics for wide, single hidden layer neural networks. In particular, we give the following contributions:
We give a mean-field formulation of policy gradient dynamics in parameter space, describing the evolution of neural network parameters in the form of a transport partial differential equation (PDE). We prove convergence of the particle dynamics to their mean-field counterpart, We further explore the structure of this problem by showing that such PDE is a gradient flow in the Wasserstein space for the appropriate energy functional.
We investigate the convergence properties of the above dynamics in the space of measures. In particular, we prove that under some mild assumptions on the initialization of the neural network parameters and on the approximating power of the nonlinearity, all fixed points of the dynamics are global optima, i.e. the approximate policy learned by the neural network is optimal.
RELATED WORKS.
Recent progress in the understanding of the parametric dynamics of simple neural networks trained with gradient descent in the supervised setting has been made in (Mei et al., 2018; Rotskoff & Vanden-Fijndenn 2018; Wei et al., 2018; Chizat, 2019; Chizat & Bach, 2020). These results have further been extended to the multilayer setting in (Nguyen & Pham, 2020). In particular, the paper (Chizat & Bach, 2018) proves optimality of fixed points for wide single layer neural networks leveraging a Wasserstein gradient flow structure and the strong convexity of the loss functional WRT the predictor. We extend these results to the reinforcement learning framework, where the convexity that is heavily leveraged in (Chizat & Bach, 2018) () is lost. We bypass this issue by requiring sufficient expressivity of the used nonlinear representation, allowing to characterize global minimizer as optimal approximators.
The convergence and optimality of policy gradient algorithms (including in the retropy-r setting) is investigated in the recent papers (Bhandari & Russo, 2019; Mei et al., 2020; Cen et al., 2020; Agarwal et al., 2019). These references establish convergence estimates through gradient domination bounds. In (Mei et al., 2020; Cen et al., 2020) such results are limited to the tabular case, while (Agarwal et al., 2019; 2020) also discuss neural softmax policy classes, but under different algorithmic update and assuming certain well-conditioning assumptions along training. Furthermore, all these results heavily leverage the finiteness of action space. In contrast, this paper focuses on the continuous space and action setting with nonlinear function approximation.
Further recent works discussing convergence properties of reinforcement learning algorithms with function approximation via neural networks include (Zhang et al., 2019: Cai et al., 2019). These results only hold for finite action spaces, and are obtained in the regime where the network behaves essentially like a linear model (known as the neural or lazy training regime). in contrast to the results of this paper, which considers training in nonlinear regime. We also note the work (Wang et al., 2019) where the action space is continuous but the training is again in an approximately linear regime.
2 MARKOV DECISION PROCESSES AND POLICY GRADIENTS
We denote a Markov Decision Process (MDP) by the 5-tuple (S P.r,,),, where S is the state space, A is the action space, P P(s, a, 8') d').,''S, a Markov transition kernel, r(s, a, Es.. is the real- valued, bounded and continuous immediate reward function and 7 E 1) is a discount factor. We will consider probabilistic policy, mapping a state a probability distribution on the action space, so that S + (A), where (.(A) denotes the space of probability measures on A. and denote for any E the corresponding density ((S))) :) A - R. The policy deñnes a state-to-state transition operator
2

---

Published conference paper ICLR 2021
P=(s,d6) S. P(s, a, (s,))(s da), and we assume that Px is Lipschitz continuous as an operator (+ + (S) wrt the policy. We further encourage exploration by defining some (relative) entropx-regulartzed rewards (Williams & Peng. 1991)
(=, r(s, a') DKL (T(s,));
where DKL. denotes the relative entropy, स a reference measure and T indicates the strength of regularization. Throughout, we choose 7 to be the Lebesgue measure on A, which we assume, like S, to be a compact subset of the Euclidean space. This regularization encourages exploration and absolute continuity of the policy wKT Lebesgue measure. Consequently, with some abuse of notation, distribution and its density in phase space. Note that the original, unregularized MDP can be recovered in the limit 0.
throughout the same notation for
In this context, given a policy æ the associated value function V S + R maps each state to the aafinit-hhorizon expected discounted reward obtained by following the policy and the Markov process defined by P:
V+(+) E
(st, at, (241)))) =
Exr Σ (r(se, (e, 8641) - (T +)+(())+())))))
where (- = s] denotes the expectation of the stochastic process 9, starting at sg s and following the (stochastic) dynamies defined recursively by the transition operator P (= ds') = P(s, a, do..ds))n (s)) da). Correspondingly, we define the Q-function Qx S as
Qx(s,))= Ex,
r(s0, ao. 81) + Σ
7 RR- (st, ae, 8141) 5,0g
F(s, a) + B (s() 8o 8, 00 a
where F(s, a) = E[r(s, , a, s')] ()]]] is the average reward from (s, a). Conversely, from the definition, we have the 7 7DkI. ( (): +):
identity for (s) E., [Q~ (50, 0o)]s0 We are interested in learning the optimal policy x² of given MDP (5, A, P,r,7), which satisfies for all E
V
max W S-M (A)
V=(s).
(3)
More specifically we would like to estimate this function through a family of approximators To S S -M (A) parametrized by vector W == RR. Note that since consider entropy-regulanee rewards, the optimal policy will be probabilistic policy (given as a Boltzmann distribution) instead of a deterministic one.
A popular algorithm to solve this problem is given by policy gradient algorithm (Sutton & Barto, 2018). Starting from an initial condition (0) E W. this algorithm updates the parameters of the predictor in the direction of steepest ascent of the average reward
w + 1) = w(t) + 3, 3. V.llerg Vr((),
for fixed absolutely continuous initial distribution of initial states do E ML((S and sequence of time steps {.:h.. Here E[-] denotes an approximation of the expected value operator.
This work investigates the regime of asymptotically small constant step-sizes 3, + 0. In this adiabatic limit, the stochastic component of the dynamics is averaged before the parameters of the model can change significantly. This allows to consider the parametric update as a deterministic dynamical system emerging from the averaging of the underlying stochastic algorithm corresponding to the limit of infinite sample sizes. This is known as the ODE method (Borkar, 2009) for analyzing stochastic approximation. We focus on the analysis of this deterministic system to highlight the core dynamical properties of policy gradients with nonlinear function approximation. The averaged, deterministic dynamics is given by the set of ODES
dt v() = Egngo Vx))) = Egng. awu V log (s, (Q,
(8,a) a) - Ilg((T (8, a)))
(5)
3

---

Published conference paper ICLR 2021
where in the second equality we have applied the policy gradient theorem (Sutton et al., 2000; Sutton Barto, 2018), defining for fixed @u E .M)(S)
00
e (50; 8) = Σ (000
e ex( - e" ()
(6)
as the (improper) discounted empirical measure. For completeness, we include a derivation of (5) in Appendix A.
SOFTMAX POLICIES IN THE MEAN-FIELD REGIME
We choose to represent our policy as a softmax policy:
exp(fu(ss.))
(s a)
SA (s,
and parametrize the energy f as a two-layer neural network in the mean-field regime, i.e.,
1 fu(s,a) = N vss a; w:)
for fixed, usually nonlinear function db components w; E Z so that W = S2 We can rewrite the above expression in terms of an empirical measure:
x A x Ω R, where we have separated w E W into N identical
N
(8) = / (ss (N))
where (dw)
Σ
pot (dw) E M4
N
This empirical measure representation removes the symmetry of the approximating functions under permutations of parameters we It also facilitates the limit N + co, when y(M) weakly, that fuIN + fu. Then, for general distribution v E .M (22 the softmax mean-field policy reads:
Tu(s, a) = exp (Jo 1h(8, a: w)u(dw))
(8)
IA exp exp (Ja (sss a; ())()))) da
Note that by our choice of softmax policy and mean-field parametrization (7) we have
Vw, log 7 TT(N) (8, a) = (iIFoNN)ss Vw; log
exp fuLNs (s, a)da
N
SA V Va, exp
N E (s,a)w)) da
V w.
(s, a: w;):))
N
SA exp FuCNY (s, a)da
/ Vi, 1b(8, a; (s,
V an ((,
N
Thus the training dynamics (5), after an appropriate rescaling of time + t/N, which due to the mean-field parametrization for fw), can be rewritten as
d dt (((()) = [ (8, Ex ,NN) Vw, x
SXA
(Qn (s, log mpiN (s, a)) TT (s,)) (ds).
(9)
The training dynamics can be more compactly represented by the evolution of the measure E M1 in parameter space, given by mean-field transport partial differential equation of the Vlasov type
d _ (w) = div (u((w)
C.. (V tb :ww, Q. log [s](s)e.. (ds))))))))
(10)
SXA
where w E Ω and we have introduced the shorthand Cx to denote the covariance operator WKT the probability measure (ss, da). Note that the above partial differential equation also captures the dynamics of the width system, Le., of the empirical measure (N) where each u, follows (9).
4

---

Published conference paper ICLR 2021
We further note that the dynamics introduced above have a gradient flow structure in the probability space (St): defining the expected value function
Elp) = Exxre [Vr, (80)]
(11)
the dynamics (10) is gradient flow for & in the Wasserstein space (see e.g., (Santambrogio. 2017) for an introduction), as prove in the appendix:
Proposition 2.1. For fixed initial distribution @u M (S), the dynamics (10) is the Wasserstein gradient flow of the energy functional (11).
Analogous dynamics equation for evolution of parameter space measure in the supervised learning case has been derived in (Mei et al., 2018; Rotskoff & Vanden-Eijnden, 2018; Chizat & Bach, 2018) and the TD learning case in (Agazzi & Lu, 2019; Zhang et al., 2020). In particular, in the case of supervised learning, the resulting dynamics is a Wasserstein gradient flow, the structure of which is used to obtain the convergence of the particle system to the mean-field dynamics. In our case, however, the energy functional is not convex WRT the policy and moreover the softmax parametrization destroys the convexity of the approximator of the policy with respect to Vt. Thus showing convergence of the dynamics becomes much more challenging.
SIMPLIFIED SETTING: THE BANDIT PROBLEM
We now introduce our results in the simple bandit setting, where state space S is one point (and will be henceforth suppressed in the notation) and without loss of generality action space A is continuous. In this case, for reward function r' and softmax policy
exp(fr(a)) (a) JA exp(ff(a))dd
we have that the value function for the regularized problem reads (we denote V,, Va to simplify notation)
Vy = / (r(a) - T log n. (
while the Q function is simply Q(a) r(a). We further note that the optimal policy in the regularized case reads:
z / (xp((=+r(a)))
x² xa = Z ir(a)),
Recalling the definition of the covariance operator Gxlffg9(s from (10), the expression for the policy gradient vector field in the latter case simplifies to
Oar = Fi v.) = DT. C.,, (wwb(aa log(Te)] - (T())))
/ (vwb(asw) / V Wub(dd:))
(12)
where DTy, DV, denote the Fréchét derivative of h., and V, WRT v and v respectively. Note that by the structure of the covariance operator C,, adding constant to the function f-( does not affect the dynamics. This reflects the fact that the softmax policy is normalized by definition.
3.1 GLOBAL OPTIMALITY OF SOFTMAX POLICY GRADIENT
We now sketch the main steps in proving that the mean-field policy gradient dynamics converge, under appropriate assumptions, to global optimizers. The proof in this simpler setting is much more transparent than the general case to be discussed in the next section, and will provide some intuition for the latter. The first part of the proof concerns the properties of fixed points of the dynamics (10), while the second part concerns the training dynamics.
STATICS
We first informally prove global optimality of any fixed point 1/" of the transport equation 4 1/ Fi) with F; from (12) such that
a) v² has full support in Ω
5

---

Published conference paper ICLR 2021
b) the nonlinearity is 1-homogeneous the first component its parameters, i.e. that writing w = (wo, w) R x & one has ob(a; w) uoσ(a;ā) @) for regular enough ø : A x 6 R
c) the span of { @) (see is dense in (A),
so that Le., Ty. (a) = (a) = 7- exp [[-]] Weaker assumptions and the general statement are given in
the next section, while the general proof appears in the appendix. directly implies that for almost all w E Q
First, we note that by assumption a), div(v=F(
F(w;V") /A Vuvb(aww) (r(a) - (a) V.-) n (da) 0.
In particular, by homogeneity assumption b), the first component of the above vector field must vanish on θ.
f o(a; (r(a))))To VV-)) n.- (da) 0.
By assumption c) that span of @ dense in (A) the above implies that
r(a) T (a) V -.e in A.
(13)
Finally, recalling that by the softmax parametrization and by the boundedness of é. n (a) > 0, we must have
r r(a) +C
which directly implies the optimality of the policy.
DYNAMICS
While it is clear that assumption b) and c) about the structure and approximating power of the nonlinearity vi hold independently of f, we want to show that assumption a) also holds uniformly in time. In this sense, the continuity of the vector field (12) will preserve the full support properties of the measure is for all t >> 0, will prove in more general framework in Lemma C.2. Consequently, any measure V respecting assumption a) at initialization will do for any finite positive time O. However, the question remains of whether this property still holds at t = oo. This is the object of Lemma C.3, where we prove that whenever the gradient approaches fixed point in parameter space, this fixed point is not a global minimizer, it must be avoided by the dynamics, and thus, the only possible fixed points of the dynamics are global minimizers.
RESULTS IN THE GENERAL SETTING
We now come back the general MDP framework introduced in Section 2.
4.1 ASSUMPTIONS
To state the main result of this section, the optimality of fixed points of (10) we need the following Assumption 1. Assume that w (wo, E R for Q 1m and vis,, a;w: wodlss 0; w) with
a) Regularity of di 6 is bounded, differentiable and Dew is Lipschitz. Also, for all J E I x A) the regular values of the map W ++ (w) == [(,, ())((,,;;))) a; are dense its range, and gr(ww) converges in C (f@ E ① ||ا|| 1}} as
to map 0y (y(w)) whose regular values dense in its range.
b) Universal approximation: the span of (ő(-,a) : w E Q} is dense in L2 A);
c) Support of the measure: There exists
0 s.f. the support of the initial condition vo is contained in O : [-r,r] x θ and separates {-r] x θ from {r} x θ i.e. any continuous path connecting {-r) x θ to (r) x θ intersects the support of Vo-
Assumption a) is a common, technical regularity assumption ensuring that (10) is well behaved and controlling the growth, variation and regularity of Φb Alternative assumptions on the case θ f Rm-1 are given in the appendix. Assumption 1 b) speaks to the approximating power of the nonlinearity, assumed to be expressive enough to approximate any function (2 x A). This condition replaces the convexity assumption from Chizat & Bach (2018), as the lack of convex structure in our setting prevents us from identifying the local and global minimization properties of fixed point. Indeed, despite the one-point convexity of E. [V>(s) as

---

Published a conference paper at ICLR 2021
functional of x (Kakade & Langford, 2002) which can be leveraged in the tabular case, such property will be lost, in general, when restricting to policies through nonlinear function approximation. We bypass this issue by requiring sufficient expressivity of the approximating function class, guaranteeing that the optimal policy can be represented with arbitrary precision. Similar assumption on approximability of neural network representation was made in recent analysis of natural policy gradient algorithm (Agarwal et al., 2019). We note that this assumption is easily satisfied by widely used nonlinearities by the universal approximation theorem (Cybenko, 1989; Barron, 1993). Examples of activation functions satisfying Assumption I a)-b) include sigmoid, tanh and Gaussian radial function nonlinearities. Extension analogous results in the ReLU case was discussed in Wojtowytsch (2020) for the supervised learning. Finally, Assumption 1 c) guarantees that the initial condition such that the expressivity from b) can actually be exploited. This condition is satisfied for example by the product of a uniform distribution on any bounded set A C IR with the normal distribution on θ or, if θ is compact, with the uniform distribution on e.
4.2 CONVERGENCE OF THE MANY-PARTICLA LIMIT
Before discussing the optimality properties of the dynamics (10), we show that this PDE accurately describes the policy gradient dynamics of a sufficiently wide, single layer neural network. To this aim, we let P2(52) be the space of probability distributions on Ω with finite second moment.
(N) Theorem 4.1. Let Assumption 1 hold and let (N) be a solution of (5) with initial condition w E W = 22M v converges to Vo E P2((2) in Wasserstein distance W2 then (N)) converges, for every t 0, to the unique solution in (10).
(N) We note that by the law of large numbers for empirical distributions, the condition of convergence of 1/) to Vo is e.g., satisfied when wo are drawn independently at random from Vo-
The proof of this result is largely standard, under the given assumptions, and is provided in the appendix for completeness. The idea of the proof is a canonical propagation of chaos argument (Sznitman, 1991). In a nutshell the first step of the argument establishes sufficient regularity of the gradient dynamics, allowing to guarantee existence and uniqueness of the solution to (10). Then, one bounds the difference in differential updates for the particle system and the mean-field dynamics by comparing them with the evolution of the particle system according to a linear, time-inhomogentoos PDE using the drift term of the mean-field model. The proof is finally concluded by application of Gronwall inequality. The main difficulty WRT similar results in the literature is to establish the needed Lipschitz continuity of the vector field driving the transport PDE: while this is an immediate consequence of assumptions on the activation functions and on the risk functional in the supervised setting, proving this type of regularity requires more effort in the RI. setting given the involved dependence of the vector field on the measure V4-
4.3 OPTIMALITY
After discussing the connection between particle dynamics and mean-field equations, we present the main convergence result of this paper:
Theorem 4.2. Let Assumption 1 hold and ve given by (10) converge to v², then Ty- = n², the optimal policy for (3).
Thus if the policy gradient dynamics (10) converges to a stationary point, that point must be a global minimizer. Again, we emphasize that in our regularized setting x² is given by probability distribution, and can thus be represented as a softmax policy. We prove this result in three steps. First, we connect the optimality of a stationary point with the support of the underlying measure in parameter space. More specifically, we show in Lemma C.I that by the expressivity of ø. the transport vector field of suboptimal fixed points of the dynamics (10) cannot vanish everywhere in parameter space. This implies that measure with sufficient support cannot correspond to suboptimal fixed point.
We then show in Lemma C.2 that such sufficient notion of support (Assumption c) is preserved by the mean-field policy gradient dynamics (10) throughout training. For any finite time, this is true by topological arguments: the separation property of the measure cannot be altered by the action of a continuous vector field such as (10). We note in particular that we do not prove that assumption a) in Section 3 holds in this case. Finally, in Lemma C.3 we combine the above results and prove that spurious fixed points are avoided by the policy gradient dynamics (10) when initialized properly. To establish this we argue by contradiction: assuming
7

---

Published conference paper ICLR 2021
10-4
550
20rr
30
10
10
0.0
02
5.4
0.6
0.8
1.0
toratters
led
ReratinG
le5
(a)
(b)
Figure 1: Evolution of E(v") (v:) as a function training time, for experiments (a) and (b) as described in the main text. Different lines correspond different random initializations of reward function and learning model. We see that the training еrror decreases monotonically (but with different rates) during trainig.
that we are approaching such a spurious fixed point at time to. we show in Lemma C.5 that the velocity field will change little for any t to- In particular, it follows that in this regime the dynamics of (10) can be approximated by the gradient descent dynamics (in particle space) of an approximately fixed potential. On the other hand, by Assumption c) and by the homogeneity of th, amount of measure i will fall in a forward invariant region where its wh component will grow linearly in t (which exists by Lemma C.I), thereby eventually contradicting the assumption that 7 is fixed point of (10).
able to show that by Lemma C.2 a positive
There are two main conceptual differences between the proof outlined above and the one carried out in the supervised learning setting. On one hand, necessary step in our proof is to establish Lipschitz continuity of the vector field defining the transport equation (10), also needed for convergence the particle dynamics as discussed above. On the other hand, the landscape of the objective function for policy gradients does not enjoy the convexity (WRT the predictor) typically assumed the supervised case. To exclude the existence of local minima assume sufficient expressivity of the activation functions, absent in the supervised analysis. This assumption is key to deduce optimality of fixed points of (10) in Theorem 4.2 in our less regular setting.
NUMERICAL EXAMPLES
To test our theoretical results in a simple setting we train a wide, single hidden layer neural network with policy gradients to learn the optimal softmax policy (8) for entropy-regularizer rewards with parametrization (7) and regularization parameter 0.2. We do so in two separate settings:
(a) S = {0). [o, 1]. This setting corresponds to bandits framework discussed in Section 3.
(b) S x A is a grid of size 100 x 100 in the set [0, 1J2. In this case, we have chosen a discount factor of
7 = 0.7 and transition process given by P(s, a, s') 0.90(s' a) + 0.1/100 (i.e., an action leads to the corresponding state s' = with probability 0.9 and is uniformly distributed with probability 0.1). At each iteration we have computed the exact distribution @n by computing the resolvent of the (weighted) transition matrix.
In both cases, we defined the optimal Q function as (*, = f (s, a), where (8 a) given by a single hidden layer neural network of width n. 5, ReLU nonlinearites and weights drawn independently and identically distributed from a centered, normal distribution with variance
4, i.e.
(" a) (s. for wi AO(00 4).
We learn the optimal policy for the problem defined above using a N 800-neurons wide single hidden layer neural network with ReLU nonlincarities in the mean-field regime (7) used as energy for a softmax policy (8). The initialization of the student network is follows: first-layer weights are initialized random drawn

---

Published a conference paper ICLR 2021
independently from a centered, normal distribution with variance o² 4, while output weights are initialized at O. The model is trained according to (4) with fixed step-size 3, 10 We report the results of this training procedure in Fig. 1, where we notice that all the paths monotonically decrease the error E(v") E(v),, as predicted by our results. Note that the convergence rate of the model varies across experiments, consistently with the purely qualitative nature of the convergence result we proved.
CONCLUSIONS AND FUTURE WORK
This work addresses the problem of optimality of policy gradient algorithms, a workhorse of deep reinforcement learning, when combined with mean-field models such as neural networks. More specifically, we provide a mean-field formulation of the parametric dynamics of policy gradient algorithms for regrlay-regulariied MDPS and prove that, under mild assumptions, all fixed points of such dynamics are optimal. This extends similar results obtained in the "neural" or "lazy" regime to the mean-field one, which is known to be much more expressive (E et al., 2019; Ghorbani et al., 2020), but also highly nonlinear. The latter feature prevents, at present, from obtaining convergence results of these models, except in very specific settings (Chizat, 2019; Javanmard et al., 2019).
Interesting avenues or future research include the relaxing the adiabaticity assumption, i.e., considering the stochastic approximation problem resulting from the finite number of samples and the finite gradient step-size, as well as establishing quantitative bounds for models with large. but finite, number of parameters. Probably the most important open question, however, concerns establishing quantitative convergence of mean-field dynamics of neural networks: even in the supervised setting, despite recent results in specific settings (Chizat, 2019; Javanmard et al., 2019), these guarantees remain mainly out of reach.
ACKNOWLEDGMENTS
AA acknowledges the support of the Swiss National Science Foundation through the grant P2GEP2-1750) and by the NSF grant DMS- 1613337 The work of JL part supported by the US National Science Foundation via grants CCF-1910571 (Duke TRIPODS) and DMS-2012286.
REFERENCES
Alekh Agarwal, Sham M Kakade, Jason D Lee, and Gaurav Mahajan. Optimality and approximation with policy gradient methods in markov decision processes. arXiv preprini arXiv: 1908. 00261, 2019.
Alekh Agarwal, Mikael Henaff, Sham Kakade, and Wen Sun. Pc-pg: Policy cover directed exploration for provable policy gradient learning. arXiv preprint arXiv:2007.08452 2020.
Andrea Agazzi and Jianfeng Lu. Temporal-diiference learning for nonlinear value function approximation in the lazy training regime. arXiv preprint arXiv:1905. 10917, 2019.
Zeyuan Allen-Zhu, Yuanzhi Li, and Zhao Song. A convergence theory for deep learning via over- parameterization. arXiv preprint arXiv: 0,0962, 2018.
Zeyuan Allen-Zhu, Yuanzhi Li, and Yingyu Liang. Learning and generalization in overparameterized neural networks, going beyond two layers. In Advances in neural information processing systems, pp. 6155-6166, 2019.
Luigi Ambrosio, Nicola Gigli, and Giuseppe Savaré. Gradient flows: in metric spaces and in the space of probability measures. Springer Science & Business Media, 2008
A. R. Barron. Universal approximation bounds for superpositions of a sigmoidal function. IEEE Transactions Information Theory, 39(3):930-944) 1993.
Jalaj Bhandari and Daniel Russo. Global optimality guarantees for policy gradient methods. arXiv preprint arXiv:1906.01 7986 2019.
Vivek S Borkar. Stochastic approximation: a dynamical systems viewpoint, volume 48. Springer, 2009.
Qi Cai, Zhuoran Yang, Jason D Lec, and Zhaoran Wang. Neural
learning converges to
global optima. In Advances Neural Information Processing Systems, pp. 11315-11322 1136,, 2019.

---

Published conference paper ICLR 2021
Shicong Cen, Chen Cheng, Yuxin Chen, Yuting Wei, and Yuejie Chi. Fast global convergence of natural policy gradient methods with entropy regularization. arXiv preprint arXh::000.005588 2020.
Lenaic Chizat. Sparse optimization on measures with over-parameterrized gradient descent. arXiv preprint arXiv:1907.10300 2019.
Lénaïc Chizat and Francis Bach. On the global convergence of gradient descent for over-parameterized models using optimal transport. In Proceedings of the 32Nd International Conference Neural Information Processing Systems, NIPS' 18, pp. 3040-3050, 2018.
Lenaic Chizat and Francis Bach. Implicit bias of gradient descent for wide two-layer neural networks trained with the logistic loss. arXiv preprint arXiv:2002. 04486, 2020.
Lenaic Chizat, Edouard Oyallon, and Francis Bach. On lazy training in differentiable programming. In Advances in Neural Information Processing Systems, pp. 2933 2933-2943, 2019.
G. Cybenko. Approximation by superpositions of sigmoidal function. Mathematics of Control, Signals and Systems, 2(4):303-314 Dec 1989.
S. S. Du, X. Zhai, B. Poczos, and A. Singh. Gradient descent provably optimizes over-parameterized neural networks, 2018. preprint, arXiv: 111v:1810000004
Simon Du, Jason Lee, Haochuan Li, Liwei Wang, and Xiyu Zhai. Gradient descent finds global minima of deep neural networks. In International Conference on Machine Learning, pp. 1675-1685, 2019.
Weinan E, Chao Ma, and Lei Wu. Barron spaces and the compositional function spaces for neural network models. arXiv preprint arXiv:1906.0800 08039, 2019.
Behrooz Ghorbani, Song Mei, Theodor Misiakiewicz, and Andrea Montanari. Linearized two-layers neural networks in high dimension, 2019. arXiv preprint arXiv: 12191.
Behrooz Ghorbani, Song Mei, Theodor Misiakiewicz, and Andrea Montanari. When do neural networks outperform kernel methods? arXiv preprint arXiv:2006. /.309, 2020.
Tuomas Haarnoja, Schoon Ha, Aurick Zhou, Jie Tan, George Tucker, and Sergey Levine. Learning to walk via deep reinforcement learning, arXiv preprint arXiv:1812.11103, 2018a.
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum
entropy deep reinforcement learning with a stochastic actor. In International Conference Machine Learning, pp. 1861-1870 2018b.
Kurt Homik. Approximation capabilities of multilayer feedforward networks. Neural Networks, 4(2):251 257, 1991. ISSN 0893-6080.
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernl:: Convergence and generalization in neural networks. Advances in neural information processing systems, pp. 8571-8580, 2018.
Adel Javanmard, Marco Mondelli, and Andrea Montanari. Analysis of a two-layer neural network via displace- ment convexity. arXiv preprint arXiv: 19010009 2019.
Sham Kakade and John Langford. Approximately optimal approximate reinforcement learning, In Proceedings of the Nineteenth International Conference on Machine Learning, pp. 267-274. Morgan Kaufmann Publishers Inc., 2002.
Jachoon Lee, Lechao Xiao. Samuel Schoenholz, Yasaman Bahri, Roman Novak, Jascha Sohl-Dickstein, and Jeffrey Pennington. Wide neural networks of any depth evolve as linear models under gradient descent, In Advances in neural information processing systems, pp. 8570-8581, 2019.
Jincheng Mei, Chenjun Xiao, Csaba Szepesvari, and Dale Schuurmans. On the global convergence rates of softmax policy gradient methods. arXiv preprint arXiv:2005.06392, 2020.
10

---

Published conference paper ICLR 2021
Song Mei, Andrea Montanari, and Phan-Minh Nguyen. A mean field view the landscape two-layer neural networks. Proceedings of the National Academy of Sciences, 115(33).E7665-777)) E7671, 2018. doi: 1L.077/ppas.18005791153
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Alex Graves, Ioannis Antonoglou, Daan Wierstra, and Martin Riedmiller. Playing Atari with deep reinforcement learning. In NIPS Deep Learning Workshop. 2013.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A. Rusu, Joel Veness, Mare G. Bellemare, Alex Graves, Martin Riedmiller, Andreas K. Fidjeland, Georg Ostrovski, Stig Petersen, Charles Beattie, Amir Sadik, Ioannis Antonoglou, Helen King. Dharshan Kumaran, Daan Wierstra, Shane Legg, and Demis Hassabis. Human-level control through deep reinforcement learning. Nature, (7540):529 -3399953333333 02 2015.
Ofir Nachum, Mohammad Norouzi, Kelvin Xu, and Dale Schuurmans. Bridging the gap between value and policy based reinforcement learning. In I. Guyon, U. V. Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017. URL https: /ptpceeeiogceeeeseeeeeeeeeoe neurips cc/paaper/2117/ELle Facf9f74 (ETEE9943338888000 3 94eee 7baal 64 6469 69 Paper pdf.
Phan-Minh Nguyen and Huy Tuan Pham. A rigorous framework for the mean field limit of multilayer neural networks. arXiv preprint arXiv:2001.1144 11.0001.11444. 2020.
Samet Oymak and Mahdi Soltanolkotabi. Towards moderate overparameterization:: global convergence guaran- tees for training shallow neural networks. IEEE Journal on Selected Areas in Information Theory, 2020.
Grant Rotskoff and Eric Vanden-Eijnden. Parameters as interacting particles: long time convergence and asymptotic eтror scaling of neural networks. In S. Bengio, Wallach, Larochelle, K. Grauman, N. Cesa- Bianchi, and R. Garnett (eds.). Advances in Neural Information Processing Systems 31, pp. 7146-7155. Curran Associates, Inc., 2018.
Grant Rotskoff, Samy Jelassi, Joan Bruna, and Eric Vanden-Eijnden Neuron birth-death dynamics accelerates gradient descent and converges asymptotically. In International Conference on Machine Learning. pp. 5508-5517, 2019.
Filippo Santambrogio. (Fuclidean, metric, and Wasserstein) gradient flows: an overview. Bulletin of Mathemat- ical Sciences, 7(1):87-154, 2017.
David Silver, Aja Huang. Christopher J. Maddison, Arthur Guez, Laurent Sifre, George van den Driessche, Julian Schritwiieser, Ioannis Antonoglou, Veda Panneershelvam, Mare Lanetot, Sander Dieleman, Dominik Grewe, John Nham, Nal Kalehbrenner, Ilya Sutskever, Timothy Lillicrap, Madeleine Leach, Koray Kavukcuoglu, Thore Graepel, and Demis Hassabis. Mastering the game Go with deep neural networks and tree search. 2016.
David Silver, Julian Schrittwieser. Karen Simonyan, Ioannis Antonoglou, Aja Huang, Arthur Guez, Thomas Hubert, Lucas Baker, Matthew Lai, Adrian Bolton, Yutian Chen, Timothy Lillicrap, Fan Hui, Laurent Sifre,, George den Driessche, Thore Graepel, and Demis Hassabis. Mastering the game of Go without human knowledge. Nature, 550:354, October 2017.
David Silver, Thomas Hubert, Julian Schrittwieser, Toannis Antonoglou, Matthew Lai, Arthur Guez, Marc Lanctot, Laurent Sifre, Dharshan Kumaran, Thore Graepel, Timothy Lillicrap, Karen Simonyan, and Demis Hassabis. general reinforcement learning algorithm that masters chess, shogi, and Go through self-play. Science, 362(6419): 199--11111-144 144, 2018
R. Sutton and A. G. Barto. Reinforcement Learning: An Introduction. MIT Press, Cambridge, MA, second edition, 2018.
Richard S Sutton, David A McAllester, Satinder P Singh, and Yishay Mansour. Policy gradient methods for reinforcement learning with function approximation. In Advances in neural information processing systems, pp. 1057-1063, 2000.
Alain-Sol Sznitman. Topics in propagation of chaos. In Ecole d' été de probabilités de Saint-Flou XIX-1989, pp. 165-251. Springer, 1991.
11

---

Published conference paper ICLR 2021
Oriol Vinyals, Igor Babuschkin, Wojciech M Czarnecki, Michaël Mathieu, Andrew Dudzik, Junyoung Chung, David H Choi, Richard Powell, Timo Ewalds, Petko Georgiev, et al. Grandmaster level in starcraft using multi-agent reinforcement learning. Nature, 575(7782)3350-344 2019.
Lingxiao Wang, Qi Cai, Zhuoran Yang, and Zhaoran Wang. Neural policy gradient methods: Global optimality and rates of convergence. arXiv preprint arXiv: 01 150, 2019.
C. Wei, J. Lee, Q. Liu, and T. Ma. On the margin theory of feedforward neural networks, 2018. preprint, arXiv:1810.05320
Ronald J Williams and Jing Peng. Function optimization using connectionist reinomrmmen learning algorithms. Connection Science, 3(3):241-266, 1991.
Stephan Wojtowytsch. On the convergence of gradient descent training for two-layer relu-networks in the mean field regime. arXiv preprint arXiv:2005. 13330, 2020.
Kaiqing Zhang, Zhuoran Yang, and Tamer Başar. Multi- Multi-agent reinforcement learning: A selective overview of 2019.
theories and algorithms. arXiv preprint
Yufeng Zhang, Qi Cai, Zhuoran Yang, Yongxin Chen, and Zhaoran Wang. Can temporal-difffrrenee and q-learning learn representation? a mean-field theory. arXiv preprint arXiv:2006.04761, 2020.
D. Zou, Y. Cao, D. Zhou, and Q Gu. Stochastic gradient descent optimizes over-parameterrrized deep ReLU networks, 2018. preprint, arXiv: 111.08888.
12

---

Published conference paper ICLR 2021
A DERIVATION OF SOFTMAX POLICY GRADIENT DYNAMICS
Lemma A.1. The gradient of the regroay-r value function can be written as
Egrugo Va = Ee aanmtt Ve log T (s, a) (Qn( a) - T log (8, (8)
and thus the policy gradient dynamics (5).
Proof. We choose throughout स as the Lebesgue measure, and use that Tw absolutely continuous wrt Taking the gradient of (1) using parametric policy Tw we obtain
Vn V= (8) = B.
7' (r(st, ar, 80+1) - (Tm(
=8
11=0
- loo Te (1) + +V (01) P(s, a1, (w))
= - JSXA
=(a1)
= -Sond r(8,,1,81) - -log (1, + 7 (51) P(s, an dsi (w
F(a1)
+ +SSAA -7V log (8,
V Van P(s a1 dsg a1,
(w da,(A.1)
#(a1)
Now, since Is P(s,a,ds) = and SA (w =1 =1 for all 8, a, w we have for the first term in brackets in the last line
SXA SSKA V. log
Tur (5, P(s, a1, dsi (s, dai) / Ww aj)da Ve / Tw ai)daaaaa #(a1)
On the other hand, we can rewrite the second term in brackets as Eng VV (51) ( so = s). and recognize the sequentially repeat the same
LHS of (A.1) evaluated at the next state 81 in the expectation. Therefore, computation as above, and recalling the definition of q²= and Q, a) in (6) and (2) we obtain
V Vo (s)
/
Σ1
r(st, ae St41) -Flog
(n a(at)
VV (8442) V Vw ww (82, day)
180
oc
Va log (st ar) To (8),
(or -T log
E
F(as)
- JSXA (Se, Flog
Tw (0((,,0) (1)
V log (s, a) xTT (w
π (a))
where in the second line we have used that, if The > on A, Tu Tw (8t, da,) Va log (w ar) = (o dae).
Proposition 2.1. For a fixed initial distribution go. the dynamics (10) is the Wasserstein gradient flow of the energy functional
E(y] = [Vr. (so)]
Proof. We find the potential of the gradient flow by functional differentiation of E:
5y Elul(w)= = / óE dāy, (s, a; w)dd da
dz (s, a)
(A.2)
5y
and consider the two terms in the integral separately, starting from the second:
8T (8, a:
6 e" (ssaaw)))dd)))) du [a e ((sas))u(dw))d da
8v
- et d(sraw))(dw)dd e² I ((s,aww)r(dw)
13

---

Published conference paper ICLR 2021
8 e W /
et I (lssaaarlld)
5c e² (((aaa))lldw) da
=T =T(M(s,050)) (((s,a;w) / (w) mol do') (=(
(A.3)
For the first term in the integrand of (A.2), we use ((s, a) as density and obtain
8n (s,a) = δτ 1 (F(s.a) - log a)) (T)) da)) (s,a)
39
SxA
= - (F(s,a) ((loog((ss. ((lo)))((logsss.o 1)) e (8'
(A.4)
/ (F(s" a²) log T("", d")) 5e "s" 5xA
(s, as)a) de")ds" oo (A.5)
dπ
where in the last line we have used that [= (s a) = log (8, (3,,a++. + We evaluate the variational derivative of e² as
00" s")
oc
/ PH(dd (s,
(8,a) Σ
8T
= I 8P, (8', 8
20
Σ
E (s(,
87 Σ qt
òu P* """" yd""
89
[ 8(5, a, 80 (s" s") ) Pa(s', 8n de (s" 6 s²)
(8, a)ds"
We further recognize the same derivative on the RHS of the above expression, allowing to write
50 (s', 8") (s,
ETpH(((,s)))) P(s, a, s"" """""""" ,
Finally, we notice that the last term in (A.4) constant in and therefore vanishes when integrated against (A.3), so we only consider
g9 (s, a) - ex(s) = ( ISXA - Flg (so,g((",cco "")) x
( + 7 a") 80 (d',"")) (0,10)
(0a.rraane (8',8)
en (o(ds")dsadd da"
x
6x
1 S 1=
Σ ( (F(s" T log T as"))
x (ssarbaa + & / P(s,a, "")= (s" s" ds" eo(ds')ds"da"
/ pt (s, (og go(ds')
t=0
= (Q (s) 7 log (s, a)) (s)
(A.6)
We conclude by noting that combining (A.3) and (A.6) we obtain C² [W((,, Q+ - log (x), and the Wasserstein gradient flow corresponding to this potential is (10).
B PROOFS OF THE MANY-PARTICLE LIMIT
Theorem 4.1. Let Assumption hold and let (N) be a solution of (5) with initial condition wh v (N) converges to Va E P2(()) in Wasserstein distance W₂ then v/ converges, for every t > 0, to the unique solution vi (10).
(N) W Ω*
14

---

Published conference paper ICLR 2021
Proof. As anticipated the main text, the proof divided in two parts:
1. We prove sufficient regularity of the dynamics (10), allowing to establish existence and uniqueness of its solution.
2. We leverage the regularity proven above to establish a propagation of chaos result, showing that the system of interacting particles behaves asymptotically as its mean-field limit.
While carrying out this proof is needed in our context since the dependence of (v) on v/ is more involved than in e.g., Mei et al. (2018); Rotskoff & Vanden- (2018); Chizat & Bach (2018), the steps of this derivation are mainly standard, see e.g., (Sznitman, 1991).
B.1 REGULARITY
We prove existence and uniqueness of the gradient flow dynamics (10) through standard arguments from the optimal transportation literature (see e.g.. (Ambrosio et al., 2008)). More specifically, recalling that π Ty we leverage the Lipschitz continuity of the vector field
F)( = | C. (Voniss sw), Q=(s,)) log (s, ex (ds)
with respect to v. To prove such regularity result, decompose
E(V) R (ʃ thr
for R() = 50
(B.1)
and S (S (M-M (A) - R maps μ + EN [y(s)]s~ 20] and (2 A) + (S (A)) is the softmax policy parametrization (8) of its argument. Recalling the definition of 2, from Assumption 1 and denoting F, = {ʃ vb : s 3 we further define the norms and constants needed in the following proof as:
D Dipwrll
D sup D wEQ,
LDS sup wweER
|| - 22
[RR. - DR
|RRl = sup ((:w) WEQ.
D
LDR = sup A.F.F.
| 1/'12
where || || denotes the operator norm. While for any directly from Assumption I, more work is needed to prove that DR < co, LDn < co. We prove this in Lemma C.5 below, and proceed with the proof of convergence of the particle dynamics. For any r and corresponding 2, from Assumption we define the set of localized functionals
the boundedness of || Ln results
(E(v) if supp (v) Q,
(v)
po clse
Furthermore, we say that coupling 7 MM (S2 2) is an admissible transport plan if both its marginals have support in Q, and finite second moments. To every admissible transport plan, for p 2 1 we associate transportation cost Cp(7) = (Se -W - ('-p//d(ww,, We prove the following results: for every > 0 we have
1. There exists X > such that for any admissible transport plan % defining the interpolation map v/2 := (tIl, + (1 - +)Il) the function +> E(x?) is differentiable with Lipschitz continuous derivative with constant (n))
2. Let Po have support in e,, Then for any given transport plan y with first marginal given by Mo. velocity field F' satisfies
/ F(u) (u u ='
{(II))))) 2 E(vo)
(B.2)
if and only if F(u) is in the subdifferential of DE, (u) == Cr Qu - log (u) for g E L2(S x A) (projected in the interior of Q, when u E ( for uo almost every u E
15

---

Published a conference paper ICLR 2021
The proof the two points above corresponds to (Chizat & Bach, 2018, Lemma B.2). We sketch the proof of these two points below, referring to the original reference for the details.
1. By the Lipschitz continuity of th 52 - (2 A) and (f( DR, C+1,,+ log æ in 2r. the energy gor) (u2) transported along an interpolating path oi is differentiable and can write its derivative as
te (0)) - R
d
( vuu)) / Dév (1wl1 (w/ w)ddr(ws,
Then, again by the Lipschitz continuity of Dijp and DR we have for 0 OS < 1" <1
de (17)) di
(vZ.))
-
(vv)) ( + / |''ef (3().t) / (Dibaaeeatew DV - w)dr((ws
1 [nec]
- Ie( (112)-1i(/ ( / DU Den-r ww)ps/ewwwwwwwdddoww w)dr(w,,
I 2,1e" - -1
for a X large enough, where in the last inequality we have used the uniform bounds on DR in Qr, that DilL-thw4oe - S (t' T & -ا w'I and we applied Hölder's inequality to
bound ((()) S CB()).
2. The proof of this result leverages an expansion of the functionals R, to the second order in their arguments:
vb(w" wlw) + Diu w) + (w, R(g) = DR -RR((()))
Recalling the Lipschitz bounds the remainders (w, < & w/P, Rn(f.9) LDRlI =1ا and combining the two expansions above we have, for transport plan y with marginals vo = (ITo) y,, 112 = (II))
(v1) (v)) / (R((VV))DVs
Duu) u)dyyu, u u')dsda R
for remainder term R. We bound such remainder term, again by the Lipschitz regularity of Dú and DR and by the boundedness of Drù DR in Q,, by C2(7)² and C1(7)² I C2(7)² thereby obtaining that
(v)) E(()) + f R (yo) u) dsda dy v/) o(C2(1))
Noting that the integrand against the coupling is the gradient flow vector field, the above uniquely characterizes the velocity field satisfying (B.2).
We note that point 1) above immediately implies that E(.) is S--memicoe along geodesics, while by 2) g(r) admits strong Wasserstein subdifferentials on its domain (Ambrosio et al., 2008, Definition 10.3.1). Combining the two results one obtains existence and uniqueness of the solutions of the Wasserstein gradient flow through (Ambrosio et al., 2008, Theorem 11.2.1).
B.2 PROPAGATION OF CHAOS
By the Lipschitz continuity of the transport field in (10) in vo with supp vo c Qr, there exists time t, > 0 such that, supp (NN) Q, for all [0, N E N. Consider now two times 0 I t < to S t,. To prove the existence of the limiting curve (Vi).. we show that the curves v (N) are uniformly in N equicontinuous in W2 and as such, possess converging subsequence by Arzela-Ascol Ascoli theorem. To show equicontinuity, bound
the the W2 distance between distributions by coupling positions of the same particles at different times and using Cauchy-Schward inequality:
N
N
W2 (N)
(2))
<
(8)
1
t2 t
d
w
(i) 12dds
N
N
ds
16

---

Published a conference paper ICLR 2021
Combining the above with the identity
N
d
E(vE
(
N Σ
Vv, E(1 )
d
w
d
dl
N i=1
we have
ne d (2 1/2 vt2 112-111/ J
1/2
(N) (N))
Cllv S vt2 -
sup SEu
8(v)
inf sup
{(V)
da
where we recall that J, = fv E (M) supp C 2r). In particular, the above continuity bound is independent on N. proving equicontinuity of (N) W2-
We now prove that the limiting point of the converging subsequence whose existence was identified above must solve (10). To do so we compare both the differential for the mean-field and particle dynamics to the one of linear inhomogneous PDE: for any bounded and continuous f R x R" + R" denoting by E = 1 Fidt and EN write
(N) F(N) dt where Fr, E(N) are the vector fields of the mean-field and particle system respectively, we V
If |SWW)(F- - EN) S &11fll>
/ |T!
Fi du, dt + |ʃ [Fid(() vz)dt
By boundedness of fFi over 2,, the second term converges by our choice of subsequence, For the first term, denoting throughout by || I|aL the bounded Lipschitz norm, we leverage the Lipschitz continuity of the vector field WRT the underlying parametric measure:
(| Felle S C21|Pt- - 2/ |BL
for C₂ > 0 large enough, again obtaining convergence by our choice of subsequence. This proves convergence of the particle model the mean-field equation (10) t,). To extend the time interval which prove convergence, we use that (v)) (and thus R) decays along trajectories of (10). Consequently, by the boundedness of the differential DR on sublevel sets of R the Lipschitz constant of {(v)) is uniformly bounded. Using again the Lipschitz continuity of I' we can show that suPweQ, || < A + Br, i.e., that particle velocities can grow most linearly in r, and an application of Gronwall allows to find, for every T > 0 that there exists r > 0 such that supp E Q for all t E [0, and propagation of chaos follows.
C PROOFS OF OPTIMALITY
Theorem 4.2. Let Assumption I hold and vi given by (10) converge to v", then Tu- TT" S x A-a.e.
in the case where O + p=-1, Our proof of the theorem above can be easily generalized to the setting of Assumption 2.
Before proceeding to prove Theorem 4.2, state the alternative form of Assumption
Assumption 2. Assume that w (wo, @) R x θ for ② C. pm-1 which is the closure of a bounded open convex set. Furthermore (bs, ()) wno(s, at w) where o is bounded, differentiable and Dow is Lipschitz. Also, for all f E I.²(S x A) the regular values of the map + g ( == I f(s, a)o(s, s, a; a adddds are dense in its range and 9y (w) satisfies Neumann boundary conditions (i.e., for all t E 0e we have dg (m)(m) = 0 where naa E pm is the normal of 00 at W).
We prove Theorem 4.2 as sketched in Section 4 by first connecting the optimality and the support of stationary measures in parametric space through Lemma C.I, and then investigating how the dynamics preserves full support property for any f > 0 in Lemma C.2 and avoids spurious minima in Lemma C.3. Before starting this program we introduce the equivalent of greedy policies in the entropy-regularize setting. For given Q(s, a). the associated Boltzmann policy TB with respect to a reference measure म is given by
TB( a) == exp [(Q(s,))) - Vo(s))/T] for Vo(s) == 7 log (exp a)]]]
and satisfies
TB( = arg max max (A) (Ea~x [Q(s,a)] DDL DDLL(T;;
One can then define the Boltzmann backup or soft Bellman backup operator T that, for a given Q and the associated Boltzmann policy TE. gives the action-value function Q associated NB:
T (s, a) F(8, a) + YTTE [log Eaie= (exp (11, 01 /
(C.1)
17

---

Published conference paper ICLR 2021
is known (Haarnoja et al., 2018b, Theorem 1) that the fixed points of the above operator are optimal, Le., they correspond to the optimal policy म π² of the To state the first partial result towards the proof of Theorem 4.2, observe that the wo-component of the transport vector field in (10) can be written as
MDP.
(s Cr. (w (8, - t log mn) eno (ds)
Co, (6, (8, - log nyl @n, (ds).
(C.2)
where we recall that [f g] is the covariance operator WRT the probability measure I(s, da) introduced below (10). We note in particular that the above expression only depends on w. With the above information hand we now proceed to prove Lemma C.1 relating the value of (C.2) and the optimality of fixed points of (10):
Lemma C.1. Let Assumption hold and let v satisfy
SSLA o(ā;s, a) (s, T Flog (s, Vx (s)) (u da) da. (ds) 0,
(C.3)
w-almost everywhere in θ. Then we have that Qu Q. holds 7t a.ee in S x A.
Proof of Lemma C.I. Assuming that (C.3) holds Lebesgue-a.e. in θ by the assumed continuity of in combined with the expressivity of ø Assumption 1 b) we must have that
(s, t log (x(s, a) Vz, (s) 0
TvlT, a.e..
We then rewrite the above condition in compact notation the fixed point equation T=Q,, (8, a) (s.
for the soft Q learning or Boltzmann backup operator 7 defined in (C.1). Since all fixed points T for < 1 are optimal (Nachum et al., 2017, Theorem 3), we must have that тe = To [=] = x for Ty Tyln, -almost every (s, a) S A. The result follows by equivalence of ny and æ.
Consequently, suboptimal fixed points of the dynamics (10) cannot satisfy (C.3) Lebesgue-a.e. in O
C.I PROOF OF THEOREM 4.2
We prove below that spurious local minima that do not satisfy (C.3) O-ae. are avoided by the dynamics. We do so by leveraging the approximate gradient structure of the policy gradient vector field when v, is close of such stationary points, as discussed in the main text. Combining this fact with the assumed convergence to 1² proves Theorem 4.2.
We note that by the assumed homogeneity of ya in its first component, if v, v' are such that
/ wowldwoo dw) = f wDy/ (dwo. dā)) a.c.
(C.4)
then
/ wool :w)y'' (dwp, d (v(.),
fe(.) fu()= worl :; ;;)) dw) =
so that in turn we have ny Te a.e.. In other words, the homogeneity of the chosen class of approximators results in degeneracy of the map vi : (f) + (S x A). To remove this degeneracy in our analysis, we identify all the distributions v, // that are equivalent under (C.4) by defining the signed measure
h. : / (dwr d
(C.5)
Leveraging this definition, prove the desired result Theorem 4.2 in two key steps: show that
the solution to (10) does not lose (projected) support for any finite time, thereby preserving the property from Assumption c),
2. stationary points & with QnD * Qx- which by Lemma C.I cannot have full projected support in θ are avoided by the dynamics.
18

---

Published a conference paper at ICLR 2021
These facts are respectively summarized in the following lemmas:
Lemma C.2. Let Assumption a) hold and let vo satisfy Assumption 1 c). then for every t > Vi solving (10) with initial condition Vu also satisfies Assumption 1 c).
Throughout, we let
Bb denote the bounded Lipschitz norm.
Lemma C.3. Let Assumption I hold and let i be a fixed point of (10) such that (C.3) does not hold a.e.. Then
there exists & > 0 such that if [1b] - -hi.,, |Br.. < f a 1 >> 0 there exists l2 >> i such that [1]] -iig-liy Rii | > e.
Proof of Lemma C.2. Analogously to (Chizat & Bach, 2018, Lemma C.13), we aim to show that the separation property Assumption c) is preserved by the evolution of va along the characteristic curves X (t, u) solving
(,, u) F:(((t, u)),
(C.6)
where F; is the transport field in (10). To reach this conclusion, the analogous result in Chizat & Bach (2018)
only relies on the continuity of the map u + X (t, u), established in (Chizat & Bach, 2018. Lemma B.4) under
Assumption 1 a). Hence, it is sufficient for our purposes to establish continuity of the map (t,, .) from (C.6). This property, however results immediately from the one- sided Lipschitz continuity of the vector field F; on 2, --,, r| O uniformly on compact time intervals, which is in turn guaranteed by the Lipschitz continuity and Lipschitz smoothness of & from Assumption and boundedness of r.
To simplify the notation in the following proof, we denote throughout
8(v) = Q5 - T .logg π. - VT and (f.9)x == [ a)s,a) (),(s da)e
JSXA
Proof of Lemma C.3. We first claim that by Lemma C.I, for any spurious fixed point (such that Qu * 2, there must exist a subset of θ with positive Lebesgue measure where i loses support and such that
(Vú, (())) v * O. This is easily proven by contradiction: if (V)), = 0 a.e. then by Lemma C.I we have
that QTo = Qu. This implies that the quantity
go (w) == (0.. ((::))
= (vi: (1, ()) ((D)) = (d(::)), (())
(C.7)
cannot vanish a.e. on O. Then, by Assumption on the regularity of g. there exists nonzero regular value -n of go (w) Assuming without loss of generality that this regular value is negative, so that 17 > 0 (else invert the signs of wa in the remainder of the proof), we define the nonempty sublevel set G = {w0,, w) E Ω go(w) - m1 and
G- = {(wo, @) E G : wa > 0}
(C.8)
Further denoting by G S ④ the projection of 9 onto θ, we have by definition that the gradient field of go is orthogonal to the level set aG, the latter being an orientable manifold of dimension - 2. Denoting by na the normal unit vector to OG the outward direction, by continuity of Vga(w) when G is compact! we can bound the scalar product between the two away from 0, ie., there exists
B = min The EEAG
Vago(w)> > 0.
We now prove that the stationarity assumption in a s-neighborhood of a spurious fixed point
|/ - h,, BL < for all t >> t
(C.9)
leads to a contradiction for small enough. To do so, by Lemma C.5 we set e(a, 17, 3) small enough so that for all 24 such that (C.9) holds we have gv, (w) < - on G and na Vagv, > B/2 on aG. Then, the two inequalities above combined with d_,, (wp, w) = e111, @) imply that the set G4 defined above is forward invariant and therefore that 0,v, (G+) 2 0 as long as (C.9) holds. Furthermore, by similar arguments we notice that characteristic trajectories cannot enter the set G \ G+ after 4,
Now, we consider two cases: either (i) a positive amount of mass is present at i in the forward invariant set G (vt, (9+)= > 0) or (ii) Mr, (G+) = 0. We discuss these two cases separately, along the lines of (Chizat & Bach, 2018, Lemmas C.4, C81 respectively.
1if G is not compact we choose v to also be regular value of the function on (a R= ||a||2 = 13 to which g converges as @ goes to infinity.
19

---

Published a conference paper at ICLR 2021
(i Assume that Mr, (G+) > O. We note that under our assumptions the first component of the velocity field in G is lower bounded by n/2, so that wa(t) (0) + tn/2 bounds from below the wuy-component of the trajectory of a test mass with initial condition with (0) E g. as long as ((t) E G. Combining this bound with the forward invariance of G.. see that if (0) G+ then (t(() > tn/ 2. Consequently, assuming that supp((s) C (-M, M) x + for every > 4 we have
heg (G) 2 2/2 - t1) (9+) + min [0, (t - f2)n/2 - M (9) (9+).
This implies linear growth of h, (G) for t > f 2M /n. contradicting the original assumption that [1g My. ns. <e for all I 4--
(ii) Consider now the complementary case Vt, (G+) = O. We proceed to show that there exists t2 > t such that Vt,, (9+) > 0. thereby reducing this case at time t to part (i). To do so, we consider w' E supp(?t,) ) such that Th E G is local minimum of ga. i.e., for which Vgs- = 0 (which exists by the preservation of the support property Assumption I c) Then, choosing E such that (w())) C G, and setting M large enough that supp (vx,) S L|,M, M] x O we prove below in Lemma C.4 that there exists l2 > (1 for which the image at t2 of (t1) == wn under the characteristic flow (C.6) is contained in G+. By continuity
of the flow map X(,, t), this conclusion extends to a neighborhood of w", with positive mass under Vt,-
We denote throughout by || lla the maximum of the supremum norm of function and the supremum norm of its gradient and recall the structure of the policy gradient vector field
Fr(w;, vn) = - v (wo)(W), ((v))) = (ww Iv, (w))
(C.
where y is defined in (C.7) and i, solves (10).
With these definitions, we proceed to prove that case (ii) in the analysis above will ultimately reduce to case (i) for large enough.
Lemma C.4. Let E M1 (S2) and 3- satisfy g (a+)) O, go())) < -n < 0 for some 1 > Then for every h, M >> 0 there exists t2, > 0 such that for all t E (0,t2) we have 1190 gu, llc. < e and wß E [-M, Oj, then the point w" is mapped, under the flow of the policy gradient vector field (C.10) at time 12 to a subset of Be((1,@"))
Proof of Lemma C.4. By homogeneity of the approximator, we can bound the first component of the velocity of particle (wo (f), (t)) under (C.10) with initial condition (0) = 5 as
dt (t) -9,, (t) 2 (w= - ())) - ga [س(=") - lgv, (o))) - gu (1)))
In the other directions, defining ((t) := (t) [=, we have
d dt (1) s (10l (=(1))
S (t) VV (o (x") + |Va (o()) - Vaga(w")] + VVag,,
- (age((()))ll
for all t E [0, where F = inf wo(() & 11], Moreover, Lipschitz continuity of the potential go (.()
and its Lipschitz smoothness imply the existence of a L. > 0 such that max (*) go (w"),. g (w)
Vga( (l)) S T -n|l_ Combining this with the assumed convergence of v4 to i, which implies 9
Iv, Ilo < e, we can bound the evolution of (wo (()) for t E [0, F] in the perturbative regime of interest as follows:
di woll) d
Lq(1)
(C.11)
dt (t) S [woo()) + Lq(t))
(C.12)
We now show that, choosing both e and neighborhood around w = (wr w+) to be small enough, the forward dynamics of w will reach the set (wo > 0} before q(t) can increase too much. More precisely, by possibly
increasing the value of such that 4/ غ, and defining T4 = inf q(t) > m/AL) we prove that there
20

---

Published a conference paper ICLR 2021
exists E E n/44 such that T > F, i.e, that the trajectory w" reaches G4 before g(t) > É. Note that as long as t t Tg and E (0, 1/44 (4) the negative terms on the RHS (C.11) can be bounded from below, and we have
wi (1) 2 wo (0) + 2'
d so that wo(t) > wo(() 2 -M. Consequently, for all t [0, A Tgl we bound the RHS of (C.12) as de g() < ME+ + LM) Using that q(0) O and Grönwall inequality we can bound the total excursion the i component g(4) I exp [LM]]. Finally, setting TO == 2(M + 1)/n Z -2(wo(0) - 1)/n > F so that wo(Tn) > 1 we are still free to set e small enough such that > TO > T. Indeed, by monotonicity of the upper-bound on (()) we have
(T)) S 2 + 1)M/n exp (M + 1))/ml s m/44,
so that setting e E (/4) concludes the proof.
We now proceed show the needed regularity of the potential 9v from (C.7) in terms of the signed measure h, defined in (C.5), In doing so, we also prove Lipschitz smoothness of the operator R defined in (B.1):
Lemma C.S. For any > O, the operator R on F, = {Jf s/u supp E 2r is Lipschitz smooth, and D DR, is bounded in the supremum norm. Furthermore, for all Co there exists a and > 0 such that for all v satisfying [1,,|BL, || < Co. 1 hByBBL < E, one has
119~ g lla s a a h IIL.
(C.13)
To prove the above lemma, we first bound some relevant quantities. Throughout, by slight abuse of notation, we denote for any function S A R
Ifll = sup / f(s, as) sES A
Lemma C.6. For all f. Fr there exists E 0, C² C" C! 0 such that if llf f'll2 has
1| T. S C/I| -1'12
(C,14)
lle. 11 S C"IIS - -2
(C.15)
11Q - 11, S C"IIlf -f'll2
(C.16)
Proof of Lemma C.6. Throughout this proof, for simplicity of notation, we will write Tr and n' Tyl Furthermore, we use that for f E Fr there exists Co so that
e
Callóblle < explf(s, a)) eColólee.
(C.17)
implying, together with the assumed absolute continuity of go that [| || 1| < co. Setting throughout
simplify the notation and combining the above with the pointwise upper bound e. 1 + K,| |z| for [z|| < Collélles obtain
exp[f(s,
explf'(s, a) expls'(s, da expls' a)
|| 7
expls[es a')]da'
expls(s,a)
expls'(s, exp(
J explf(s,a')daa
cxplf((s a dd
explff al d')d Texplf'ss
explf((s.
(11- [[1-expif(s,d)) (f f's + | (s, - 11)
eexp ass)d'
J
TexpIf(s, a'))da
2Colollee
2Collllee
- (s, 1)da 1))
(KIlff(sa) f's, a)112
AA
|A]
2Caldllee
K,(1 + 4Collo lo (C11+c40olle
Mlf -f'll2 C'Il. -f'lla
(C.18)
[A
21

---

Published a conference paper at ICLR 2021
where we have denoted by |A| the Lebesgue measure of the action space A. We now proceed to establish the second bound in the statement of the lemma. In this case, denoting the t-steps transition probability as (s, = Sst Pr (s, Ps (s1, d82). ... Pa (s) ds;) we have
1l Px 111 =
co(dso) (Pu ds) P (su, ds))
06
t-1
e0 Po Pj (Pr - Pr.) P
-
j=0
Observing that for any smooth e E M (S) for the operator norm of the difference in the above sum we have
/ e(ds) / P(s, a, ds') (TTs, da) - x (s, da))
lle(Pss - P
(1-2""0'11 C"Ilf flll2
|ell|| l 5
?
for large enough C", where used (C.18) in the last line and the Lipschitz continuity of P in its second argument, and [|P|| is the operator norm of the transition operator SΛ P(s, a, ds') "ree)r"(s,, da) .M (A) +
M4 (A), which is equal to 1. From this we conclude
(1 7)2
lle+ em 112 S
Σ
f211
Finally, defining for notational convenience R₁ : TDKL(T' R) we write:
> / d(s.c).s)) -ve
1|x, Q 112
|>> / P(s,a,ds') (f (st R,,("", ,''ie'"""""" (s', (d. ||
$ > / als, [R d')e" ((dd) - R=(s", d')e" (s', ds" {ddd))
$ / (Rl - RR) (8", d'))o" (s', as mddd' + -11 a'''",")(e"" (s', ds (Td'))
11 1 ('") (8, ds)))" x x'dd"))
(C.19)
and bound each term separately letting CI",O"" Ctal > 0 be large enough constants. For the first, we have:
|ʃ 1/(R=- (R, (2" a do""")))) da')
1
|| |2el log π - log 7' 1|2
1
I eel( da
<
|||||2
[|. s'll2 + ||og j ee (sial 1|2
1
< 11 - -1'
(C.20)
where we have bounded the log term as follows:
f eal
|| log f fellon)) - - (s
|oog
I da ||2
(el + +1/12
f (el(sa)) (c)((en)) 1) da
log
1
Sello.a
22

---

Published conference paper ICLR 2021
le le- 110 "1o0|1e 1000 / (cltiss)) - 1) da +1 ll2
log
A
|e 11 ||ee 1lo1eelloo 1lo0 K, [f(s.))- f(s ll2
log
[A|
log
|le |le-5" |1 11,0 Krll((((s ('( + +1) ll2
x |le |l 11 K-lls - 111l2
(C.21)
For the second term in (C.19), using the boundedness of ||/0||2 < On and that e" e-' 7 (er ewr) for certain, do depending on s' we write
| |2 d') e² (s' ds" (=(",, da""
1 R2 ,d))ll2 le 2rr
S C7 I| - 11
(C.22)
We finally bound the third term by writing
// RC (s ew (s', dod (T - xt (3) da')
- I 1l fl11
",",,,")))) dllll= - 2'1l2
(C.23)
and obtain (C.16) by combining (C.20)
Proof of Lemma C.5. We first establish the desired properties the functional R. To do so we differentiate (8)
at f E (S x A)
e J f(x.a)u(dw)
8 (Laer)da)
eτf
/
σʃ (5,0) = ō. [A eTIlaaaiaa IA oT et
e'a
dv
-t [
()))
(C.24)
Then, combining the above with (A.6) and Hölder inequality we obtain
|D sup / (DR)(ss dsda < jeF,
dsda
where we have used that all the terms appearing in DR2 are bounded for every choice of > 0. To establish Lipschitz smoothness of R, letting f' E F, and denoting, to simplify notation, æ TT and n' t'',, we proceed bound the operator norm by splitting the RHS as
||DR, -Ry S
sup
[C Qn 7 Flog ক| Crr Qr 7 log T'1l
(EL2(SXA)
||
< sup [(1) - (11) + (111)
(Ex)(SxA) ||
and considering the resulting terms separately. First all, defining throughout by slight abuse of notation 8(T) = Qn - T log n, we have
(1) (1) = -K=K/ / ()s, (s, da) - / (() ((x')))
/ ( ((s, ())..))(((s))) ( r(s, (/ (( (((r)))(((a))) e (s)
||[|||||= - "=
x C[][[|llff fl12.
(C.25)
23

---

Published conference paper ICLR 2021
where the last line was obtained using (C.14) and boundedness from above and below n, Qx Fr. We further write, for K < D0 large enough
(1T) = (((s, / l(s,,
- ( - f ((s,
(s, +1 ((s, (xr)))
(s) (s)dss
- [10(T)) l1llee|aa (1 - 2 + ee ex
| f f'll2,
(C.26)
where have used Mauchy-Schwarte inequality together with (C.14) and (C.15). Finally, using (C.16) and (C.21), we bound for K < oo possibly larger than above
/ ((s, ((s, da'), 8(T) )((T))
(111) (((s,,))
- 5(z')112
< ||<|||||=|||||00| (12, Qr + ||V - 122 + T log (12)
||=1111/~1|||000 (11 || Qr + ||0 ~
- 7 log - l12))
K |||e||/|s 12
(C.27)
Combining (C.25), (C.26) and (C.27) we obtain that
|DR, DR 0c f.EET |If -
LDR = sup
proving the Lipschitz smoothness claim.
Proceeding the proof of (C.13), combining the Lipschitz smoothness of R bounded sets, the identity (ss 0; (1,
= ((s a; () and the boundedness of the set { v/v v E P2( |h!l < Co} we obtain
/ ((:w)(u-
Il - full2
(C.28)
-hb (oww))) hb) (dz)
sup ( / (1 (dw)
< || !|, - BL,
which combined with the -Liscchtt continuity of the map f f(ss, a)d(s, a ddsdd and with the Lipschitz smoothness of R concludes the proof.
24


Reviews:
[{"review_id": "sllISWu9cTf", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "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.\n\nPros:\n\n1. This paper introduces the mean-field formulation into the reinforcement learning framework. The technical proof seems highly challenging.\n2. 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.\n\nCons:\n1. 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.\n2. The paper only shows convergence, but fails to give more detailed properties like convergence rate, etc. \n", "extra_scores": {}, "extra_text": {}}, {"review_id": "rY5wzU8UYzp", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "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.\n\nOverall, 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:\n\n- 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.\n\n- 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.", "extra_scores": {}, "extra_text": {}}, {"review_id": "mvNw91nauQS", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "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. \n\nAlthough 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. \n\nThe 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?", "extra_scores": {}, "extra_text": {}}, {"review_id": "3Fp7-D8ou1j", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "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.\n\nThe 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.\n\nOn 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.\n\nAdditionally, 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.", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

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)