Policy Optimization as Wasserstein Gradient Flows

Abstract

Policy optimization is a core component of reinforcement learning (RL), and most existing RL methods directly optimize parameters of a policy based on maximizing the expected total reward, or its surrogate. Though often achieving encouraging empirical success, its underlying mathematical principle on {\em policy-distribution} optimization is unclear. We place policy optimization into the space of probability measures, and interpret it as Wasserstein gradient flows. On the probability-measure space, under specified circumstances, policy optimization becomes a convex problem in terms of distribution optimization. To make optimization feasible, we develop efficient algorithms by numerically solving the corresponding discrete gradient flows. Our technique is applicable to several RL settings, and is related to many state-of-the-art policy-optimization algorithms. Empirical results verify the effectiveness of our framework, often obtaining better performance compared to related algorithms.

arXiv commentAccepted by ICML 2018; Initial version on Deep Reinforcement Learning Symposium at NIPS 2017

Cited by in corpus (10)

• Improving Sequence-to-Sequence Learning via Optimal Transport
• Regularizing activations in neural networks via distribution matching with the Wasserstein metric
• A contribution to Optimal Transport on incomparable spaces
• Sqrt(d) Dimension Dependence of Langevin Monte Carlo
• Hessian-Free High-Resolution Nesterov Acceleration for Sampling
• The Wasserstein Proximal Gradient Algorithm
• Learning Policies through Quantile Regression
• Generalized Talagrand Inequality for Sinkhorn Distance using Entropy Power Inequality
• Distributionally-Constrained Policy Optimization via Unbalanced Optimal Transport
• Improving Adversarial Text Generation by Modeling the Distant Future