Global Convergence of Wasserstein Policy Gradient for Entropy-Regularized Reinforcement Learning
Zhaoyu Zhu, Rui Gao, Shuang Li
原文摘要
Wasserstein policy gradient (WPG) optimizes stochastic policies in the optimal-transport geometry of action distributions. In entropy-regularized reinforcement learning, WPG combines transport driven by gradients of the soft Q-function with entropy-induced diffusion. The resulting dynamics are coupled through Bellman policy evaluation, and the policy optimization objective is generally nonconvex. We establish a unified global convergence theory of WPG for infinite-horizon discounted Markov decision processes with general state spaces and continuous action spaces. We identify a statewise gradient-dominance inequality by exploiting the Bellman structure together with log-Sobolev inequalities, which we show to hold uniformly along the trajectory. This inequality relates the squared norm of the Wasserstein policy gradient to the value gap in the sup norm. In continuous time, we prove well-posedness of the flow and exponential convergence of the value gap in the sup norm and of the policy in statewise KL divergence. In discrete time, by additionally controlling the regularity and the discretization error, we obtain geometric convergence to a neighborhood of the optimum with constant step sizes and an $O(\varepsilon^{-1})$ iteration complexity with diminishing step sizes for fixed problem parameters. Finally, we establish faster local convergence rates for both dynamics, governed by the geometry of the optimal policy.
cs.LG