Global optimality of softmax policy gradient with single hidden layer neural networks in the mean-field regime

Abstract

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

arXiv comment22 pages, 1 figure

References in corpus (11)

• Playing Atari with Deep Reinforcement Learning
• Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor
• Neural Tangent Kernel: Convergence and Generalization in Neural Networks
• A Mean Field View of the Landscape of Two-Layers Neural Networks
• On Lazy Training in Differentiable Programming
• Neural Policy Gradient Methods: Global Optimality and Rates of Convergence
• On the Global Convergence Rates of Softmax Policy Gradient Methods
• Fast Global Convergence of Natural Policy Gradient Methods with Entropy Regularization
• Learning to Walk via Deep Reinforcement Learning
• PC-PG: Policy Cover Directed Exploration for Provable Policy Gradient Learning
• Can Temporal-Difference and Q-Learning Learn Representation? A Mean-Field Theory