Wasserstein Proximal Policy Gradient
This work addresses the problem of efficient policy optimization in continuous-action reinforcement learning for researchers and practitioners in the field.
The authors tackled the problem of continuous-action, entropy-regularized reinforcement learning and achieved a global linear convergence rate for their proposed method, Wasserstein Proximal Policy Gradient (WPPG), with competitive performance on standard continuous-control benchmarks. WPPG attains this result without evaluating the policy's log density or its gradient.
We study policy gradient methods for continuous-action, entropy-regularized reinforcement learning through the lens of Wasserstein geometry. Starting from a Wasserstein proximal update, we derive Wasserstein Proximal Policy Gradient (WPPG) via an operator-splitting scheme that alternates an optimal transport update with a heat step implemented by Gaussian convolution. This formulation avoids evaluating the policy's log density or its gradient, making the method directly applicable to expressive implicit stochastic policies specified as pushforward maps. We establish a global linear convergence rate for WPPG, covering both exact policy evaluation and actor-critic implementations with controlled approximation error. Empirically, WPPG is simple to implement and attains competitive performance on standard continuous-control benchmarks.