{"id":"2df87cc3-731e-4d13-b51f-01c1d0896867","arxiv_id":"2505.00663","paper_version":1,"verdict":"CONDITIONAL","confidence":"MODERATE","novelty_score":7.0,"correctness_risk":"medium","formal_verification":"none","parameter_count":3,"one_line_summary":"WPO derives a closed-form policy update from Wasserstein gradient flows, which for Gaussian policies coincides with the standard policy gradient in expectation but with lower variance, and works for arbitrary stochastic policies without reparameterization.","lead":"Wasserstein Policy Optimization (WPO) is a new reinforcement learning algorithm that updates stochastic policies using the gradient of the action-value function, without needing the reparameterization trick. It matches or outperforms standard continuous control methods on benchmarks and may scale better to very high-dimensional action spaces.","discovery_kind":"new_method","skeptic_critique":{"model":"deepseek-v4-flash","headline":"The 'completely general' claim rests on an unproven projection: for non-Gaussian policies the WPO update (Eq. 6) is not shown to ascend J, and Eq. 5's integration by parts has a nonzero boundary term for constrained-support policies, including the paper's own exponential example.","rationale":"This paper has real independent support: the Gaussian-case equivalence is proven with care (Sec. 4.1, App. A.3); the c-Wasserstein extension (App. C.3) is principled; the MPO-hyperparameter ablation (App. C.2) strengthens the baseline comparison; the Control Suite and fusion experiments are extensive; and the Acme implementation exists even though the paper honestly notes it is not the exact experimental code (footnote, Sec. 2).\n\nThe load-bearing weakness is the step from the nonparametric Wasserstein flow to the parametric update of Eq. 6. The projection in Sec. 2.3 minimizes local KL/Fisher distance to the flow-updated distribution; because the Wasserstein metric differs from the KL metric, and because exact representability of the flow by the parametric family is the exception rather than the rule, the projected direction need not be an ascent direction for J. The paper proves the expected update equals the policy gradient for Gaussians — hence locally ascends J — but for non-Gaussian families it provides only the exponential-family computation (Sec. 4.1) and one mixture-of-Gaussians experiment (Sec. 4.2). Moreover, my recomputation of Eq. 5 with the boundary term kept shows the identity fails for the exponential family on [0,∞), so even the paper's own illustration of 'arbitrary distributions' is not, as computed, the projection of the Wasserstein flow.\n\nThe reader's CONDITIONAL verdict is well calibrated; my review sharpens the condition into two checkable items: the sign of ⟨F^{-1}h, ∇_θJ⟩ over non-Gaussian policy families, and the boundary-term correction to Eq. 5. If the first test yields any non-positive directional derivative, the abstract's 'completely general' claim must be re-scoped (e.g., to large-support, boundary-vanishing policies or to Gaussian families), while the algorithm itself, its Gaussian analysis, and its empirical results remain valuable. Because the fix is re-scoping and verification rather than abandonment, the verdict stays CONDITIONAL.","tokens_in":19803,"tokens_out":34766,"duration_ms":318420,"concrete_test":"Primary check: for a two-component Gaussian-mixture policy as in Sec. 4.2, using exact 1-D quadrature (or ≥10^7 samples) for g = E[Q∇_θ log π], h = E[∇_θ∇_a log π ∇_aQ], and the exact Fisher FIM, compute s = ⟨F^{-1}h, g⟩ over a grid of (ρ, μ_i, σ_i) for Q(a) = −a⁴/100 + a² and for several random cubic polynomials; any configuration with s ≤ 0 shows the projected-flow update can descend J for a non-Gaussian policy, refuting the implied generality claim. Secondary check: recompute F_{tθ} for the exponential family on [0,∞) keeping the a = 0 boundary term in App. A.2 Eq. 26; if F_{tθ} = (E[∇_aQ] − ∇_aQ(0))/β², then Eq. 5 fails for the paper's own Sec. 4.1 example.","verdict_should_be":"UNCHANGED","load_bearing_attack":"The central claim — a 'simple and completely general closed-form update' derived from a Wasserstein gradient flow (Abstract; Sec. 2.3) — requires that the parametric projection convert the nonparametric ascent flow ∂π/∂t = −∇·(π∇_aQ) into a parameter-space step that increases J[π_θ]. That conversion is verified only for Gaussian policies (Sec. 4.1), where the expected WPO update coincides with the policy gradient. For non-Gaussian policies no ascent property is established: the projection minimizes D_KL(π_θ || π_θ + (∂π/∂t)dt − ∇_θπ Δθ) under the local Fisher metric, and the resulting update Δθ = F^{-1}h with h = E[∇_θ∇_a log π ∇_aQ] is not analyzed for the sign of the directional derivative ⟨∇_θJ, F^{-1}h⟩. Ascent of the continuous-time flow in the Wasserstein metric does not automatically survive projection in the KL metric, nor does it constrain the discrete-time step, so the 'completely general' claim rests on an unproven assumption; the mixture-of-Gaussians experiment (Sec. 4.2, Fig. 3) is a single favorable example.\n\nSecond, the derivation of Eq. 5 (App. A.2, Eq. 26) discards the boundary term under the assumption that ∇_θπ vanishes at the action-space boundary. For any policy with nonzero boundary density — including the exponential policy π(a) = (1/β)e^{−a/β} on [0,∞) used in Sec. 4.1 — the boundary term is −∇_aQ(0)/β² ≠ 0 in general, so F_{tθ} = (E[∇_aQ] − ∇_aQ(0))/β², not E[∇_aQ]/β²; the claimed identity in Eq. 5 is false for this family, and the update computed in Sec. 4.1 is not the projection of the Wasserstein flow the derivation claims to implement. This is a concrete, checkable failure of the generality claim, independent of the metric-mismatch concern.","agreement_with_reader":"agree"},"referee_report":{"model":"deepseek-v4-flash","summary":"The manuscript proposes Wasserstein Policy Optimization (WPO), an actor-critic algorithm whose update is obtained by projecting the Wasserstein gradient flow of the return functional onto a parametric policy family. The central object is the closed-form update Δθ = F^{-1} E_π[∇θ∇a log π ∇a Q] (Eq. 6). The paper shows that for Gaussian policies this update coincides with the standard policy gradient and with reparameterized DPG/SVG-style updates, analyzes a mixture-of-Gaussians example, and then introduces a practical diagonal-Gaussian implementation with variance rescaling and KL regularization. Experiments cover DeepMind Control Suite, replicated high-dimensional tasks, and a TCV tokamak magnetic-control task, reporting competitive or faster learning relative to MPO, SAC, and DDPG.","tokens_in":20181,"tokens_out":10367,"duration_ms":109805,"significance":"WPO is an interesting and potentially useful update: it gives a stochastic-policy analogue of deterministic policy gradients without the reparameterization trick, and the Gaussian analysis is clean. The paper contributes a broad empirical comparison, an open-source Acme implementation, and an exact derivation for full-support Gaussian policies. However, the headline claim of a 'completely general' update is currently stronger than what is proved: the derivation has a boundary-term flaw for constrained-support policies, and the non-Gaussian projection step is not shown to preserve ascent. If the generality claim were restricted or proved, the paper would be a solid contribution; as written, the theoretical part needs revision.","major_comments":[{"comment":"The integration by parts leading to Eq. (5) drops the boundary term under the assumption that ∇θπθ(a|s) vanishes at infinity. This fails for any policy with constrained support. A concrete counterexample is the exponential policy π(a)=β^{-1}e^{-a/β} on [0,∞) used in §4.1: ∇βπ(0)=-β^{-2}, so the corrected identity is F_tβ=(Eπ[∇aQ]-∇aQ(0))/β², not Eπ[∇aQ]/β², and the update shown in that section is not the WPO update for this family. The 'completely general' claim in the abstract and §2.3 should be restricted to policies whose support is the full action space with appropriate decay, or the update should be modified to include the boundary term.","section":"§2.3 / App. A.2, Eq. (26)"},{"comment":"The projection step from the nonparametric Wasserstein flow to the parametric update minimizes a local KL/Fisher quadratic, but no argument shows that the resulting discrete step increases J[πθ] for non-Gaussian policies. The continuous-time flow is an ascent in the Wasserstein metric on densities; ascent does not automatically survive projection in the KL/Fisher metric, and the paper provides no bound on or sign analysis of ⟨∇θJ, F^{-1}F_tθ⟩ outside the Gaussian case. The mixture-of-Gaussians example in §4.2 is a single favorable numerical instance, so the 'completely general' update is not established by the evidence presented.","section":"§2.3"},{"comment":"The non-Gaussian evidence is weakened by the fact that the mixture-of-Gaussians experiment uses the heuristic rescaling of gradients by σ_i² rather than the FIM of the mixture, and the practical algorithm in §5 additionally restricts to diagonal Gaussian policies and replaces the full Fisher by a variance rescaling. The text is transparent about these choices, but the experimental results therefore do not validate Eq. (6) for arbitrary stochastic policies; the abstract and conclusion should align the generality claim with what is actually implemented and tested.","section":"§4.2 / §5"}],"minor_comments":[{"comment":"The caption refers to the 'left plot' twice and is inconsistent about which panel shows the policy-gradient evolution versus the WPO evolution; please relabel or re-describe the panels.","section":"Fig. 3 caption"},{"comment":"The caption says 'fusion task discussed in Section 6.2', but the fusion task is presented in Section 6.3.","section":"Fig. 6 caption"},{"comment":"The claim that the Gaussian equivalences extend 'straightforwardly' to multivariate normals would benefit from an appendix derivation, because the covariance parameterization introduces off-diagonal terms.","section":"§4.1"},{"comment":"The footnote states that the public Acme implementation is not the version used for the experiments; please state clearly in the main text whether the exact experimental configuration will be released.","section":"Footnote 1"},{"comment":"The integration-by-parts assumption should be stated as a decay condition on the full boundary term, e.g., ∇θlogπ(a|s) π(a|s) ∇aQ(s,a) → 0, rather than only on ∇θπ(a|s); this is the condition actually needed for Eq. (30) to hold.","section":"App. A.2, Eq. (26)"}],"recommendation":"major_revision","confidential_remarks":"This is a substantial empirical paper from a strong group, and the Gaussian analysis plus the broad benchmark evaluation are valuable. The main risk is the gap between the theoretical generality claim and what is proved. The boundary-term issue is concrete and must be fixed, either by restricting the claim to full-support policies or by incorporating boundary corrections. I would not reject if the claims are scaled back appropriately; the projection-ascent question could be answered with a proof for a useful policy family or by an explicit statement that the non-Gaussian case is only heuristic."},"author_rebuttal":null,"desk_editor":{"model":"deepseek-v4-flash","letter":"Two things to know about this paper. It derives a parametric policy update from Wasserstein gradient flow—Eq. (6)—that is new to RL and genuinely elegant: F^{-1} E_π[∇_θ∇_a log π ∇_a Q]. It avoids reparameterization and can in principle handle arbitrary stochastic policies. The Gaussian analysis is the cleanest part: they show the expected WPO update coincides with the standard policy gradient, and that DPG/SVG(0) coincide with WPO in the Gaussian case. That is a real result and worth citing.\n\nThe experiments are extensive and fair: 49 Control Suite tasks, a high-dimensional combined-replica task, and a tokamak fusion task, with careful hyperparameter checks against MPO. The paper is honest about sensitivity to network nonlinearity and squashing. The Acme implementation is released, though it is not the exact experimental code.\n\nBut the central claim is overstated. First, the projection onto parametric policies is done with the KL divergence and a local Fisher approximation, not the Wasserstein distance. For non-Gaussian policies there is no proof that the update ascends J; the mixture-of-Gaussians example is a single favorable simulation. That is a caveat, not a fatal flaw—natural gradient methods have the same gap.\n\nSecond, and more concretely: the derivation of Eq. (5) in Appendix A.2 drops the boundary term when integrating by parts, assuming ∇_θ π vanishes at the action-space boundary. That assumption fails for the paper's own exponential-policy example in Sec. 4.1, which has nonzero density at a=0. The correct expression for that family is (E[Q'] − Q'(0))/β^2, not E[Q']/β^2. So the update they compute for the exponential policy is not the projection of the Wasserstein flow. This is a checkable error in the paper's own illustrative example, and it undercuts the 'completely general' claim in the abstract.\n\nThe practical algorithm is also a long way from the idealized update: diagonal Gaussians, a variance-rescaling heuristic, tuned KL penalties, and a cube-root squashing function on the fusion task. The paper is transparent about this, but it means the empirical results validate a hybrid update, not the derived one.\n\nMy recommendation: send it to peer review. The core idea is sound and novel, the empirical work is solid, and the flaws are fixable. Ask the authors to soften the generality claim, redo the boundary term carefully, and either prove or drop the ascent claim for non-Gaussians. This is a good paper that needs revision, not rejection.","headline":"A genuinely new parametric update for stochastic control with a clean Gaussian analysis, but the 'completely general' claim is undercut by a dropped boundary term in Eq. 5 and a missing ascent proof for non-Gaussians.","tokens_in":20849,"tokens_out":5096,"would_cite":true,"duration_ms":48009,"reading_group":"yes","serious_thinker":"yes","would_accept_peer_review":true},"rs_alignment":null,"lean_confirmation":null,"pith_extraction":{"msc":[],"pacs":[],"model":"deepseek-v4-flash","headline":"Wasserstein Policy Optimization claims a single closed-form update that uses action-value gradients yet works for arbitrary stochastic policies without reparameterization.","keywords":["Wasserstein gradient flow","policy optimization","actor-critic","continuous control","deterministic policy gradient","reparameterization trick","stochastic policies","high-dimensional action spaces"],"falsifier":"Take a one-state MDP with action-value $Q(a) = -a^4/100 + a^2$ and a two-component Gaussian-mixture policy initialized with means at $\\pm 1$ and large variances, the setting of the paper's Figure 3; compute the exact expected WPO update with the full Fisher matrix and no KL penalty, and check whether the return $J[\\pi]$ increases for a sufficiently small step size. A step that decreases $J$, or fails to increase it as the step size tends to zero, would contradict the claim that WPO is an ascent direction.","tokens_in":19532,"feed_emoji":"🎛️","tokens_out":10251,"duration_ms":92785,"temperature":0.7,"pith_summary":"The paper sets out to give reinforcement learning in continuous action spaces a policy update that is a true gradient step in the space of probability distributions over actions, not just in parameter space. It derives Wasserstein Policy Optimization (WPO) as the parametric projection of the Wasserstein gradient flow that maximizes expected return, and arrives at a closed-form formula: the natural-gradient-rescaled expectation of $\\nabla_\\theta \\nabla_a \\log \\pi$ times the action-value gradient $\\nabla_a Q^\\pi(s,a)$. If this derivation is right, one update rule combines the two families that were previously separate: it exploits action-value gradients (as deterministic policy gradient does) while still updating any stochastic policy by plain score-function sampling, with no reparameterization trick. This matters because deterministic methods explore poorly and reparameterized stochastic methods exclude distributions like mixtures of Gaussians; WPO claims to avoid both limitations. The paper supports this claim with a Gaussian-case equivalence result, a mixture-of-Gaussians illustration, and experiments on continuous-control benchmarks and a tokamak magnetic-control task.","feed_headline":"Wasserstein flow yields closed-form update for any stochastic policy","feed_subtitle":"The rule borrows deterministic policy gradients' action-value gradients and adds stochastic exploration without reparameterization.","key_machinery":"The load-bearing object is the Wasserstein gradient-flow PDE $\\partial \\pi / \\partial t = -\\nabla_a \\cdot (\\pi \\nabla_a Q^\\pi)$ (with sign chosen for ascent), the steepest descent of expected return under the 2-Wasserstein metric. The paper's step is to project this nonparametric flow onto a parametric family by minimizing the KL divergence between the updated and the target policy; expanding that KL to second order gives the Fisher matrix $F_{\\theta\\theta}$ and the cross term $F_{t\\theta} = \\mathbb{E}_{a\\sim\\pi}[\\nabla_\\theta \\nabla_a \\log \\pi_\\theta(a|s) \\, \\nabla_a Q^\\pi(s,a)]$, obtained by integration by parts. The resulting closed-form update $\\theta \\leftarrow \\theta + F_{\\theta\\theta}^{-1} F_{t\\theta}$ is what carries the argument. In practice the paper specializes to diagonal Gaussian policies, rescales mean and variance gradients by $\\sigma_i^2$ and $\\tfrac{1}{2}\\sigma_i^2$ to approximate the Fisher inverse, and adds a KL penalty against the previous policy to prevent premature collapse; it also notes that the c-Wasserstein extension replaces $\\nabla_a Q$ by a squashed version $\\nabla c^*(\\nabla_a Q)$.","core_discovery":"The central claim is that WPO is a simple and completely general closed-form policy update that is a parametric approximation to the Wasserstein gradient flow of the expected-return functional. For any policy $\\pi$, the update is $\\theta \\leftarrow \\theta + F_{\\theta\\theta}^{-1} \\mathbb{E}_{a\\sim\\pi}[\\nabla_\\theta \\nabla_a \\log \\pi_\\theta(a|s) \\, \\nabla_a Q^\\pi(s,a)]$, where $F_{\\theta\\theta}$ is the Fisher information matrix of the policy; the mixed score-function derivative lets the algorithm carry the action-value gradient through the policy without reparameterizing the sampling distribution. The functional derivative $\\delta J/\\delta \\pi = (1/(1-\\gamma)) Q^\\pi d^\\pi$ makes $Q^\\pi$ the driving field, so the update inherits DPG's use of $\\nabla_a Q^\\pi$, while the expectation over $a \\sim \\pi$ keeps it a stochastic-policy method. For Gaussian policies the expected WPO update equals both the classic policy gradient and the reparameterized DPG update, so the paper argues that the differences surface only for non-Gaussian distributions such as mixtures of Gaussians, where WPO is applicable and behaves distinctly. Empirically, WPO is competitive with DDPG, SAC, and MPO across the DeepMind Control Suite, learns faster on replicated high-dimensional control tasks (up to 105 action dimensions), and matches MPO on a simulated tokamak magnetic-control task while driving the policy variance down.","pith_inferences":["If the KL projection is a faithful surrogate for the Wasserstein flow, WPO's dynamics are governed by the 2-Wasserstein geometry rather than by the Fisher metric; a testable consequence is that WPO should be more robust to reparameterizations of the policy than natural policy gradient, a contrast the paper's exponential-policy example begins to illustrate.","The paper's observation that sampled WPO updates have zero variance for locally linear $Q$ suggests that its variance reduction should become more pronounced as the action dimension grows; one can test this directly by measuring gradient variance of WPO versus classic policy gradient on the replicated high-dimensional tasks.","The c-Wasserstein squashing variant (cube root applied to $\\nabla_a Q$) opens a family of updates, each a valid descent direction for a different cost function; comparing squashing functions on tasks with rapidly changing $Q$ could identify which geometric penalization is best for stability.","Because WPO is presented as a genuine gradient method in distribution space, it may combine naturally with entropy or KL bonuses as regularizers rather than as parts of the objective, potentially yielding a cleaner exploration-exploitation trade-off than a tuned entropy coefficient."],"forward_implications":["WPO gives a single update rule that reduces to the standard policy gradient for Gaussian policies while matching the reparameterized DPG update, so the method inherits exploration from stochastic sampling and the variance reduction of action-value gradients.","Because no reparameterization is needed, any policy with a computable mixed score-function gradient—including mixture-of-Gaussians or exponential distributions—can be trained by WPO, opening a class of policies that stochastic DPG methods cannot handle exactly.","On high-dimensional continuous control (replicated Humanoid tasks with 21, 65, and 105 action dimensions), the paper reports that WPO learns faster than DDPG, SAC, and MPO, with the advantage growing as the action dimension grows.","On the tokamak magnetic-control task, WPO matches MPO's reward while its policy variance shrinks toward deterministic over training, which the paper interprets as the expected behavior for a fully observed environment."],"supporting_citations":[{"why":"Supplies the Wasserstein gradient-flow theory (Eq. 1) from which the WPO update is derived.","marker":"Ambrosio et al., 2008"},{"why":"States the deterministic policy-gradient theorem that WPO generalizes to stochastic policies.","marker":"Silver et al., 2014"},{"why":"Introduces SVG(0), the reparameterized stochastic DPG method whose update WPO matches in the Gaussian case.","marker":"Heess et al., 2015"},{"why":"Provides the SAC algorithm and reparameterized stochastic actor used as a baseline and comparison point.","marker":"Haarnoja et al., 2018"},{"why":"Gives the same parametric projection of a Wasserstein flow and the c-Wasserstein extension used for gradient squashing.","marker":"Neklyudov et al., 2023"},{"why":"Defines the classic score-function policy gradient and the variance-rescaling heuristic adopted in WPO's practical update.","marker":"Williams, 1992"},{"why":"Supplies the MPO baseline and the KL-regularization form WPO adopts to prevent premature policy collapse.","marker":"Abdolmaleki et al., 2018"},{"why":"Gives the functional-derivative form of expected return that identifies $Q^\\pi$ as the driving field in the flow.","marker":"Agarwal et al., 2021"},{"why":"Provides the DeepMind Control Suite tasks on which WPO is compared with baseline algorithms.","marker":"Tassa et al., 2018"},{"why":"Defines the tokamak magnetic-control task and training setup on which WPO is evaluated against MPO.","marker":"Tracey et al., 2024"}],"fun_headline_variants":["WPO does DPG without reparameterization","Wasserstein flow: deterministic gradient, stochastic policies","Closed-form WPO update for any policy distribution","No reparameterization: WPO for stochastic policies","Wasserstein flow unifies DPG and stochastic gradients"],"cache_read_input_tokens":3200,"weakest_assumption_plain":"The central assumption is that replacing the true Wasserstein steepest-descent direction with its best local KL-divergence approximation inside the parameterized policy family still points uphill in expected return; the paper gives no proof of this for non-Gaussian policies.","fun_headline_variants_meta":{"raw":{"variants":["WPO does DPG without reparameterization","Wasserstein flow: deterministic gradient, stochastic policies","Closed-form WPO update for any policy distribution","No reparameterization: WPO for stochastic policies","Wasserstein flow unifies DPG and stochastic gradients"]},"model":"deepseek-v4-flash","effort":"low","cost_usd":0.001737,"raw_usage":{"total_tokens":6881,"prompt_tokens":981,"completion_tokens":5900,"prompt_tokens_details":{"cached_tokens":384},"prompt_cache_hit_tokens":384,"prompt_cache_miss_tokens":597,"completion_tokens_details":{"reasoning_tokens":5825}},"tokens_in":597,"tokens_out":5900,"duration_ms":40569,"temperature":1.0,"reasoning_tokens":5825,"cache_read_input_tokens":384,"cache_creation_input_tokens":0},"cache_creation_input_tokens":0},"created_at":"2026-08-16T04:36:29.928289+00:00","model_set":{"reader":"deepseek-v4-flash"},"falsifier":"Take a one-state MDP with action-value $Q(a) = -a^4/100 + a^2$ and a two-component Gaussian-mixture policy initialized with means at $\\pm 1$ and large variances, the setting of the paper's Figure 3; compute the exact expected WPO update with the full Fisher matrix and no KL penalty, and check whether the return $J[\\pi]$ increases for a sufficiently small step size. A step that decreases $J$, or fails to increase it as the step size tends to zero, would contradict the claim that WPO is an ascent direction.","supporting_citations":[{"cited_title":"Deterministic policy gradient algorithms","cited_arxiv_id":null,"evidence_quote":"States the deterministic policy-gradient theorem that WPO generalizes to stochastic policies."},{"cited_title":"Learning continuous control policies by stochastic value gradients","cited_arxiv_id":null,"evidence_quote":"Introduces SVG(0), the reparameterized stochastic DPG method whose update WPO matches in the Gaussian case."},{"cited_title":"Wasserstein quantum M onte C arlo: a novel approach for solving the quantum many-body schr \\\"o dinger equation","cited_arxiv_id":null,"evidence_quote":"Gives the same parametric projection of a Wasserstein flow and the c-Wasserstein extension used for gradient squashing."},{"cited_title":"T., Tassa, Y., Munos, R., Heess, N., and Riedmiller, M","cited_arxiv_id":null,"evidence_quote":"Supplies the MPO baseline and the KL-regularization form WPO adopts to prevent premature policy collapse."}],"review_version":1}