Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Yufeng Zhang, Siyu Chen, Zhuoran Yang, Michael I. Jordan, Zhaoran Wang

TL;DR

The paper develops a mean-field analysis of neural actor-critic learning with two-layer overparameterized networks and two-timescale updates, linking the critic's TD updates to a Wasserstein gradient flow and the actor's PPO updates to replicator dynamics. In the infinite-width, continuous-time regime, it proves global optimality at a sublinear rate and shows the critic's representation can adapt within an $O(1/\alpha)$ neighborhood of the initial features. The analysis extends beyond NTK by allowing data-dependent representation learning and introduces a restarting mechanism to control the evolution of the parameter distribution. These results provide finite-time, non-asymptotic convergence guarantees for neural AC in a mean-field setting and offer theoretical insight into representation learning in policy-based RL.

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

TL;DR

The paper develops a mean-field analysis of neural actor-critic learning with two-layer overparameterized networks and two-timescale updates, linking the critic's TD updates to a Wasserstein gradient flow and the actor's PPO updates to replicator dynamics. In the infinite-width, continuous-time regime, it proves global optimality at a sublinear rate and shows the critic's representation can adapt within an neighborhood of the initial features. The analysis extends beyond NTK by allowing data-dependent representation learning and introduces a restarting mechanism to control the evolution of the parameter distribution. These results provide finite-time, non-asymptotic convergence guarantees for neural AC in a mean-field setting and offer theoretical insight into representation learning in policy-based RL.

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.
Paper Structure (20 sections, 15 theorems, 133 equations)

This paper contains 20 sections, 15 theorems, 133 equations.

Table of Contents

  1. Introduction
  2. Background
  3. Policy Optimization and Actor-Critic
  4. Wasserstein Space
  5. Algorithm
  6. Main Result
  7. Convergence of Mean-field PPO
  8. Global Optimality and Convergence of Two-timescale AC
  9. Acknowledgement
  10. Proofs of Main Results
  11. Proof of Theorem \ref{['thm:ppo']}
  12. Detailed Statement of Lemma \ref{['lem:RAVF']}
  13. Proof of Theorem \ref{['thm:td']}
  14. Proof of Theorem \ref{['thm:ttac']}
  15. Proofs of Supporting Lemmas
  16. ...and 5 more sections

Key Result

Theorem 4.1

Let $\pi^* = \mathop{\mathrm{argmax}}_{\pi} J(\pi)$ be the optimal policy and $\pi_0$ be the initial policy. Then, it holds that where $\widetilde{\phi}^{\pi_t}\in \mathscr{P}({\mathcal{S}}\times\mathcal{A})$ is an evaluation distribution for the policy evaluation error and $\zeta = \mathbb{E}_{s\sim \mathcal{E}_{\mathcal{D}_0}^{\pi^*}}\Bigl[ {\mathrm{KL}}\bigl(\pi^*(\cdot {\,|\,} s) \,\|\, \pi_0

Theorems & Definitions (16)

  • Theorem 4.1: Convergence of MF-PPO
  • Lemma 4.4: Regularity of Representation of $Q^\pi$
  • Theorem 4.5: Upper Bound of Policy Evaluation Error
  • Theorem 4.6: Global Optimality and Convergence Rate of Two-timescale AC with Restarting Mechanism
  • Lemma A.1: Regularity of Representation of $Q^\pi$
  • Lemma C.1
  • Lemma C.2
  • Lemma C.3
  • Lemma C.4
  • Lemma C.5
  • ...and 6 more