Research map / Reinforcement learning / Policy gradient and actor-critic
Wasserstein Policy Optimization
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes Wasserstein Policy Optimization (WPO), an actor-critic method that projects the Wasserstein gradient flow of expected return onto a parametric policy, yielding a closed-form update that uses action-value gradients with arbitrary stochastic policies.
Problem
Policy optimization for continuous action spaces: existing methods either use classic stochastic policy gradients that ignore action-value gradients, or DPG-style methods that are restricted to deterministic policies or require the reparameterization trick.
Approach
Derives a parametric policy update by projecting the 2-Wasserstein gradient flow of the expected return onto the policy parameters via the Fisher information, giving a closed-form update theta += F^-1 E[grad_theta grad_a log pi * grad_a Q] that resembles DPG but works for any stochastic policy without reparameterization; practical version uses diagonal Gaussian scaling and KL regularization.
Key results
- On DeepMind Control Suite, WPO takes off robustly where DDPG and SAC often fail and is in the same range as the best method across nearly all tasks.
- On combined Humanoid Stand tasks with 1, 3, and 5 replicas (21, 65, 105 action dims), WPO learns faster than MPO, DDPG and SAC, with the margin growing with replica count.
- On the TCV fusion task, WPO achieves slightly higher reward than MPO while its policy std decreases toward deterministic behavior.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use WPO instead of MPO | continuous-control tasks in the DeepMind Control Suite · DeepMind Control Suite | return | Mixed |
| Use WPO instead of DDPG or SAC | continuous-control tasks in the DeepMind Control Suite · DeepMind Control Suite | learning progress | Helped |
| Use WPO instead of MPO | replicated Humanoid Stand tasks · Humanoid Stand, 1, 3, and 5 replicas | learning speed | Helped |
| Use WPO instead of MPO | TCV tokamak control · Fusion shape_70166 variant | reward | Helped |
| Use WPO instead of policy gradient | a symmetric two-component Gaussian mixture policy · One-dimensional non-concave action-value example | convergence speed | Helped |
| Replace MPO's tuned hyperparameters with WPO-matched hyperparameters (ablation) | MPO · DeepMind Control Suite | performance | No clear effect |
| Add variance rescaling to MPO with WPO-matched hyperparameters (ablation) | MPO · DeepMind Control Suite | learning speed | Helped |
| Use WPO-matched hyperparameters without variance rescaling (ablation) | MPO · Humanoid CMU and Hopper | learning speed | Hurt |
| Replace ELU with SiLU in the actor and critic networks (ablation) | WPO · DeepMind Control Suite | return | Mixed |
| Apply elementwise cube-root squashing to action-value gradients (ablation) | WPO · DeepMind Control Suite | return | Mixed |
| Use a soft KL penalty instead of a hard KL constraint (ablation) | WPO · DeepMind Control Suite | performance | No clear effect |
| Add a KL penalty on the policy mean (ablation) | WPO · DeepMind Control Suite | convergence speed | Hurt |
Surprising findings
- For Gaussian policies, the expected natural WPO update coincides with the standard policy gradient in expectation, despite the different sampled estimators.
- Without KL regularization WPO prematurely collapses to a deterministic policy and fails to learn, especially on the fusion task.
Limitations
- Diagonal Fisher approximation ignores parameter-to-distribution Jacobians, a heuristic rather than the exact natural update.
- Can collapse to deterministic policy without KL regularization; WPO struggled on Dog domain where state space exceeds observation space.
- Results on fusion and combined tasks are limited to Gaussian policies and a single set of hyperparameters.
Experimental setup
Datasets: DeepMind Control Suite TCV tokamak simulator (shape_70166 task)
Benchmarks: DeepMind Control Suite Humanoid Stand combined replicas TCV fusion shape_70166
Metrics: episode return
Baselines: MPO DDPG SAC SVG(0)
Builds on: deterministic policy gradient MPO natural policy gradient SVG(0)
Closest papers by meaning
- Policy Optimization as Wasserstein Gradient Flows (ICML 2018 · similarity 0.91)
- ReLU to the Rescue: Improve Your On-Policy Actor-Critic with Positive Advantages (ICML 2024 · similarity 0.88)
- Actor-Free Continuous Control via Structurally Maximizable Q-Functions (NeurIPS 2025 · similarity 0.88)
- Value Improved Actor Critic Algorithms (NeurIPS 2025 · similarity 0.87)
- Mean Field Langevin Actor-Critic: Faster Convergence and Global Optimality beyond Lazy Learning (ICML 2024 · similarity 0.87)
- Actor-Critics Can Achieve Optimal Sample Efficiency (ICML 2025 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.