Wasserstein Policy Optimization

Report | Slides | GitHub repository

I completed this study project for MATH 277A at UC San Diego, taught by Prof. Xiaochuan Tian. The method is not my own: my goal was to understand Wasserstein Policy Optimization (WPO) from the underlying optimal-transport geometry, then reproduce its qualitative behavior in small numerical experiments with a simplified implementation.

Why transport a policy?

Ordinary policy gradient improves a stochastic policy by reweighting the log-likelihood of sampled actions according to scalar value estimates. Deterministic policy gradient instead differentiates the critic with respect to the action, which uses local geometry in action space but requires a deterministic policy and a separate exploration mechanism. WPO combines aspects of both viewpoints: it retains a stochastic policy while moving its probability mass along directions determined by the action-value gradient.

For a fixed state $s$, regard $\pi(\cdot\mid s)$ as a probability measure over actions. Informally, its Wasserstein gradient flow follows the continuity equation

This equation describes probability mass being transported in action space along the velocity field $\nabla_a Q^{\pi_t}(s,a)$. Unlike pointwise likelihood reweighting, the update knows that nearby actions are geometrically related.

With a parameterized policy $\pi_\theta$, the functional transport direction must be projected back onto the finite-dimensional parameter space. The resulting WPO update combines mixed derivatives of the policy log-density with the critic’s action gradient, together with a Fisher-type preconditioner. For a diagonal Gaussian policy, the repository uses a simplified analytic scaling that avoids constructing the full neural-network Fisher matrix.

Experiments

The repository contains my small implementation of two experiments based on the WPO paper. These are pedagogical reproductions, not new benchmark results.

1. Mixture of Gaussians. In a one-dimensional objective with two optima, I compared policy gradient, natural policy gradient, and WPO. The WPO update transported mass toward both modes in a stable and balanced manner, while the other updates were more susceptible to premature variance shrinkage or mode collapse. This reproduces the qualitative observation in the original paper rather than serving as a broad benchmark.

2. Inverted pendulum. I compared policy gradient, deterministic policy gradient, and WPO inside a common DDPG-style actor-critic implementation. A raw critic-weighted likelihood gradient was unstable in this off-policy setting, so I also evaluated a sample-based value baseline. With that correction, WPO achieved learning behavior comparable to deterministic policy gradient and regularized policy gradient without manually specified exploration noise.

The pendulum experiment also exposed an important limitation: without KL or entropy regularization, the stochastic WPO policy quickly became nearly deterministic. This is harmless for a simple task with an almost deterministic optimal controller, but it can cause premature collapse in more difficult problems. My implementation intentionally omits several stabilization techniques from the official implementation, so I interpret the results as a geometric demonstration rather than evidence that WPO should replace established actor-critic algorithms.

Reflection

The most useful aspect of this reproduction was conceptual. WPO treats policy optimization as motion in a space of probability measures, where the metric reflects the geometry of the action space. Working through the derivation and experiments connected optimal transport, gradient flows, and continuous-control reinforcement learning in a concrete way. It also showed me where the elegant functional picture becomes difficult in practice: one still needs an accurate critic, a tractable projection into the parameterized policy family, and regularization that prevents the transported distribution from collapsing too early.