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

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 . 1 Introduction . 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 .

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 .