Wasserstein Flow Meets Replicator Dynamics: A Mean-Field Analysis of Representation Learning in Actor-Critic
Authors: Yufeng Zhang
Siyu Chen
Zhuoran Yang
Michael Jordan
Zhaoran Wang
Venue: NeurIPS 2021 Poster (2021)
Keywords: reinforcement learning
policy optimization
optimal transport

Abstract: Actor-critic  (AC) algorithms, empowered by neural networks, have had significant empirical success in recent years. However, most of the existing theoretical support for AC algorithms focuses on the case of linear function approximations, or linearized neural networks, where the feature representation is fixed throughout training. Such a limitation fails to capture the key aspect of representation learning in neural AC, which is pivotal in practical problems. In this work, we take a mean-field perspective on the evolution and convergence of feature-based neural AC. Specifically, we consider a version of  AC where the actor and critic are represented by overparameterized two-layer neural networks and are updated with two-timescale learning rates. The critic is updated by temporal-difference (TD) learning with a larger stepsize while the actor is updated via proximal policy optimization (PPO) with a smaller stepsize. In the continuous-time and infinite-width limiting regime, when the timescales are properly separated, we prove that neural AC finds the globally optimal policy at a sublinear rate. Additionally,  we prove that the feature representation induced by the critic network is allowed to evolve within a neighborhood of the initial one. 

Wasserstein Flow Meets Replicator Dynamics: A Mean-Field Analysis of Representation Learning in Actor-Critic
Yufeng Zhangt
Siyu Chent Tsinghua University
Northwestern University (rtengghhnng033000.nnottttttttte edu
chensy180malls wtt.://0mmnlsssssnnnnaaaae edu. cn
Zhuoran Yang Princeton University zy6@princeton. edu
Michael I. Jordan
Zhaoran Wang
UC Berkeley
Northwestern University
jordanEcs. www.eeer..eeeeoooeeoooo edu
zhaoranvang@gaiil.com com
Abstract
Actor-critic (AC) algorithms, empowered by neural networks, have had signifi- cant empirical success in recent years. However, most of the existing theoretical support for AC algorithms focuses on the case of linear function approximations, or linearized neural networks, where the feature representation fixed through- out training. Such a limitation fails to capture the key aspect of representation learning in neural AC, which is pivotal in practical problems. this work, we take a mean-field perspective on the evolution and convergence of ceature-based neural AC. Specifically, we consider version of AC where the actor and critic are represented by overparameterized two-layer neural networks and updated with two-timescale learning rates. The critic is updated by learning with larger stepsize while the actor is updated via proximal policy opti- mization (PPO) with a smaller stepsize. In the continuous-time and infinite-width limiting regime, when the timescales are properly separated, we prove that neural AC finds the globally optimal policy at a sublinear rate, Additionally, we prove that the feature representation induced by the critic network is allowed evolve within neighborhood of the initial one.
(TD)
Introduction
In reinforcement learning (RL) [56], agent aims learn the optimal policy that maximizes the expected total reward by interacting with the environment. Policy-based RL algorithms achieve such goal by directly optimizing the expected total reward as function of the policy, which often involves two components: policy evaluation and policy improvement. Specifically, policy evaluation refers to estimating the value function of the current policy, which characterizes the performance of the current policy and reveals the updating direction for finding better policy, which is known as policy improvement. Algorithms with these two ingredients also called actor-critic (AC) methods [36], where the actor and the critic refer the policy and its corresponding value function, respectively.
Equal contribution.
35th Conference Neural Information Processing Systems (NeurIPS 2021).

---

Recently, in RL applications with large state spaces, actor-critic empowered by expressive function approximators such as neural networks have achieved striking empirical successes [3, 4. 9, 20, 51, 52, 60J. These successes benefit from the data-dependent representations learned by neural networks. Unfortunately, however, the theoretical understanding of this data-dependent benefit is very limited. The classical theory of AC focuses on the case of linear function approximation, where the actor and critic are represented using linear functions with the feature mapping fixed throughout learning [10, 11, 36]. Meanwhile, few recent works establish convergence and optimality of AC with overparameterized neural networks [26, 39, 61], where the neural network training is captured by the Neural Tangent Kernel (NTK) [30]. Specifically, with properly designed parameter initialization and stepsizes, and sufficiently large network widths, the neural networks employed by both actor and critic can be assumed to be well approximated by linear functions of random feature determined by initial parameters. In other words, concerning representation learning, the features induced by these algorithms are by assumption infinitesimally close to the initial featural representation, which is data-independent.
In this work, we make initial steps towards understanding how representation learning comes into play in neural AC. Specifically, we address the following questions:
Going beyond the NTK regime, does neural AC provably find the globally optimal policy? ? How does the feature representation associated with the neural network evolve along with neural AC?
We focus on a version of AC where the critic performs cemporal-differeee (TD) learning [55] for policy evaluation and the actor improves its policy via proximal policy optimization (PPO) [49], which corresponds to Kullback-Leibler (KL) divergence regularized optimization problem, with the critic providing the update direction. Moreover, we utilize two-timescale updates where both the actor and critic are updated at each iteration but with the critic having a much larger stepsize. In other words, the critic is updated at faster timescale. Meanwhile, we represent the critic explicitly as a two-layer overparameterized neural network and parameterize the actor implicitly via the critic and PPO updates. To examine convergence, we study the evolution of the actor and critic in the continuous- time limiting regime with the network width going to infinity. In such regime, the actor update is closely connected replicator dynamics [12, 28, 50] and the critic update is captured by semigradient flow in the Wasserstein space [59]. Moreover, the semigradient flow runs at faster timescale according to the two-timescale mechanism.
It turns out that the separation of timescales plays an important role in the convergence analysis. In particular, in the continuous- time limit, it enables us to first separately analyze the evolution of actor and critic and then combine these results to get final theoretical guarantees. Specifically, focusing solely on the actor, we prove that the time-averaged suboptimality of the actor converges sublinearly to zero up to the time-averaged policy evaluation error associated with critic updates. Moreover, for the critic, under proper regularity conditions, we connect the Bellman error to the Wasserstein distance and show that the time-averaged policy evaluation error also converges sublinearly to zero. Therefore, we show that neural AC provably achieves global optimality at sublinear rate. Furthermore, regarding representation learning, we show that the critic induces a data-dependent feature representation within an (11/a) neighborhood of the initial representation in terms of the Wasserstein distance, where a is a sufficiently large scaling parameter.
The key to our technical analysis reposes on three ingredients: (i) diinite-dimensionall variational inequalities with a one-point monotonicity [27], (ii) mean-field perspective neural networks [19, 41, 42, 53, 54], and (iii) the two-timoe -wo-timescale stochastic approximation [13, 37J. In particular, in the infinite-widtt limit, the neural network and its induced feature representation are identified with a distribution over the parameter space. The mean-field field perspective enables us to characterize the evolution of such a distribution within the Wasserstein space via partial differential equation (PDE) [5. 6, 58, 59]. For policy evaluation, such a PDE given by the semigradient flow induced by TD learning. We characterize the error of policy evaluation by showing that mean-field Bellman error satisfies version of one-point monotonicity tailored to the Wasserstein space. Moreover, our actor analysis utilizes the geometry of policy optimization, which shows that the expected total reward,

---

as function of the policy, also enjoys the property of one-point monotonicity in the policy space. Finally, the actor and critic errors are connected via two-timescale stochastic approximation. To the best of our knowledge, this is the first time that convergence and global optimality guarantees have been obtained for neural AC.
Related Work. AC with linear function approximation has been studied extensively in the literature. In particular, using two-timescale stochastic approximation via ordinary differential equations, limiting regime.
[10, 11, 36] establish asymptotic convergence guarantees in the continuous More recently. using more sophisticated optimization techniques, various works [29, 35, 64-66] have established discrete-time convergence guarantees that show that linear AC converges sublinearly to either a stationary point or the globally optimal policy. Furthermore, when overparameterized neural networks are employed, [26, 39, 61] prove that neural AC converges to the global optimum at sublinear rate. In these works, the initial value of the network parameters and the learning rates are chosen such that both actor and critic updates are captured by the NTK. In other words, when the network width is sufficiently large, such version of neural AC is well approximated by its linear counterpart via the neural tangent feature. In comparison, we establish mean-field analysis that has different scaling than the NTK regime. We also establish finite-time convergence to global optimality, and more importantly, the feature representation induced by the critic is data-dependent and allowed to evolve within a much larger neighborhood around the initial one.
Furthermore, our work is also related to the recent line of research on understanding stochastic gradient descent (SGD) for supervised learning problems involving an overparameterizee two-layer neural network under the mean-field regime. See, e.g., [16, 19, 21, 22, 31, 40-42, 53, 54, 63] and the references therein. In the continuous-time and infinite-width limit, these works show that SGD for neural network training is captured by a Wasserstein gradient flow [5, 6, 59] of an energy function that corresponds to the objective function in supervised learning. In contrast, our analysis combines such a mean-field analysis with TD learning and two-timescale stochastic approximation, which are tailored specifically to AC. Moreover, our critic is updated via TD learning, which is semigradient algorithm and there is no objective functional making TD learning gradient- algorithm. Thus, in the mean-field regime, our critic is given by Wasserstein semigradient flow, which also differs from these existing works.
Additionally, our work is closely related to [1, 69]. who provide mean-field analyses for neural TD-learning and Q-learning [62]. In comparison, we focus on AC, which is a two-timescale policy optimization algorithm. Finally, [2] studies softmax policy gradient with neural network policies in the mean-field regime, where policy gradient is cast as Wasserstein gradient flow with respect to the expected total reward. The algorithm assumes that the critic directly gets the desired value function and thus the algorithm is single-timescale. Moreover, the convergence guarantee in [2] is asymptotic. In comparison, our AC is two-timescale and we establish non-asymptotic sublinear convergence guarantees to global optimality.
Notation. We denote by (X') the set of probability measures over the measurable space X. Given curve p IR- + X, we denote by P. = Oept. tms its derivative with respect the time. For an operator F X' + X' and measure μ E (X'), we denote by FęJ = o F-1 the push forward of μ through F. We denote by (²)) the chi-squared divergence between probability measures p and μ. which is defined as x²(x (1 f(// - 1)ddu. Given two probability measures and ju, we denote the Kullback-Leibler divergence or the relative entropy from μ to p by KL(p || (4) = log(e//)dd For 1/1, #2 (X),, we define the (()) weighted homogeneous Soboley norm as [M [p..) Ip... (f ((dz)) v2llæ 1(x) = sup {Kff. - re)llllfll I 1]. We denote by [I.
the 6 -n-or with respect to probability measure We denote by @ the semidirect product, i.e.,
8 K = Kyy x))(x) for Ja E (X) and transition kernel K : X + (P)). For a function f:X X + IR, we denote by fip(f) = sup Supx. [f(x)- - fy)ll/////=xxx-- - yl its Lipschitz constant. We
denote normal distribution on R D by (p, E), where is the mean value and Σ is the covariance matrix.
3

---

2
Background
In this section, we first introduce the policy optimization problem and the actor-critic method. We then present the definition of the Wasserstein space.
2.1 Policy Optimization and Actor-Critic
We consider a Markov decision process (MDP) given by (S, P,r, Do), where S S R is the state, A C Rda is the action space, 7 E (0, 1) is the discount factor, P S x A + (SS) is the transition kernel, r S x A + R is the reward function, and Do (S)) is the initial state distribution. Without loss of generality, we assume that S x A C R and [I(s, a) S 1, where d1 + d2. We remark that as long as the state-action space is bounded, we can normalize the space to be within the unit sphere. Given a policy takes an action am at state 8m according T(- and observes reward I'm = r(mm am). The environment then transits the next state 8m++1 according to the transition kernel P(- 11,mm am). Note that the policy induces Markov chains both S and S x A. Considering the Markov chain on S, we denote the induced Markov transition kernel by p : S + (S), which is defined as (s'(s () [ P(s' s, a) (a)a((a)))) s). Likewise, we denote the Markov transition kernel on S x A by P" S x + (S x A), which defined as pa (s', d' |s, a) = x x(d' P(s [ss a). Let D be a probability measure on A. We then define the visitation measure induced by policy ι and starting from D
A + (S), at the mtth step, the agent
(d(s,a)) = (1 (1) Σ
nth th .((mm am) E d(s, a) (so, 0o) DD)
(2.1)
m20
where (Sm; am) is the trajectory generated by starting from (so, an) D and following policy thereafter. If D D @ a holds, we then denote such visitation measure by & Furthermore, we denote by {ds) = [A Edss, da) the marginal distribution of visitation measure & with respect to S. In particular, when (sg, 0o) optimization, we aim maximize the expected total rewards (TT) defined follows,
@ x holds in (2.1), it follows that Eg EB @ æ. In policy
J(T) = E"
Σ 7 r(Sm, 4m) so ~ Do
m20
where we denote by E the expectation with respect to x(. (m) and 8m+1 ~ P(- Sm: Qm) for m 2 0. We define the action value function Q S x A + IR and the state value function
V== : :S + R as follows,
Q (s. = ET
Σ ymh rr(m am) so -m2
s, ag = a
(*(()= (()))))(( ( (2.2)
where we denote by (,, ) the inner product on the action space A. Correspondingly, the advantage
function A" S x A + IR is defined as
AT a) = Q (s.
(T()).
It is known that the action value function Q is the unique global minimizer to the following mean-squared Bellman error (MSBE),
MSBE(Q: _ - (Q(s,a) r(s, a) - - "
(2.3)
where कhth a weighting distribution depending on policy x and is with full support, i.e., supp S A. Therefore, the policy optimization problem can be written as the following bilevel optimization problem,
max ((T) = E.NDO [ (s, (, (.1))) subject to Q = argmin MSBE(Q; m).
(2.4)
Q
The inner problem in (2.4) is known as policy evaluation subproblem, while the outer problem is the policy improvement subproblem. One of the most popular way to solve the policy optimization problem is actor-critic (AC) methods [56], where the job of the critic is to evaluate current policy and then the actor updates its policy according to the critic's 's evaluation.

---

2.2
Wasserstein Space
Let θ C DRP be a Polish space. We denote by (2) S. ((0) the set of probability measures with finite second moments. Then, the Wasserstein-2 (W2) distance between μ, v E P2(() defined as follows,
W2(u, -in)=iff[[[[I- - law(X) μu, law (Y) = =
where the infimum taken over the random variables X and on O and we denote by law(X) the distribution of a random variable X. We call M (0)) W2) the Wasserstein (W2) space, which an infinite-dimensional manifold [59]. See $A.1 for more details.
Algorithm
Two-timescale Actor-critic. We consider two-timescale Actor-critic (AC) algorithm [34, 45] for the policy optimization problem in (2.4). For policy evaluation, we parameter the critic Q with neural network and update the parameter via temporal-difference (TD) learning [55]. For policy improvement, we update the actor policy æ via proximal policy optimization (PPO) [49]. Our algorithm is two-timescale since both the actor and critic are updated at each iteration with different stepsizes. Specifically, we parameterize the critic Q by the following neural network with width M and parameter = (0(1),. ... ((M)) E DDX
M
Qg(s,a) = M
a
Σ (s. a; 0(1).
(3.1)
Here (s, a; S A RD - R is the activation function and a > is the scaling parameter. Such a structure also appears in [17, 18, 41]. a descrete-time finite-width (DF) scenario, at the kth iteration, the critic and actor are updated follows,
DF-TD: (1 = 0% (Qā. (sk: ax) r(sk, ax) - -2 (5 V Voo 0
(()))
(3.2)
DF-PPO: R (s(s) argmax (Qa. (s (.(())))) - - (LL((LL(T) 11)
(3.3)
where (sk, a&) Th and (s/, of.)) pa sky ax). Here म.. the policy for the actor at the kth iteration, &h is the corresponding weighting distribution, E and are the stepsizes for the DF-PPO update and the DF-TD update, respectively. In (3.2), the scaling a arises since our update falls into the lazy- werytirannin regime [18]. In the sequel, we denote by 1) = E'/e the relative TD timescale. Note that in double-loop AC algorithm, the critic can usually be solved with high precision. In the two-timescale AC however, even with the Kl--divergence term in (3.3) which regularizes the policy update and helps to improve the local estimation quality of the TD update, the critic Qù for updating the actor's policy T. can still be far from the true action value function Q+. Since the policy evaluation problem is not fully solved at each iteration, the two-timescale AC can be more efficient in computation while more challenging establish theoretical guarantee.
Mean-field (MF) Limit. To analyze the convergence the two-timescale AC with neural networks, employ the analysis that studies the mean-field limit regime [41, 42). Here, by saying the mean- field limit, we refer to the infinite-width limit, i.e., M - for the neural network width M in
(3. 1),, and the continuous-time limit, i.e., / ks where & -> 0 for the stepsize in (3.2) and (3.3). For 0 = [ J==1 independently sampled from a distribution we can write the infinite-width limit of (3.1) as
Q(s,a:p) a a o(ss ()d()))
(3.4)
In the sequel, we denote by Pe the distribution of 7(() for the anfnite-width limit of the neural network at the kth iteration. We further let Pt and be the continuous-time limits of Pa and âg. respectively.
5

---

As studied in [69]. the mean- field limit of the DF-TD update in (3.2) is
MI-TD:
Dept -1 (t,
(3.5)
where 1) is the relative TD timescale and
g0:: - 3- ap)
- Q(s's a -a r(s, -
(3.6)
is a vector field. Here -- is taken with respect to (s,a) a) ~ कh and (s', (s',a')= (* a,)) It remains to characterize the mean-field limit of the DF-PPO update in (3.3). By solving the maximization problem in (3.3), the infinite-width limit of DE-PPO update can be written in closed form as
log [ 6)] log [ax sS]] = Q8,, (z) Zs(s),
where Zx(s) is the normalizing factor such that f f 8) = 1 for any s S. By letting ke and E - O, we have O, log n. Qt Ze which can be further written as 0, = n. - (Q; Zz). Here we have Qe(a, s) = Qa, 8; pr) and Ze is the continuous-tin limit of Żg. Furthermore, noting that 0g (e(da 8) = 0, the mean-field limit of the DF-PPO update in (3.3) is
d
where Ar(s, a) Qc(s,a) - f or (da
MF-PPO:
di πe = Te Ar,
(3.7)
The two updates (3.5) and (3.7) correspond to the mean-field limits of (3.2) and (3.3), respectively, and together serve as the mean-field limit of the two-timescale AC. In particular, remark that the MF-TD update in (3.5) for the critic is captured by semigradient flow in the Wasserstein space [59] while the MF-PPO update in (3.7) for the actor resembles the replicator dynamics [12, 28, 50). Note that such a framework is applicable to continuous state and action space. In this paper, we aim to provide a theoretical analysis of the mean-field limit of the two-timescale AC.
4
Main Result
In this section, we first establish the convergence of the MF-PPO update in $4.1. Then, with additional assumptions, we establish the optimality and convergence of the mean-field two-timescale AC in $4.2.
4.1 Convergence of Mean-field PPO
For the MF-PPO update in (3.7), we establish the following theorem on its global optimality and convergence rate.
Theorem 4.1 (Convergence of MF-PPO), Let x" = argmax, [(x) be the optimal policy and xa be the initial policy. Then, it holds that
T // (J(=*) - -((t))J((z)))))
5 4x T / 11Q Q 12 dt, T
1
(4.1)
policy evaluation error
where Tev E (S x A) is an evaluation distribution for the policy evaluation error and Ç = E E KL(T* x (-|s) no (-]s)) is the expected KL-divergence between nt" and To. Furthermore, Do_ letting कt 1 فu & Te. where Φo E (S A) is base distribution and Φo [A 5,(, da). the concentrability coefficient is then given by
Proof. See $ for detailed proof.

---

The concentrability coefficient commonly appears in the reinforcement learning literature [7, 23, 24, 38, 39, 43, 48, 57, 61]. In contrast to a more standard concentrability coefficient form, note that is irrelevant to the update of the algorithm. To show the convergence of the MF-PPO, our condition here is much weaker since we only need a given base distribution do such that K < co.
Theorem 4.1 shows that the MF-PPO converges the globally optimal policy rate of O(T-1) up to the policy evaluation error. Such theorem implies the global optimality and convergence of double-loop AC algorithm, where the critic Qe is solved to high precision and the policy evaluation error is sufficiently small. In the sequel, we consider a more challenging setting, where the critic Qt is updated simultaneously along with the update of the actor's 's policy Tt.
4.2 Global Optimality and Convergence of Two-timescale AC
In what follows, we aim to characterize the upper bound of the policy evaluation error when the critic and the actor are updated simultaneously. Specifically, the actor is updated via MF-PPO in (3.7) and the critic Q, (=: pt) updated via the MF-TD in (3.5). For the smooth function o in the parameterization of the function in (3.4), we consider it to be the following two-layer neural network,
(8, a; 0) = B3 -(6) ((w (s, 0, 1));
(4.2)
where ā -R is the activation function, an odd and invertible function with scaling hyper-parameter B3 0. It then holds that D d d where d and D are the dimensions of (s.a) a) and 6, respectively. It is worth noting that the function class of f o a; ((d())) for PRR is the same
(b, w) the parameter, and B : R + (- 1, (1,, is
F =={ 3 3 (s, a, udd's dw) v P2((-BB,B)) x }
(4.3)
which captures a vast function class because of the universal function approximation theorem [8, 47]. We remark that we introduce the rescaling function 33 in (4.2) to avoid the study of the space of probability measures over (-B3, B3) Rd+1 in (4.3), which has boundary and thus lacks the regularity in the study of optimal transport. Furthermore, note that we introduce hyper-parameter
> 1 in the Q function in (3.4). Thus, we are using a - F to represent F, which causes an "over- representation" when > 1. Such over-representation appears to be essential for our analysis. For a brief peek, we remark that a actually controls the gap in the average total reward over time when the relative time-scale 77 is properly selected according to Theorem 4.6. Furthermore, such an influence imposed through Lemma 4.4, which shows that the Wasserstein distance between Po and Pr. is upper bounded by (1//x). In what follows, we consider the initialization of the TD update to be Pn (0, ID), which implies that Q(s, a; po) 0. We next impose the following assumption on the two-layer neural network a.
Assumption 4.2 (Regularity of the Neural Network). For the two-layer neural network defined in (4.2), we assume that the following properties hold.
(i) The rescaling function 3 : R + (-1, is odd, Lipschitt continuous, L1,3 smooth, and invertible. Meanwhile, the inverse 3-1 locally Lipschitz continuous. In particular, we assume that 3 is Lipschitz continuous in 12/3,22/3
(ii) The activation function : RR + R is odd, bounddd, Lo.a -iiscchit continuous, and L smooth.
We remark that Assumption 4.2 is not restrictive and can be satisfied by large family of neural networks, e.g., (xx) tanh(x) and ((b) tanh(b). Noting that [|(s, a)ll2 < 1, Assumption 4.2 implies that the function ass, a; f) in (4.2) is odd with respect to w and b and is also bounded, Lipschitz continuous, and smooth in the parameter domain, that is,
Voo a; f) < B1, VV 0;0) < Bg.
(4.4)
We then impose the following assumption on the MDP.

---

Assumption 4.3 (Regularity of the MDP). For the MDP (S, Do), we assume the following properties hold.
(i) The reward function r and the transition kernel admit the following representations with respect to the activation function o,
r(s,0) = B, 1)
(4.5)
P(s' = ((s,,,,)) = pls
(4.6)
where and +s); are probability measures in (IR for any s' E S. B, is positive scaling parameter, and pps'): S - R+ is nonnegative function.
The reward function r satisfies that r(s, a) 2 0 forany (s, a) E SxA. For the representation of in (4.5) and the representation of the transition kernel P in (4.6), we assume that
( ll Pw,0) < Mμ: J pllodds S Mi
x (v((s)) Pw,0) < Mp, E
f p(s))dss M2,0: M2.p:
where Pmr.o is the marginal distribution of Po with respect to w, i.e., Pw:0 po(db, -), x? is the chi-squared divergence, and Mp. My. Mi.e. M2 are absolute constants.
(iii) We assume that there exists an absolute constant & such that
[[:(s: (6': <9, [|((()) 1 <G, [s;((;;)) 1 ( < <9, [| - (''': "(d':-)ll < <9; Vs,6', E S
where |1:1l--a0) is the weighted homogeneous Sobolev norm.
We remark that by assuming v to probability measure and that ppss) 2 in (4.6), the represen- tation of the transition kernel does not lose generality. Specifically, the function class of (4.6) is the same
p ( ()T)) dw)) (5':--) is signed measure for any s' E s).
(s'
See IC.1 for a detailed proof. Assumption 4.3 generalizes the linear MDP in [14, 32, 67, 68). In contrast, our representation of the reward function and the transition kernel benefits from the universal function approximation theorem and is thus not as restrictive as the original linear MDP assumption. Note that the infinite-widtt neural network has a two-layer structure by (4.2). We establish the following lemma the regularity of the representation of the action value function Q" by such neural network.
Lemma 4.4 (Regularity of Representation of Q*). Suppose that Assumptions 4.2 and 4.3 hold. For any policy π, there exists probability measure P* E (RR for the representation of Q with the following properties.
(i) For function Q a; Px) defined by (3.4) with p Ps and the action value function Q" (s, a) defined by (2.2), we have Q(s, Px) = Q= a) for any (s,a) E S x A.
(ii) By letting B3 2 2(B, + 7(1 -1)1-))-BBB,,,,0 for the neural network defined in (4.2) and ~ (N TD) for the initial distribution, we have W2(Pn, Po) & D for any policy n, where define W2(-.) aW2(+,) as the scaled W2 metric. Here constant D depends the discount factor and the absolute constants La,s, In,s, lg. B,, Mas Ma. M1 M2 defined in Assumptions 4.2 and 4.3.
Proof. See $B.2 for detailed proof.

---

Property (i) of Lemma 4.4 shows that the action value function Q" can be parameterized with the infinite-width two-layer neural network Q:= px) in (3.4). Note that a larger B3 captures a larger function class in (4.3). Without loss of generality, we consider that B3 2 2 + (1-)(1111111 B,,))) holds in the sequel. Hence. by Property (ii), it holds that W2(Pm, pb) S ((1) for any policy и. In particular, it holds by Property (i) of Lemma 4.4 that ||Qr - QTh 1-1 = (:Q pt) P(; 112,))1,,5... and we have the following theorem to characterize such an error with regard to the W2 space.
Theorem 4.5 (Upper Bound of Policy Evaluation Error). Suppose that Assumptions 4.2 and 4.3 hold and Po ~ NN, ID) is the initial distribution. We specify the weighting distribution Th in MF-TD (3.5) as T E emror in Theorem 4.1. Then, it holds that
# where कhth E (S x A) is the evaluation distribution for the policy evaluation
cl W3 (Pts PT.)) + A.:
(1 - (1-v)) - 11 Q"t 2, S
(4.7)
dt
21)
where
A. 22 12 'BB, W2(pe: Po) Wy 2 Pri)
+ aB (4B) max [P Po), (Pt...))) + B, WP
Here B1 and B2 defined in (4.4) of Assumption 4.2, 77 is the relative TD timescale, is the scaling parameter of the neural network, and W2 W22 is the scaled W22 metric. Moreover, constant B depends on the dicount factor 7, the scaling parameter B3 in (4.2), and the absolute constants 13 Br, p: G defined in Assumptions 4.2 and 4.3.
e
Proof. See $B.3 for detailed proof.
Here we give nonrigorous discussion on how upper bound A, in (4.7). If W2 Po) I O()) holds for any E [0, by W2(Pn. Po) I ((1) in Lemma 4.4 and the triangle inequality of W2 distance [59], it follows that W2(pt, < O(1) and A, I 0(2 + a=1). Taking a time average of integration on both sides of (4.7), the policy evaluation error + f 11Q Q" l1 di is then upper bounded by O07''T + a /22 - + +a =)) Inspired by such fact, introduce the following restarting mechanism to ensure W2(Pt, S O(1).
Restarting Mechanism. Let Wo = AD be threshold, where D is the upper bound for (2 Po) by Lemma 4.4, A 2 3 is a constant scaling parameter for the restarting threshold, P: is the distribution of the parameters in the neural network at time t, and Pu is the initial distribution. Whenever we detect that W2(pt, Po) reaches Wo in the update, we pause and reset Pt to Pa by resampling the parameters from Po. Then, we reset the critic with the newly sampled parameters while keeping the actor's 's policy Te unchanged and continue the update.
The restarting mechanism guarantees (P),:.Po) I AD by restricting the distribution Ps of the parameters to be close to Po- Moreover, by letting 2 3, we ensure that Pr, realizable by p since W2 Po) < D < XD, which means that the neural network is capable of capturing the representation of the action value function Q We remark that by letting Wo O(1), we allow pt to deviate from Po up W2(Pe, Po) I O1)-1) )) in the restarting mechanism. In contrast, the NTK regime| [15] which corresponds to letting
VM in (3.1) only allows pt deviate from po by the chi-squared divergence (pt Po) < O(M-1)= = o(1). That is, the NTK regime fails to induce feature representation significantly different from the initial one. Before moving on, we summarize the construction of the weighting distribution th Theorem 4.1 and 4.5 as follows,
1 1~
1
Φu / To( da),
TT
0
+ 00 @ Te,
(4.8)
2
2
where do is the base distribution. Now we have the following theorem that characterizes the global optimality and convergence of the two-timescale AC with restarting mechanism.

---

Theorem 4.6 (Global Optimality and Convergence Rate of Two-timescale AC with Restarting Mechanism). Suppose that (4.8) and Assumptions 4.2 and 4.3 hold. With the restarting mechanism, it holds that
L' (((+))((m)))) d(t))) v C + - a
aala 11 a1 S2 11 11 2T (5)
(4.9)
(a)
(b)
where we have
Ex
E 5 (TT'ss)))))) (.oo)//Tooo1 (o(.)()))))))
Do
S = (1 D B,) S2 1 VY
2BBI A(1 + A)D2
V7
Here B,, B1 and B2 are defined in Assumption 4.2 and 4.3, D is the upper bound for W2(P7. Po) in Lemma 4.4, B depends on the discount factor 7 and the absolute constants defined in Assumption 4.2 and 4.3, and is the scaling parameter for the restarting threshold. Besides, it holds for the total restarting number that
N N ( 2)- + V2) +
Proof. See $B.4 for detailed proof.
Note that for given MDP with starting distribution Do. the expected KL-divergence Ç and the concentrability coefficient x are both independent of the wo-timessale update. We remark that our condition for (4.9) to be bounded not restrictive. Specifically, we only need a given To and Φu such that the KL--divergence Ç and the concentrability coefficient co, which is weaker than the concentrability coefficient used in [7, 23, 24, 38, 39, 43, 48, 57, 61].
The first term (a) on the right-hand side of (4.9) diminishes as - 0o. The second term (b) corresponds to the policy evaluation error. We give an example to demonstrate the convergence of the two-time AC. We let the scaling parameter = for the restarting threshold. By letting 1/2) as a - co. Thus, we have that + $ (J(T") - J(T))) dt
3/2 it holds that (b) = -1a-1/2) descends at rate of O O(a-0aaaa0 (1/2) critic has a larger relative TD timescale in (3.5). As for the total number of restartings N. holds that N < (2T as cc, which induces tradeoff, i.e., larger a guarantees a smaller gap in dt but yields more restartings and larger relative TD timescale.
(//a-3/AT-1/3/ 1/3).. Note that 7) 3/2 shows that the
5 Acknowledgement
Zhaoran Wang acknowledges National Science Foundation (Awards 2048075, 2008827, 2015568, 1934931), Simons Institute (Theory of Reinforcement Learning), Amazon, J.P. Morgan, and Two Sigma for their supports. Zhuoran Yang acknowledges Simons Institute (Thcory of Reinforcement Learning). This work was also supported in part by the Vannevar Bush Faculty Fellowship program under grant number 11001421-1--941..
References
[1] Agazzi, A. and Lu, (2019). Temporal-difference learning for nonlinear value function approximation in the lazy training regime. arXiv preprint 1091
[2] Agazzi, A. and Lu, J. (2020). Global optimality of softmax policy gradient with single hidden layer neural networks in the mean-field regime. arXiv preprint arXiv:2010.11858.
10

---

[3] Agostinelli, F., McAleer, S., Shmakov, A. and Baldi, P. (2019). Solving the rubik's 's cube with deep reinforcement learning and search. Nature Machine Intelligence, 356-363.
[4] Akkaya. I., Andrychowicz, M., Chociej, M., Litwin, M., McGrew, B., Petron, A., Paino, A., Plappert, M., Powell, G., Ribas, R. et al. (2019). Solving rubik's 's cube with robot hand. arXiv preprint arXiv:1910.07113.
[5] Ambrosio, L. and Gigli, N. (2013). A user's guide to optimal transport. In Modelling and Optimisation of Flows Networks. Springer, 1-155,
[6] Ambrosio, L., Gigli, N. and Savaré, G. (2008). Gradient flows: In metric spaces and in the space of probability measures. Springer.
[7] Antos, A., Szepesvári, and Munos, R. (2008). Learning near-optimal policies with Bellman- residual minimization based fitted policy iteration and single sample path. Machine Learning,
71 89-129.
[8] Barron, (1993). Universal approximation bounds for superpositions of sigmoidal function. IEEE Transactions on Information Theory, 39 930-945.
[9] Berner, C., Brockman, G., Chan, B., Cheung, V.. Debiak, P., Dennison, C., Farhi, D.,
Fischer, Q., Hashme, S., Hesse, et (2019). Dota with large scale deep reinforcement learning. arXiv preprint 1912v:1912.06660
[10] Bhatnagar, S., Ghavamzadeh, M., Lee, M. and Sutton, R. S. (2008). Incremental natural actor- critic algorithms. In Advances Neural Information Processing Systems.
[11] Bhatnagar, S., Sutton, R. S., Ghavamzadeh, M. and Lee, M. (2009). Natural actor-critic algo- rithms. Automatica, 45 2471-2482
[12] Börgers, T. and Sarin, R. (1997). Learning through reinforcement and replicator dynamics. Journal of economic theory, 77 1-14.
[13] Borkar, V. S. (2009). Stochastic approximation: A dynamical systems viewpoint. Springer.
[14] Cai, Q., Yang, Z., Jin, C. and Wang, Z. (2019). Provably efficient exploration in policy opti- mization. arXiv preprint
[15] Cai, Q., Yang, Z., Lee, D. and Wang, Z. (2019). Neural temporal-difference learning con- verges global optima. In Advances Neural Information Processing Systems.
[16] Chen, M., Wang, Y., Liu, T., Yang, Z., Li, X., Wang, Z. and Zhao, (2020). On com- putation and generalization of generative adversarial imitation learning. arXiv preprint wwXiv:2001.0279
[17] Chen, Z., Cao, Y., Gu, Q. and Zhang, T. (2020). Mean-field analysis of two-layer neural networks: Non-asymptotic rates and generalization bounds. arXiv preprint arXiv:2002.044 04026.
[18] Chizat, L. and Bach, F. (2018), A note on lazy training in supervised differentiable program- ming. arXiv preprint 1rXiv:1812.07956
[19] Chizat, L. and Bach, F. (2018). On the global convergence of gradient descent for over- parameterized models using optimal transport. In Advances in Neural Information Processing Systems.
[20] Duan, Y., Chen, X., Houthooft, R., Schulman, J. and Abbeel, P. (2016), Benchmarking deep reinforcement learning for continuous control. In International Conference on Machine Learning.
[21] Fang, C., Dong, H. and Zhang, T. (2019). Over parameterized two-level neural networks can learn optimal feature representations. arXiv preprint arXiv:1910.11508.

---

[22] Fang, C., Gu, Y., Zhang, W. and Zhang, T. (2019). Convex formulation of overparameterizee deep neural networks. arXiv preprint arXiv:1911.0762)
[23] Farahmand, A.- Ghavamzadeh, M., Szepesvári, C. and Mannor, S. (2016). Regularized policy iteration with nonparametrie function spaces. The Journal of Machine Learning Research, 17 4800--8744
[24] Farahmand, A.-m., Szepesvári, C. and Munos, R. (2010). Error propagation for approximate policy and value iteration. In Advances in Neural Information Processing Systems.
[25] Friedrichs, K. O. (1944). The identity of weak and strong extensions of differential operators. Transactions of the American Mathematical Society, 55 132-151.
[26] Fu, Z., Yang, Z. and Wang. Z. (2020). Single-timescale actor-critic provably finds globally optimal policy. arXiv preprint arXiv:2008.00483
[27] Harker, P. T. and Pang. J.-S. (1990). Finite-dimensional variational inequality and nonlinear complementarity problems: A survey of theory, algorithms and applications. Mathematical Programming, 48 161-220.
(28) Hennes, D., Morrill, D., Omidshafiei, S., Munos, R., Perolat, J., Lanctot, M., Gruslys, A., Lespiau, J.-B., Parmas, P., Duéiiez-Guzmán, E. et al. (2020). Neural replicator dynamics: Mul- tiagent learning via hedging policy gradients. In International Conference on Autonomous Agents and MulriAgent Systems.
[29] Hong, M., Wai, H.-T., Wang. Z. and Yang, Z. (2020). A two-timescale framework for bilevel optimization: Complexity analysis and application to actor-critic. arXiv preprint arXiv:2007.05170.
[30] Jacot, A., Gabriel, F. and Hongler, C. (2018). Neural tangent kernel: Convergence and general- ization in neural networks. In Advances in Neural Information Processing Systems.
[31] Javanmard, A., Mondelli, M. and Montanari, A. (2019), Analysis of two-layer neural network via displacement convexity. arXiv preprint arXiv:1901.01375
[32] Jin, C., Yang, Z., Wang, Z. and Jordan, M. I. (2020). Provably efficient reinforcement learning with linear function approximation. In Conference on Learning Theory.
[33] Kakade, S. and Langford, J. (2002). Approximately optimal approximate reinforcement learn- ing. In International Conference on Machine Learning, vol. 2.
[34] Kakade, S. M. (2002). A natural policy gradient. In Advances Neural Information Processing Systems.
[35] Khodadadian, S., Doan, T. T., Maguluri, S. T. and Romberg. J. (2021). Finite sample analysis of two-time-scale natural actor-critic algorithm. arXiv preprint arXiv:2101.10506.
[36] Konda, V. R. and Tsitsiklis, J. N. (2000). Actor-critic algorithms. In Advances in Neural Information Processing Systems.
[37] Kushner, H. and Yin, G. G. (2003). Stochastic approximation and recursive algorithms and applications. Springer.
[38] Lazaric, A., Ghavamzadeh, M. and Munos, R. (2016). Analysis of classification-basse policy iteration algorithms. The Journal of Machine Learning Research, 17 583-612.
[39] Liu, B., Cai, Q., Yang, Z. and Wang, Z. (2019). Neural proximal/tusst region policy optimiza- tion attains globally optimal policy. arXiv preprint arXiv:1906.10300
12

---

[40] Lu, Y., Ma, C., Lu, Y., Lu, J. and Ying, L. (2020). A mean field analysis of deep resnet and beyond: Towards provably optimization via overparameterization from depth. In International Conference on Machine Learning. PMLR.
[41] Mei, S., Misiakiewicz, and Montanari, A. (2019). Mean-field theory of two-layers neural networks: dimension-free bounds and kernel limit. In Conference on Learning Theory. PMLR.
[42] Mei, S., Montanari, A. and Nguyen, P.-M. (2018). A mean field view of the landscape of two- layer neural networks. Proceedings of the National Academy of Sciences, 115 E7665-E7671.
[43] Munos, R. and Szepesvári, C. (2008). Finite-time bounds for fitted value iteration. The Journal of Machine Learning Research, 815-857.
[44] Otto, F. and Villani, C. (2000). Generalization of an inequality by Talagrand and links with the logarithmic Sobolev inequality. Journal of Functional Analysis, 173 361-4664
[45] Peters, J. and Schaal, S. (2008). Natural actor-critic. Neurocomputing, 71 1180-1 190.
[46] Peyre, R. (2011). Comparison between W2 distance and H-1 norm, and localisation of wasserstein distance. arXiv preprint arXiv:/104.460
[47] Pinkus, A. (1999). Approximation theory of the MLP model in neural networks. Acta Numerica, 8 143-195.
[48] Scherrer, B., Ghavamzadeh, M., Gabillon, V., Lesner, B. and Geist, M. (2015). Approximate modified policy iteration and its application to the game of Tetris. The Journal of Machine Learning Research, 16 1629-1676.
[49] Schulman, J., Wolski, F., Dhariwal, P., Radford, A. and Klimov, O. (2017). Proximal policy optimization algorithms. arXiv preprint arXiv: 1707.06777
[50] Schuster, P. and Sigmund, K. (1983). Replicator dynamics. Journal of theoretical biology, 100 533-538.
[51] Silver, D., Huang, A., Maddison, C. Guez, A., Sifre, L., Van Den Driessche, G., Schrittwieser, J., Antonoglou, I., Panneershelvam, V., Lanctot, M. et al. (2016). Master- ing the game of Go with deep neural networks and tree search. Nature, 529 484 489.
[52] Silver, D., Schrittwieser, J., Simonyan, K., Antonoglou, I., Huang. A.. Guez, A., Hubert, T., Baker, I.,, Lai, M., Bolton, A. al. (2017). Mastering the game of Go without human knowl- edge. Nature, 550 354.
[53] Sirignano, J. and Spiliopoulos, K. (2020). Mean field analysis of neural networks: A central limit theorem. Stochastic Processes and their Applications, 130 1820-1852.
[54] Sirignamo, and Spiliopoulos, K. (2020). Mean field analysis neural networks: A law of large numbers. SIAM Journal on Applied Mathematics, 80 725-752.
[55] Sutton, R. S. (1988). Learning to predict by the methods of temporal differences. Machine Learning, 9-44.
[56] Sutton, R. S. and Barto, A. G. (2018). Reinforcement learning: An introduction. MIT press.
[57] Szepesvári, C. and Munos, R. (2005). Finite time bounds for sampling based fitted value iteration. International Conference on Machine Learning. ACM.
[58] Villani, C. (2003). Topics in optimal transportation. American Mathematical Society.
[59] Villani, C. (2008). Optimal transport: Old and new. Springer.
13

---

[60] Vinyals, O., Babuschkin, I., Czarnekki, W. M., Mathieu, M., Dudzik, A., Chung, J., Choi, D. H., Powell, R., Ewalds, T., Georgiev, P. et al. (2019). Grandmaster level in StarCraft II using multi-agent reinforcement learning. Nature, 575 350-354.
[61] Wang, L.,, Cai, Q. Yang, Z. and Wang, Z. (2019). Neural policy gradient methods: Global optimality and rates of convergence. arXiv preprini 190X.019090011550
[62] Watkins, C. and Dayan, P. (1992). Q-learning. Machine Learning. 8 279-292..
[63] Wei, C., Lee, J. D., Liu, Q. and Ma, T. (2019). Regularization matters: Generalization and optimization of neural nets vs their induced kernel. In Advances in Neural Information Processing Systems.
[64] Wu, Y., Zhang, W., Xu, P. and Gu, Q. (2020). A finite time analysis of two time-scale actor critic methods. arXiv preprint arXiv:2005.01350.
[65] Xu, T., Wang, and Liang, Y. (2020). Improving sample complexity bounds for actor-critic algorithms. arXiv preprint urXiv:2004.12555
[66] Xu, T., Wang, Z. and Liang, Y. (2020). Non-asymptotic convergence analysis of two time-scale (natural) actor-critic algorithms. arXiv preprint arXiv:2005.03557.
[67] Yang, L. F. and Wang, M. (2019). Reinforcement learning in feature space: Matrix bandit, kernels, and regret bound. arXiv preprint
[68] Yang, L. and Wang, M. (2019). Sample-optimal parametric q-learning using linearly additive features. arXiv preprint ar (900:0472707777
[69] Zhang, Y., Cai, Q., Yang, Z., Chen, Y. and Wang. Z. (2020). Can temporal-difference and q-learning learn representation? a mean- field theory. arXiv preprint arXiv:2006.04761
14


Reviews:
[{"review_id": "wQkNTIj1xpl", "reviewer": "Reviewer_ozyQ", "summary": "The paper studies convergence of an Actor-Critic algorithm for an MDP problem with infinite one-hidden layer neural network approximation of Q function in the regime mean-field  (Theorem 4.5) and lazy type regime (Theorem 4.6).    By bringing (heavy) machinery from gradient flows literature, the authors perform a similar analysis to the one performed by a group of A. Montanari or Chizat/Bach in the context of supervised learning problem with one-hidden layer approximation which more recently has been adapted by Agazzi an Lu to study popular Policy Gradient algorithm. \n", "questions": "", "limitations": "", "rating": 6, "confidence": 5, "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": "Originality:  The results seem to be new, and the main contribution is the study of actor-critic algorithm represented by gradient flow equations (3.5) MF-TD and (3.7) MF-PPO. However, the dynamics of 3.7 is very simple and one can see from the error in Theorem 3.7 that the crux of the matter is the estimate on Q functions (where grad flow  3.7 plays a small role).  At that point, it would be very helpful if authors presented a thorough comparison to the work [69] of Zhang et al. and possibly [1] Agazzi et al. Especially result from [69] seems very relevant (see their Section 4), though I only had a chance to skim through that other paper.   \n\nQuality: This is theoretical work and as already pointed the results stated in section 4 follow from a heavy analysis of gradient flow. In my opinion, authors did not make sufficient effort to help readers follow the proofs. Even though I'm rather well familiar with the analysis of this type I found reading the paper difficult.   \n\nThere are a number of things that would require further comments. Let me give a few examples (impossible to list them all here):\n- Heuristics in section 3 are helpful, but there isn't any comments on which of these approximation errors would be easy to quantify and which not.\n- in 3.3, entropy regularisation appears without any comments.  This doesn't correspond to classical an actor-critic algorithm. Is this needed? How does this help?\n-  [top page 7] when giving a reference to technical results, please specify chapter, section, concrete results. Giving a reference to a whole book is useless. \n- I have some issues with assumptions 4.2: why the activation function need to be odd? And Lipschtnes in [-2/3,2/3] seems too prescriptive.  \n- In 4.3: why it is sensible to assume that both reward and transition functions need to admit one-hidden layer representation and with the same activation function? This seems cooked up and little effort has been made to explain why these are sensible assumptions. \n- Theorem 4.5 and a discussion that follows: One would expect the bound to decay with t, but instead authors stated results that says that one need to send \\eta and \\alpha  to \\infty. Now \\eta in 3.5 corresponds to the time change and \\alpha dictates scaling of the neural network approximation. Hence, if I'm reading this correctly, this means that in effect, authors study lazy regime in the end... This is somehow also implied by 'restoring mechanism.'\n- Authors mention replicator dynamics and even put this in the title, but there isn't anything in the paper saying why this is relevant.\n\nClarity: As already stated, I would appreciate more discussion around assumptions and main results and how they relate to what's in the literature. The analysis in the appendix is hard to read, and significantly more effort is needed to make it accessible to NIPS community. \n\nSignificance: While I appreciate the importance of this type of work to shed more light on why neural networks based algorithms work, I don't feel this work adds significantly in comparison to what already has been done in the study of on-hidden layer neural networks in mean-field or lazy regimes.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "n/a", "review_text": "", "extra_scores": {"ethical_concerns": "n/a", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3h"}, "extra_text": {}}, {"review_id": "iUUe2uwVnb", "reviewer": "Reviewer_R6mo", "summary": "This paper studies actor-critic (AC) algorithms with neural networks from a theoretical perspective. \n\nThe authors consider an AC algorithm variant that updates its actor with PPO with a small stepsize, and its critic updated via TD and a larger stepsize. Under certain conditions, the authors prove that this AC reaches the optimal policy with sublinear rate. The authors additionally show that the feature representations learned by the critic can move further away from the initialization. ", "questions": "", "limitations": "", "rating": 7, "confidence": 2, "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": "- Originality: authors approached the neural AC problem from very interesting angles and the approach as well as the results seem to be novel\n- Quality: The theoretical results in the paper seem interesting and important\n- Clarity: the paper is clearly written\n- Significance: significant results that improves understanding of how actor-critic methods learn with neural networks\n\nQuestions/comments for the authors:\n- Why would we want to update the critic much faster than the actor? In many recent paper on deep reinforcement learning (for example, TD3, SAC), the actor and critic have similar or same learning rates. And if the two are updated at the same rate, how would that affect the results proven in the paper? \n- The results on representation can evolve into a larger neighborhood seems to be very interesting, but why is this significant? What does it tell us about how actual neural networks should be treated differently from NTK? In my opinion it would be great if the authors can provide some more discussion here. \n- The paper mentions many times the close relationship to replicator dynamics, but it seems to me the paper does not have a good background introduction to exactly what this relationship is, and this might be a bit confusing for readers who are not that familiar with replicator dynamics. \n\nMinor comments:\n- there is a typo at line 178 \"the the critic\"\n\n==================\n\npost rebuttal: I have read the responses from the authors, I feel my questions are addressed and my evaluation remains the same", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Since the authors are analyzing neural AC algorithms, it might be really interesting to have some numerical experiments, especially on the relative learning rate of actor and critic. And it would be interesting to have analysis on how this relates to other problems in RL such as bias accumulation. However, the theoretical results in this paper seem to be already quite important. ", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3"}, "extra_text": {}}, {"review_id": "WqpDEncuUQw", "reviewer": "Reviewer_5c4N", "summary": "The authors analyse a mean-field limit of an actor-critic algorithm for which the critic is parametrized with a single-hidden-layer neural network, and the policy is tabular. The analysis is carried out in continuous time, and in the limit of infinite network width.", "questions": "", "limitations": "", "rating": 6, "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": "The analysis of reinforcement learning algorithms under non-linear function approximation is an important direction for better understanding the behavior of deep reinforcement learning agents. As such, the topic the authors study is likely to be of interest to the (theoretical side of) the NeurIPS RL community.\n\nThe mean-field limit considered here is interesting and elegant, and broadly the mathematical presentation is precise. A wide variety of related work is presented. However, at the moment I think the paper is currently lacking a broader discussion of the context of these results. This is the primary reason for the current rating for the paper, and I would be open to revising the authors if the authors can address these concerns.\n\nAt present, a lot of space is given to detailing aspects of the convergence proof, at the expense of e.g. a conclusion at the end of the main paper, and more discussion/motivation of the various limits described in Section 3, and what these results have to say about Q-learning with neural networks. In particular, I was left with several questions about the motivation and utility of the main results:\n * To obtain the mean-field limit Eqns (3.5) and (3.7), several limits are taken (infinite width, and stepsizes/continuous-time updates). I was unsure whether the order in which these limits are taken affects the analysis, and how closely we should expect the analysis of these systems to relate to finite-width networks updated in discrete-time, and what the analysis can tell us about these systems. Some additional discussion of these connections would be useful, and small-scale experiments to understand when the limit is likely to \"kick in\" may also be useful.\n * There are some seemingly strong assumptions made in the course of the analysis, such as a tabular policy, bounded output weights for the neural network in Eqn (4.2), structure assumptions on the MDP (Assumption 4.3), and an algorithmic restarting mechanism, presumably to avoid divergence. I would expect these to limit the applicability of the main theoretical result, but there is not much discussion of the practical implication of these assumptions in the paper.\n\nMinor technical comments\n\nWhy can we assume (s, a) has uniformly bounded norm without loss of generality? It seems this induces some topological assumptions on S x A (precompactness, for example) - can the authors make clear why their arguments would still apply to state-action spaces which aren't precompact? Is the argument at Equation (4.4) affected by this?\n\nUnder the general assumptions on the structure of state-action space, I don't think it holds that Q^pi is the unique minimizer of MSBE - some kind of qualification allowing  \\tilde{\\Phi}-almost sure equality is required (i.e. consider modifying Q^pi at a single state-action pair that has zero probability under \\tilde\\{Phi}^pi).\n\nLine 245: Is there a typo here? Should (-B_beta, B_beta) be [-B_beta, B_beta]?\n\nLine 248: \"Such over-representation appears to be essential for our analysis\". Can more be said about this?\n\n---\n\nAFTER REBUTTAL:\n\nI thank the authors for their clarifications on many of the questions I mentioned in the original review, and have increased my rating for the paper. There are several points of discussion (for example, the boundedness of state-action features) here that would benefit the paper to be included in later versions.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "As mentioned above, more discussions of the assumptions made for the main theoretical results (including their limitations) would be useful.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "5"}, "extra_text": {}}, {"review_id": "Dux-9HWx2US", "reviewer": "Reviewer_yWAT", "summary": "This paper analyzes the optimality and convergence of the two-scale update for an actor-critic algorithm where the actor and critic are parameterized by two layer NNs (which is a generalization of traditional linear function approximation guarantees in RL). The authors analyze for the PPO as actor and minimizer of the the MSBE as the critic. Considering the actor-critic updates as gradient flow differential equations in the L_2 Wasserstein space, in following steps, the authors show: (1) optimality and convergence rate of policy error proportional to inverse of time and the critic error (2) bounding the critic error based on the Wasserstein error of the distribution of NN weights (3) applying (1) and (2) with Restarting mechanism (resampling the critic weights from an initial distribution when they are sufficiently close to the initial distribution), they prove the global optimality and convergence rate of AC method .", "questions": "", "limitations": "", "rating": 6, "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": "The paper is clear and enjoyable to read. The idea of the paper is interesting and the convergence analysis of an actor critic algorithm with two layer neural networks potentially helps strengthening RL theoretical foundation.\n\nMy main concern about this paper is the update 3.2 that is a based on minimizing MSBE which suffers from double sampling and I couldn’t find it being incorporated/or discussed in the analysis of the paper. I am interested in knowing if and how much this affects the convergence analysis in this paper?\n\nLine 276: does not lose…", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "They discussed potential analytical limitations in the conclusion in the appendix", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4"}, "extra_text": {}}]

Author rebuttal:

Meta-review: All reviewers thought this paper made useful contributions around strengthening the theoretical foundations of RL through an analysis of an AC algorithm with neural networks. AC is a very popular approach in RL and this will likely be of interest to the large body of researchers interested in RL. The authors are encouraged to address the reviewers' feedback in their camera ready paper. 
Decision: Accept (Poster)