Wasserstein Flow Meets Replicator Dynamics: A Mean-Field Analysis of Representation Learning in Actor-Critic

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.

text: 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.

heading: 1 Introduction
text: In reinforcement learning (RL) [56], an agent aims to learn the optimal policy that maximizes the expected total reward by interacting with the environment. Policy-based RL algorithms achieve such a goal by directly optimizing the expected total reward as a 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 a better policy, which is known as policy improvement. Algorithms with these two ingredients are also called actor-critic (AC) methods [36], where the actor and the critic refer to the policy and its corresponding value function, respectively.
† Equal contribution.
35th Conference on 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, 60]. 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, a 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 a 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 temporal-difference (TD) learning [55] for policy evaluation and the actor improves its policy via proximal policy optimization (PPO) [49], which corresponds to a 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 a 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 a regime, the actor update is closely connected to replicator dynamics [12, 28, 50] and the critic update is captured by a semigradient flow in the Wasserstein space [59]. Moreover, the semigradient flow runs at a 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 a sublinear rate. Furthermore, regarding representation learning, we show that the critic induces a data-dependent feature representation within an O(1/α) neighborhood of the initial representation in terms of the Wasserstein distance, where α is a sufficiently large scaling parameter.
The key to our technical analysis reposes on three ingredients: (i) infinite-dimensional variational inequalities with a one-point monotonicity [27], (ii) a mean-field perspective on neural networks [19, 41, 42, 53, 54], and (iii) the two-timescale stochastic approximation [13, 37]. In particular, in the infinite-width limit, the neural network and its induced feature representation are identified with a distribution over the parameter space. The mean-field perspective enables us to characterize the evolution of such a distribution within the Wasserstein space via a partial differential equation (PDE) [5, 6, 58, 59]. For policy evaluation, such a PDE is given by the semigradient flow induced by TD learning. We characterize the error of policy evaluation by showing that mean-field Bellman error satisfies a 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 a 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 a two-timescale stochastic approximation via ordinary differential equations, [10, 11, 36] establish asymptotic convergence guarantees in the continuous-time limiting regime. 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 a 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 a version of neural AC is well approximated by its linear counterpart via the neural tangent feature. In comparison, we establish a mean-field analysis that has a 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 overparameterized 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 a semigradient algorithm and there is no objective functional making TD learning a gradient-based algorithm. Thus, in the mean-field regime, our critic is given by a 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 a 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 P(X ) the set of probability measures over the measurable space X . Given a curve ρ : R → X , we denote by ρ̇s = ∂tρt | t=s its derivative with respect to the time. For an operator F : X → X and a measure µ ∈ P(X ), we denote by F♯µ = µ◦F−1 the push forward of µ through F . We denote by χ2(ρ ∥µ) the chi-squared divergence between probability measures ρ and µ, which is defined as χ2(ρ ∥µ) = ∫ (ρ/µ− 1)2dµ. Given two probability measures ρ and µ, we denote
the Kullback-Leibler divergence or the relative entropy from µ to ρ by KL(ρ ∥µ) = ∫ log(ρ/µ)dρ. For ν1, ν2, µ ∈ P(X ), we define the Ḣ−1(µ) weighted homogeneous Sobolev norm as ∥ν1 − ν2∥Ḣ−1(µ) = sup { |⟨f, ν1 − ν2⟩|
∣∣∥f∥Ḣ1(µ) ≤ 1}. We denote by ∥f(x)∥p,µ = (∫ |f(x)|pµ(dx))1/p the ℓp-norm with respect to probability measure µ. We denote by ⊗ the semidirect product, i.e., µ ⊗ K = K(y |x)µ(x) for µ ∈ P(X ) and transition kernel K : X → P(Y). For a function f : X → R, we denote by Lip(f) = supx,y∈X ,x ̸=y |f(x)−f(y)|/∥x−y∥ its Lipschitz constant. We denote a normal distribution on RD by N (µ,Σ), where µ is the mean value and Σ is the covariance matrix.

heading: 2 Background
text: In this section, we first introduce the policy optimization problem and the actor-critic method. We then present the definition of the Wasserstein space.

heading: 2.1 Policy Optimization and Actor-Critic
text: We consider a Markov decision process (MDP) given by (S,A, γ, P, r,D0), where S ⊆ Rd1 is the state, A ⊆ Rd2 is the action space, γ ∈ (0, 1) is the discount factor, P : S × A → P(S) is the transition kernel, r : S × A → R+ is the reward function, and D0 ∈ P(S) is the initial state distribution. Without loss of generality, we assume that S × A ⊆ Rd and ∥(s, a)∥2 ≤ 1, where d = 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 π : S × A → P(S), at the mth step, the agent takes an action am at state sm according to π(· | sm) and observes a reward rm = r(sm, am). The environment then transits to the next state sm+1 according to the transition kernel P (· | sm, am). Note that the policy π induces Markov chains on both S and S ×A. Considering the Markov chain on S, we denote the induced Markov transition kernel by Pπ : S → P(S), which is defined as Pπ(s′ | s) = ∫ A P (s
′ | s, a)π(da | s). Likewise, we denote the Markov transition kernel on S × A by P̃π : S × A → P(S × A), which is defined as P̃π(s′, a′ | s, a) = π(a′ | s′)P (s′ | s, a). Let D̃ be a probability measure on S ×A. We then define the visitation measure induced by policy π and starting from D̃ as
ẼπD̃ ( d(s, a) ) = (1− γ) · ∑ m≥0 γm · P ( (sm, am) ∈ d(s, a) | (s0, a0) ∼ D̃ ) , (2.1)
where (sm, am) is the trajectory generated by starting from (s0, a0) ∼ D̃ and following policy π thereafter. If D̃ = D ⊗ π holds, we then denote such a visitation measure by ẼπD. Furthermore, we denote by E(ds) = ∫ A Ẽ(ds,da) the marginal distribution of visitation measure Ẽ with respect to S. In particular, when (s0, a0) ∼ D ⊗ π holds in (2.1), it follows that ẼπD = EπD ⊗ π. In policy optimization, we aim to maximize the expected total rewards J(π) defined as follows,
J(π) = Eπ [∑ m≥0 γm · r(sm, am) ∣∣ s0 ∼ D0],
where we denote by Eπ the expectation with respect to am ∼ π(· | sm) and sm+1 ∼ P (· | sm, am) for m ≥ 0. We define the action value function Qπ : S × A → R and the state value function V π : S → R as follows,
Qπ(s, a) = Eπ [∑ m≥0 γm · r(sm, am) ∣∣ s0 = s, a0 = a], V π(s) = 〈Qπ(s, ·), π(· | s)〉A, (2.2) where we denote by ⟨·, ·⟩A the inner product on the action space A. Correspondingly, the advantage function Aπ : S ×A → R is defined as
Aπ(s, a) = Qπ(s, a)− V π(s). It is known that the action value function Qπ is the unique global minimizer to the following mean-squared Bellman error (MSBE),
MSBE(Q;π) = 1
2 E(s,a)∼Φ̃π
[( Q(s, a)− r(s, a)− γE(s′,a′)∼P̃π(· | s,a)[Q(s ′, a′)] )2] , (2.3)
where Φ̃π is a weighting distribution depending on policy π 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 π
J(π) = Es∼D0 [〈 Qπ(s, ·), π(· | s) 〉 A ] , subject to Qπ = argmin
Q MSBE(Q;π). (2.4)
The inner problem in (2.4) is known as a 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 evaluation.

heading: 2.2 Wasserstein Space
text: Let Θ ⊆ RD be a Polish space. We denote by P2(Θ) ⊆ P(Θ) the set of probability measures with finite second moments. Then, the Wasserstein-2 (W2) distance between µ, ν ∈ P2(Θ) is defined as follows,
W2(µ, ν) = inf { E [ ∥X − Y ∥2 ]1/2 ∣∣∣ law(X) = µ, law(Y ) = ν}, where the infimum is taken over the random variables X and Y on Θ and we denote by law(X) the distribution of a random variable X . We call M = (P2(Θ),W2) the Wasserstein (W2) space, which is an infinite-dimensional manifold [59]. See §A.1 for more details.

heading: 3 Algorithm
text: Two-timescale Actor-critic. We consider a 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 a 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 θ̂ = (θ̂(1), · · · , θ̂(M)) ∈ RD×M ,
Qθ̂(s, a) = α
M M∑ i=1 σ(s, a; θ̂(i)). (3.1)
Here σ(s, a; θ) : S × A × RD → R is the activation function and α > 0 is the scaling parameter. Such a structure also appears in [17, 18, 41]. In a descrete-time finite-width (DF) scenario, at the kth iteration, the critic and actor are updated as follows,
DF-TD: θ̂(i)k+1 = θ̂ (i) k −
ε′
α
( Qθ̂k(sk, ak)− r(sk, ak)− γQθ̂k(s ′ k, a ′ k) ) ∇θσ(s, a; θ̂(i)k ), (3.2)
DF-PPO: π̂k+1(· | s) = argmax π
{〈 Qθ̂k(s, ·), π(· | s) 〉 A − ε −1 ·KL ( π(· | s) ∥ π̂k(· | s) )} , (3.3)
where (sk, ak) ∼ Φ̃π̂k and (s′k, a′k) ∼ P̃ π̂k(· | sk, ak). Here π̂k is the policy for the actor at the kth iteration, Φ̃π̂k is the corresponding weighting distribution, ε and ε′ are the stepsizes for the DF-PPO update and the DF-TD update, respectively. In (3.2), the scaling of α−1 arises since our update falls into the lazy-training regime [18]. In the sequel, we denote by η = ε′/ε the relative TD timescale. Note that in a 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θ̂k for updating the actor’s policy π̂k can still be far from the true action value function Qπ̂k . 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 to establish a theoretical guarantee.
Mean-field (MF) Limit. To analyze the convergence of the two-timescale AC with neural networks, we employ the analysis that studies the mean-field limit regime [41, 42]. Here, by saying the meanfield 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., t = kε where ε→ 0 for the stepsize in (3.2) and (3.3). For θ̂ = {θ̂(i)}Mi=1 independently sampled from a distribution ρ, we can write the infinite-width limit of (3.1) as
Q(s, a; ρ) = α ∫ σ(s, a; θ)ρ(dθ). (3.4)
In the sequel, we denote by ρ̂k the distribution of θ̂ (i) k for the infinite-width limit of the neural network at the kth iteration. We further let ρt and πt be the continuous-time limits of ρ̂k and π̂k, respectively.
As studied in [69], the mean-field limit of the DF-TD update in (3.2) is MF-TD: ∂tρt = −η div ( ρt · g(·; ρt, πt) ) , (3.5)
where η is the relative TD timescale and
g(θ; ρ, π) = −Eπ Φ̃π
{[ Q(s, a; ρ)− r(s, a)− γ ·Q(s′, a′; ρ) ] · α−1∇θσ(s, a; θ) } (3.6)
is a vector field. Here Eπ Φ̃π is taken with respect to (s, a) ∼ Φ̃π and (s′, a′) ∼ P̃π(· | s, 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 DF-PPO update can be written in closed form as
ε−1 · { log [ π̂k+1(a | s) ] − log [ π̂k(a | s) ]} = Q(s, a; ρ̂k)− Ẑk(s),
where Ẑk(s) is the normalizing factor such that ∫ π̂k(da | s) = 1 for any s ∈ S. By letting t = kε and ε → 0, we have ∂t log πt = Qt − Zt, which can be further written as ∂tπt = πt · (Qt − Zt). Here we have Qt(a, s) = Q(a, s; ρt) and Zt is the continuous-time limit of Ẑk. Furthermore, noting that ∂t ∫ πt(da | s) = 0, the mean-field limit of the DF-PPO update in (3.3) is
MF-PPO: d
dt πt = πt ·At, where At(s, a) = Qt(s, a)−
∫ Qt(s, a)πt(da | s). (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, we remark that the MF-TD update in (3.5) for the critic is captured by a 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.

heading: 4 Main Result
text: 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.

heading: 4.1 Convergence of Mean-field PPO
text: 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 π∗ = argmaxπ J(π) be the optimal policy and π0 be the initial policy. Then, it holds that
1
T ∫ T 0 ( J(π∗)− J(πt) ) dt ≤ ζ T + 4κ · 1 T ∫ T 0
∥Qt −Qπt∥2,ϕ̃πtdt︸ ︷︷ ︸ policy evaluation error , (4.1)
where ϕ̃πt ∈ P(S × A) is an evaluation distribution for the policy evaluation error and ζ = Es∼Eπ∗D0
[ KL ( π∗(· | s) ∥π0(· | s) )] is the expected KL-divergence between π∗ and π0. Furthermore,
letting ϕ̃πt = 12 ϕ̃0 + 1 2ϕ0 ⊗ πt, where ϕ̃0 ∈ P(S ×A) is a base distribution and ϕ0 = ∫ A ϕ̃0(·,da), the concentrability coefficient κ is then given by
κ = ∥∥∥∥ Ẽπ∗D0 ϕ̃0 ∥∥∥∥ ∞ .
Proof. See §B.1 for a 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 ϕ̃0 such that κ <∞.
Theorem 4.1 shows that the MF-PPO converges to the globally optimal policy at a rate of O(T−1) up to the policy evaluation error. Such a theorem implies the global optimality and convergence of a double-loop AC algorithm, where the critic Qt 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 policy πt.

heading: 4.2 Global Optimality and Convergence of Two-timescale AC
text: 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 Qt = Q(·; ρt) is updated via the MF-TD in (3.5). For the smooth function σ in the parameterization of the Q function in (3.4), we consider it to be the following two-layer neural network,
σ(s, a; θ) = Bβ · β(b) · σ̃ ( w⊤(s, a, 1) ) , (4.2)
where σ̃ : R → R is the activation function, θ = (b, w) is the parameter, and β : R → (−1, 1) is an odd and invertible function with scaling hyper-parameter Bβ > 0. It then holds that D = d+ 2, where d and D are the dimensions of (s, a) and θ, respectively. It is worth noting that the function class of ∫ σ(s, a; θ)ρ(dθ) for ρ ∈ P2(RD) is the same as
F = {∫ β′ · σ̃ ( w⊤(s, a, 1) ) ν(dβ′,dw) ∣∣∣ ν ∈ P2((−Bβ , Bβ)× Rd+1)}, (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 β in (4.2) to avoid the study of the space of probability measures over (−Bβ , Bβ) × Rd+1 in (4.3), which has boundary and thus lacks the regularity in the study of optimal transport. Furthermore, note that we introduce a hyper-parameter α > 1 in the Q function in (3.4). Thus, we are using α · F to represent F , which causes an “overrepresentation” when α > 1. Such over-representation appears to be essential for our analysis. For a brief peek, we remark that α actually controls the gap in the average total reward over time when the relative time-scale η is properly selected according to Theorem 4.6. Furthermore, such an influence is imposed through Lemma 4.4, which shows that the Wasserstein distance between ρ0 and ρπt is upper bounded by O(1/α). In what follows, we consider the initialization of the TD update to be ρ0 = N (0, ID), which implies that Q(s, a; ρ0) = 0. We next impose the following assumption on the two-layer neural network σ.
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 β : R → (−1, 1) is odd, L0,β-Lipschitz continuous, L1,β-smooth, and invertible. Meanwhile, the inverse β−1 is locally Lipschitz continuous. In particular, we assume that β−1 is ℓβ-Lipschitz continuous in [−2/3, 2/3].
(ii) The activation function σ̃ : R → R is odd, Bσ̃-bounded, L0,σ̃-Lipschitz continuous, and L1,σ̃-smooth.
We remark that Assumption 4.2 is not restrictive and can be satisfied by a large family of neural networks, e.g., σ̃(x) = tanh(x) and β(b) = tanh(b). Noting that ∥(s, a)∥2 ≤ 1, Assumption 4.2 implies that the function σ(s, a; θ) 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,
|∇θσ(s, a; θ)| < B1, |∇2θθσ(s, a; θ)| < B2. (4.4)
We then impose the following assumption on the MDP.
Assumption 4.3 (Regularity of the MDP). For the MDP (S,A, γ, P, r,D0), we assume the following properties hold.
(i) The reward function r and the transition kernel P admit the following representations with respect to the activation function σ̃,
r(s, a) = Br · ∫ σ̃ ( (s, a, 1)⊤w ) µ(dw), (4.5)
P (s′ | s, a) = ∫ σ̃ ( (s, a, 1)⊤w ) φ(s′)ψ(s′; dw), (4.6)
where µ and ψ(s′; ·) are probability measures in P2(Rd+1) for any s′ ∈ S , Br is a positive scaling parameter, and φ(s′) : S → R+ is a nonnegative function.
(ii) The reward function r satisfies that r(s, a) ≥ 0 for any (s, a) ∈ S×A. For the representation of r in (4.5) and the representation of the transition kernel P in (4.6), we assume that
χ2(µ ∥ ρw,0) < Mµ, χ2 ( ψ(s; ·) ∥ ρw,0 ) < Mψ, ∀s ∈ S,∫
φ(s)ds ≤M1,φ, ∫ φ(s)2ds ≤M2,φ,
where ρw,0 is the marginal distribution of ρ0 with respect to w, i.e., ρw,0 = ∫ ρ0(db, ·), χ2 is the chi-squared divergence, and Mµ, Mψ , M1,φ, M2,φ are absolute constants.
(iii) We assume that there exists an absolute constant G such that∥∥ψ(s; ·)− ψ(s′; ·)∥∥ Ḣ−1(µ) < G, ∥∥ψ(s; ·)− µ∥∥
Ḣ−1(µ) < G,∥∥ψ(s; ·)− µ∥∥ Ḣ−1(ψ(s′;·)) < G, ∥∥ψ(s; ·)− ψ(s′; ·)∥∥ Ḣ−1(ψ(s′′;·)) < G, ∀s, s ′, s′′ ∈ S,
where ∥·∥Ḣ−1(·) is the weighted homogeneous Sobolev norm.
We remark that by assuming ψ to be a probability measure and that φ(s′) ≥ 0 in (4.6), the representation of the transition kernel does not lose generality. Specifically, the function class of (4.6) is the same as
P = {∫ σ̃((s, a, 1)⊤w)ψ̃(s′; dw) ∣∣∣ ψ̃(s′; ·) is a signed measure for any s′ ∈ S}.
See §C.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-width neural network has a two-layer structure by (4.2). We establish the following lemma on the regularity of the representation of the action value function Qπ by such a neural network.
Lemma 4.4 (Regularity of Representation of Qπ). Suppose that Assumptions 4.2 and 4.3 hold. For any policy π, there exists a probability measure ρπ ∈ P2(RD) for the representation of Qπ with the following properties.
(i) For functionQ(s, a; ρπ) defined by (3.4) with ρ = ρπ and the action value functionQπ(s, a) defined by (2.2), we have Q(s, a; ρπ) = Qπ(s, a) for any (s, a) ∈ S ×A.
(ii) By letting Bβ ≥ 2(Br + γ(1− γ)−1BrM1,φ) for the neural network defined in (4.2) and ρ0 ∼ N (0, ID) for the initial distribution, we have W̃2(ρπ, ρ0) ≤ D̄ for any policy π, where we define W̃2(·, ·) = αW2(·, ·) as the scaled W2 metric. Here constant D̄ depends on the discount factor γ and the absolute constants L0,β , L1,β , lβ , Br,Mµ,Mψ,M1,φ,M2,φ defined in Assumptions 4.2 and 4.3.
Proof. See §B.2 for a 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(·; ρπ) in (3.4). Note that a larger Bβ captures a larger function class in (4.3). Without loss of generality, we consider thatBβ ≥ 2(Br+γ(1−γ)−1BrM1,φ) holds in the sequel. Hence, by Property (ii), it holds that W̃2(ρπ, ρ0) ≤ O(1) for any policy π. In particular, it holds by Property (i) of Lemma 4.4 that ∥Qt−Qπt∥2,ϕ̃πt = ∥Q(·; ρt)−Q(·; ρπt)∥2,ϕ̃πt 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 ρ0 ∼ N (0, ID) is the initial distribution. We specify the weighting distribution Φ̃πt in MF-TD (3.5) as Φ̃πt = Ẽπt
ϕ̃πt , where ϕ̃πt ∈ P(S ×A) is the evaluation distribution for the policy evaluation
error in Theorem 4.1. Then, it holds that
(1−√γ) · ∥Qt −Qπt∥22,ϕ̃πt ≤ − d
dt
W̃ 22 (ρt, ρπt)
2η +∆t, (4.7)
where
∆t =2α 1/2η−1BB1 · W̃2(ρt, ρ0)W̃2(ρt, ρπt) + α−1B2 · ( 4B1 max { W̃2(ρπt , ρ0), W̃2(ρt, ρ0) } +Br ) W̃2(ρt, ρπ) 2.
Here B1 and B2 are defined in (4.4) of Assumption 4.2, η is the relative TD timescale, α is the scaling parameter of the neural network, and W̃2 = αW2 is the scaledW2 metric. Moreover, constant B depends on the dicount factor γ, the scaling parameter Bβ in (4.2), and the absolute constants lβ , Br,M1,φ,G defined in Assumptions 4.2 and 4.3.
Proof. See §B.3 for a detailed proof.
Here we give a nonrigorous discussion on how to upper bound ∆t in (4.7). If W̃2(ρt, ρ0) ≤ O(1) holds for any t ∈ [0, T ], by W̃2(ρπt , ρ0) ≤ O(1) in Lemma 4.4 and the triangle inequality of W2 distance [59], it follows that W̃2(ρt, ρπt) ≤ O(1) and ∆t ≤ O(α1/2η−1 + α−1). Taking a time average of integration on both sides of (4.7), the policy evaluation error 1T ∫ T 0 ∥Qt −Qπt∥2,ϕ̃πtdt is then upper bounded by O(η−1T−1 + α1/2η−1 + α−1). Inspired by such a fact, we introduce the following restarting mechanism to ensure W̃2(ρt, ρ0) ≤ O(1).
Restarting Mechanism. Let W̃0 = λD̄ be a threshold, where D̄ is the upper bound for W̃2(ρπ, ρ0) by Lemma 4.4, λ ≥ 3 is a constant scaling parameter for the restarting threshold, ρt is the distribution of the parameters in the neural network at time t, and ρ0 is the initial distribution. Whenever we detect that W̃2(ρt, ρ0) reaches W̃0 in the update, we pause and reset ρt to ρ0 by resampling the parameters from ρ0. Then, we reset the critic with the newly sampled parameters while keeping the actor’s policy πt unchanged and continue the update.
The restarting mechanism guarantees W̃2(ρt, ρ0) ≤ λD̄ by restricting the distribution ρt of the parameters to be close to ρ0. Moreover, by letting λ ≥ 3, we ensure that ρπt is realizable by ρt since W̃2(ρπt , ρ0) ≤ D̄ ≤ λD̄, which means that the neural network is capable of capturing the representation of the action value function Qπt . We remark that by letting W̃0 = O(1), we allow ρt to deviate from ρ0 up to W2(ρt, ρ0) ≤ O(α−1) in the restarting mechanism. In contrast, the NTK regime [15] which corresponds to letting α = √ M in (3.1) only allows ρt to deviate from ρ0 by the chi-squared divergence χ2(ρt ∥ ρ0) ≤ O(M−1) = o(1). That is, the NTK regime fails to induce a feature representation significantly different from the initial one. Before moving on, we summarize the construction of the weighting distribution Φ̃πt in Theorem 4.1 and 4.5 as follows,
Φ̃πt = Ẽπt ϕ̃πt
, ϕ̃πt = 1
2 ϕ̃0 +
1 2 ϕ0 ⊗ πt, ϕ0 = ∫ A ϕ̃0(·,da), (4.8)
where ϕ̃0 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
1
T ∫ T 0 ( J(π∗)− J(πt) ) dt ≤ ζ
T︸︷︷︸ (a) +4κ
√ α−1S1 + α1/2η−1S2 + η−1D̄2
2T (1−√γ)︸ ︷︷ ︸ (b) , (4.9)
where we have
ζ = Es∼Eπ∗D0
[ KL ( π∗(· | s) ∥π0(· | s) )] , κ = ∥∥∥∥ Ẽπ∗D0 ϕ̃0 ∥∥∥∥ ∞ ,
S1 = (1 + λ)2D̄2B2(4B1λD̄ +Br)
1−√γ , S2 =
2BB1λ(1 + λ)D̄2
1−√γ .
Here Br, B1 and B2 are defined in Assumption 4.2 and 4.3, D̄ is the upper bound for W̃2(ρπ, ρ0) in Lemma 4.4, B depends on the discount factor γ 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 N that
N ≤ (λ− 2)−1 ( (α−1ηS1 + α 1/2S2)2TD̄ −2(1−√γ) + 1 ) .
Proof. See §B.4 for a detailed proof.
Note that for a given MDP with starting distribution D0, the expected KL-divergence ζ and the concentrability coefficient κ are both independent of the two-timescale update. We remark that our condition for (4.9) to be bounded is not restrictive. Specifically, we only need a given π0 and ϕ̃0 such that the KL-divergence ζ <∞ and the concentrability coefficient κ <∞, 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 T → ∞. 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 λ = 3 for the restarting threshold. By letting η = α3/2, it holds that (b) = O(α−1/2) as α → ∞. Thus, we have that 1T ∫ T 0 ( J(π∗) − J(πt) ) dt descends at a rate of O(T−1 + O(α−1/2) + O(α−3/4T−1/2)). Note that η = α3/2 shows that the critic has a larger relative TD timescale in (3.5). As for the total number of restartings N , it holds that N ≤ O(α1/2T ) as α→ ∞, which induces a tradeoff, i.e., a larger α guarantees a smaller gap in 1 T ∫ T 0 ( J(π∗)− J(πt) ) dt but yields in more restartings and a larger relative TD timescale.

heading: 5 Acknowledgement
text: 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 (Theory of Reinforcement Learning). This work was also supported in part by the Vannevar Bush Faculty Fellowship program under grant number N00014-21-1-2941.

Review summaries:
This paper studies actor-critic (AC) algorithms with neural networks from a theoretical perspective.   The 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. 

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 .