Wasserstein Policy Optimization
Authors: David Pfau
Ian Davies
Diana L Borsa
João Guilherme Madeira Araújo
Brendan Daniel Tracey
Hado van Hasselt
Venue: ICML 2025 poster (2025)
Keywords: Policy Optimization
Wasserstein metric
Optimal Transport
Gradient Flow
Deep Reinforcement Learning
Actor-Critic
Continuous Control
Primary area: reinforcement_learning
TL;DR: We derive a novel policy gradient algorithm from Wasserstein gradient flows and show that it is simple and effective at deep reinforcement learning for continuous control.

Abstract: We introduce Wasserstein Policy Optimization (WPO), an actor-critic algorithm for reinforcement learning in continuous action spaces. WPO can be derived as an approximation to Wasserstein gradient flow over the space of all policies projected into a finite-dimensional parameter space (e.g., the weights of a neural network), leading to a simple and completely general closed-form update. The resulting algorithm combines many properties of deterministic and classic policy gradient methods. Like deterministic policy gradients, it exploits knowledge of the *gradient* of the action-value function with respect to the action. Like classic policy gradients, it can be applied to stochastic policies with arbitrary distributions over actions -- without using the reparameterization trick. We show results on the DeepMind Control Suite and a magnetic confinement fusion task which compare favorably with state-of-the-art continuous control methods.

Wasserstein Policy Optimization
David Pfau
Ian Davies
Diana Borsa
João G. M. Araújo'
Brendan Tracey' Hado van Hasselt
Abstract
the policy update. The majority of these methods use pol- icy update derived from the classic policy gradient theorem for stochastic policies (Sutton et al., 1999), which applies to both discrete and continuous action spaces.
We introduce Wasserstein Policy Optimization (WPO), an actor-critic algorithm for reinforce- ment learning in continuous action spaces. WPO can be derived as an approximation to Wasserstein gradient flow over the space of all policies pro- jected into a finite-diment diteessiimnii parameter space (e.g., the weights of a neural network), leading to a simple and completely general closed-form update. The resulting algorithm combines many properties of deterministic and classic policy gra- dient methods. Like deterministic policy gradi- ents, it exploits knowledge of the gradient of the action-value function with respect to the action. Like classic policy gradients, it can be applied to stochastic policies with arbitrary distributions over actions without using the reparameteriza- tion trick. We show results on the DeepMind Control Suite and a magnetic confinement fusion task which compare favorably with state-of-the- art continuous control methods.
One notable exception is the class of algorithms derived from the deterministic policy gradient (DPG) theorem (Sil- ver et al., 2014; Lillicrap, 2015). These use information about the gradient of the value in action space update the policy, but are limited to deterministic policies, hindering exploration. Adding in stochasticity can be difficult to tune, and extensions that learn the variance (Heess et al., 2015; Haarnoja et al., 2018) rely on the reparameterization trick, which limits the class of policy distributions that can be used.
Here present new policy gradient method that uses gradients of the action value but can learn arbitrary stochas- tic policies. The method can be derived from the theory of Wasserstein gradient flows (Ambrosio et al., 2008), and projecting the nonparametric flow onto the space of para- metrie functions (e.g. neural networks) leads to a simple closed-form update that combines elements of stochastic and deterministic policy gradients. We give high-level illustration of this in Fig. 1. The update is natural to imple- ment as an actor-critic method. We call the resulting method Wasserstein Policy Optimization (WPO).
1. Introduction
Reinforcement learning has made impressive progress in controlling complex domains with continuous actions such as robotics (Haarnoja et al., 2024), magnetic confinement fusion (Degrave al., 2022) and game playing (Farebrother & Castro, 2024). A significant portion of this progress can be attributed to policy optimization methods, which directly optimize the parameters of a policy by stochastic gradient ascent on the expected long term return (e.g., Schulman et al., 2015a; 2017; Abdolmaleki et al., 2018). While un- biased estimates of the policy gradient can be computed directly from returns (Williams, 1992), many practical deep reinforcement learning algorithms for continuous control employ an actor-critic approach (Barto et al., 1983), which also estimates a value function to diminish the variance of
The paper is organized as follows. First, we review policy optimization for continuous control, and Wasserstein gra- dient flows. We then show how the latter can be used to derive WPO when applied to the former. We review prior work in this area and analyze the dynamics of WPO and how it relates to other policy gradients. We describe several extensions needed to convert WPO into a practical and com- petitive deep reinforcement learning algorithm. Finally, we compare its performance against several baseline methods on the DeepMind Control Suite (Tassa et al., 2018; Tun- yasuvunakool et al., 2020) and a task controlling magnetic coils in simulated magnetic confinement fusion device (Tracey et al., 2024). On a task where the task dimension can be varied arbitrarily, we find that WPO learns faster than baseline methods by an amount that increases as the task dimension grows, suggesting it may work well in very high dimensional action spaces. An open- implementation
'Google DeepMind, London, UK. Correspondence to: David Pfau Bacaue@googg...om comm>.
Proceedings of the 42nd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Wasserstein Policy Optimization
(aQ"(s,a)
Q a)
VaQ a)
Policy Gradient a)(s,a)V,logg
DPG v₂0"(s, a
WPO V aVV, Vglogn
Figure 1. Conceptual illustration of how WPO combines elements of stochastic and deterministic policy gradient methods. Left: "classic" policy gradient. Samples are taken from a stochastic policy. Each sample contributes a scalar Q" (s, a) factor to the gradient. Middle: deterministic policy gradient (DPG). A deterministic action is chosen and the policy gradient depends on the gradient QT a). Right: Wasserstein policy optimization (WPO). Samples are taken from a stochastic policy, as in classic policy gradient, but depend on the gradient of Q" with respect the action, as DPG.
of WPO is available in Aeme (Hoffman et al., 2020)'.
et al., 2018). This is especially critical for achieving good performance with deep reinforcement learning.
2. Method
One notable exception to this is deterministic policy gra- dients (DPG) (Silver et al., 2014; Lillicrap, 2015; Barth- Maron et al., 2018), which can be seen as the limit of the policy gradient update as the policy becomes deter- ministic. These algorithms were developed as early as the 1970s (Werbos, 1974) and were known under names such as 'action-dependent heuristic dynamic programming' (Prokhorov & Wunsch, 1997) or 'gradient ascent on the value' (van Hasselt & Wiering, 2007). As the name sug- gests, this only applies to deterministic policies that map state vectors to single action vector, so a, (s:)). Then the deterministic policy gradient has the form VeJ [Tg] = Eg,x VVa, Q(s,, a, [s;(s where VoT(S:) is the Ja- cobian of the deterministic policy. The appearance of the gradient of the value in action space in the policy gradi- ent is potentially useful in high-dimensional action spaces. However, the restriction to deterministic policies makes ex- ploration difficult. While there are extensions to DPG for learning a stochastic policy, such as SVG(0) (Heess et al., 2015) and SAC (Haarnoja et al., 2018), they rely on the reparameterization trick, which limits their generality.
2.1. Policy Gradient and Iteration Methods
We aim to find policy <aass) which is distribution over continuous space of actions a E "l"" conditioned on state s E R that maximizes the expected long-term discounted return J [T] = Ea ET >'re] for a Markov de- cision process with transition distribution P(s'|s, a) and reward function r(s, where 7, is shorthand for ((S,,a)).
While "classic" policy gradient methods come in many forms (Schulman et al., 2015b), they are mostly variants of the basic update VaJ[me]
EarNT. Et V, (.Vologr((ayss)) (ayls)) and only differ in the choice of scalar V,. If U is entirely a function of the returns like - 7f Te', then the policy can be optimized directly, as in REINFORCE (Williams, 1992). If V is the action-value function (st ar) or some transformation of the action-value function like the advantage function, or soft- max of the action-value, then both a policy and value func- tion must be estimated simultancously, which is standard in actor-critic methods (Barto et al., 1983). Additionally, it is common to add some form of trust region regulariza- tion to prevent the policy update from changing too much on any step (Schulman et al., 2015a; 2017; Abdolmaleki
Separately, policy iteration algorithms (Howard, 1960; Sut- ton & Barto, 2018) estimate the value of the current policy and then derive improved policy based on these values. Iterating this converges to the optimal values and policy, These algorithms do not necessarily follow the gradient of the value. A modern example is MPO (Abdolmaleki et al., 2018), which updates its parametric policy (i.e., a neural network) towards the exponentiation of the current action values, using a target policy π ataa x exp( ((s.a)
httpp:: / //i thub. wtt.c/omc///o// Note that the implementation in Acme for the experiments implementation on DeepMind Control Suite tasks and found qualitatively similar performance to that reported here.
goog 1 deepmi nd/ cme.
the version used this paper. However we have run this
2

---

Wasserstein Policy Optimization
and minimizing KL divergence with respect to this target.
with the environment, and (1 - > is just constant. In what follows, we foeus per-state updates where these terms do not appear.
2.2. Wasserstein Gradient Flows
is possible to derive a true policy gradient for stochastic policies with form similar DPG, based on Wasserstein gradient flows. The theory of gradient flows is discussed in depth in (Ambrosio et al., 2008), and we review the relevant results here. Although the discussion here is fully general, we keep the notation as close to the RL notation as possible to avoid confusion.
2.3. Wasserstein Policy Optimization
To convert the theoretical results in the previous section into practical algorithm, need to approximate the PDE in Eq. with an update to parametric function such as neural network. Starting at Ty from the parametric family of functions, for given infinitesimal dt and flow 2ma, ät we solve for the choice of infinitesimal AO which minimizes the KL divergence between the original distribution and the updated distribution. This gives the optimal direction to cancel out the flow, hence the minus sign in the KL below. It is well known that the KL divergence is approximated locally by a quadratic form with the Fisher information matrix (Pascanu & Bengio, 2013):
Consider an arbitrary functional J[r] of a probability den- sity o (a)) and let 5.7 be the functional derivative of J. Then minimum of J:
8т the following PDE will converge
Ov Öt
(T
(T Va
57 8T
=
Va
(1)
If we think of V 67 as a velocity field, then this may di be recognizable as the continuity equation from fluid me- chanics, the drift term in the Tokker-Planck equation, or in a machine learning context as one way of writing the expression for the likelihood of neural ODE (Chen al., 2018). Problems in optimal transport can also be framed in terms of the continuity equation (Benamou & Brenier, 2000). The 2- Wasserstein distance, conventionally defined as
ÔTO dt VOTO Öt
DKL
TT
To +
dt
T
Fir
Fif
dt
2
(-10)
A0
Frg
Foo
(-10)
alogoo (als))
27
Fu = Er
Ot
Alogaa (als)
Fi = Er
V VologTo(als))
/ p(a, b) b||2dadb (2)
Ot
inf (T m1)= (xo.))))))
ime (als)
Ol
where T [TT, m ) is the space of all joint distributions with marginals and can also be expressed as:
Fou = E, Vologr which is minimized at A = Fog Fig-
W:2 (TO, T1) inf /'' Earn [a)l(( di (3) VVEV(To..))
Next, we derive a simple expression for Feg. We plug in Ora 3t
V. (Ta (als) (a. for ascent on J[r]:
where V(TO 11) is the set of velocity fields ve such that, if and On. de V. (Te (V..; (a)),, then
ne = no T. - ग at 1. From this dynamic formulation, it can be
Fro V VologTo @mo(als)
da
at
shown that the flow in Eq. is the steepest descent the
= / V VologTo (T -mg- VoValogo((als) (s.
functional I in the space of probability densities under the
a)) da
metric induced locally by the 2- 2-Wasserstein distance.
(5)
In our case J J[m] is the expected return in an MDP and we want to maximize this quantity. The functional derivative (wrt σ) has the form:
This derivation is expanded in Sec. A.2 in the appendix, but mainly follows from integration by parts. This leads to simple closed-form update for parametric policies which we call the Wasserstein Policy Optimization update:
8J Ju
1 -
ss a) =
(" (s) 1-7
(4)
0 F Ex
[ V (aloo
(6)
where d= (s) (1- 21
s) is the discounted state occupancy function. A derivation is given in Sec. A. of the appendix. This is the functional generalization of the policy gradient in the tabular setting (see Agarwal et al. (2021), 1), Eq. 7). The d= (s) term typically emerges implicitly in the update as the sampling frequency when interacting
A similar parametric approximation was derived for finding the ground state energy of quantum systems (Neklyudov et al., 2023). While this application is quite different, the derivation follows from the same theory of gradient flows, arriving very similar expression. To the best of our
3

---

Wasserstein Policy Optimization
Sampled WPO updates
Expected WPO update
Expected WPO update
((a)
0.8
(())
0.8
(())
0.8
updated sta)
0.6
0.6
0.6
-2
0.4
-2
0.4
0.4
-3
0.2
-3
0.2
1-
0.2
[(6)
0.0
0.0
3 0.0
-2
1
2
3
-2
-1
0
2
3
-2
-1
2
Figure 2. Concrete WPO updates for a single-variate normal policy for two different action-value functions. In the left and middle plots we consider ((a) -a" /2. with obvious optimum = 0, and policy with μ mean at several sampled actions. These are averaged to produce the update, moving the mean towards the optimal action. The expected WPO update, shown in the middle plot, is then Aμ I and Ao -a Both the mean and variance will decrease (or, generally, move probability mass to the optimal action), as shown. Conversely, in the right plot we consider Q(a) = a /2 then Ap = and Ar a, and both the mean and variance will increase, as shown. In all cases, the expected value of the resultant policy is increased.
1. The left plot shows the gradient the
knowledge we are the first to use this update for policy optimization reinforcement learning. This somewhat ide- alized update requires the full Fisher information matrix and lacks some extensions needed to make policy optimization methods work in practice. We describe how to extend this update into practical deep RL algorithm in Sec. 5.
amount to different way to update the critic, whereas WPO is novel way to update the actor.
The idea of policy optimization flow has appeared few times in the literature. Richemond & Maginnis (2018) show that policy optimization with an entropy bonus can be written as a PDE similar to the expres- sion in Sec. 2.2 with an additional diffusion term, but they stop short of deriving an update for parametric policies. The work of Zhang al. (2018) also formulates policy optimiza- tion as a regularized Wasserstein gradient flow, but arrives at more complicated DPG-like update that uses the repa- rameterization trick, similarly to SVG(0) and SAC. There is close connection between Wasserstein gradient flows and Stein variational gradient descent (Chewi et al., 2020), but methods based on SVGD require expensive particle-based approximations (Messaoud et al., 2024).
Wasserstein gradient
3. Related Work
The general idea of using the Wasserstein metric in reinforce- ment learning has appeared in many forms. In Abdullah et al. (2019) the 2-Wasserstein distance is used to define Wasserstein Robust Reinforcement Learning, an algorithm that trains policy which is robust to misspecification of the environment for which is being trained. The Wasserstein distance appears here as a constraint on the transition dy- namies of the environment, rather than as way of defining learning dynamics. In Moskovitz al. (2020), the Wasser- stein metric is used locally as alternative preconditioner in place of the Fisher information matrix in natural policy gradient (Kakade, 2001), while the conventional form of the policy gradient is still used. The Wasserstein distance has also been explored as an alternative to the KL divergence as a regularizer to prevent the policy from changing too quickly (Pacchiano et al., 2020).
4. Analysis
4.1. Gaussian Case
To better understand the mechanics of the WPO update, we analyze the simplified case where the policy is a single- variate normal distribution and the policy and value are not state-dependent. In this case we have:
The Wasserstein metric has also been used to define dis- tances between state distributions, for various ends (Ferns et al., 2004; He et al., 2021; Castro et al., 2022), whereas we are focused on gradient flows in the space of actions.
a -
V log ((a) =
Va (gg)(a),
(a
1
In a combination of the above aims, (Metelli al., 2019) and (Likmeta et al., 2023) use the Wasserstein distance define method to propagate uncertainty across state-action pairs in the Bellman Equation, with the aim of using that quantified uncertainty to better deal with the exploration- exploitation trade-off inherent to online reinforcement learn- ing. the context of actor-critic methods, this would
g
Fug =
Fa
Fap = 0.
Let A, and A 0 be the contributions to the update due to the gradients of the mean and variance, respectively, such

---

Wasserstein Policy Optimization
MoG Volue Punction
Policy Grodiont
WPO
30
20
20
Q(a)
0.3
I7(a)
10
10
20
0.2
0-
10
0.1
-10
-10
0.0
-20
-20
-10
2000
4000 6000 8000 10000
2000
4000 6000 8000 10000
Iteraaions
Iterations
Figure 3. Concrete WPO learning for a one dimensional mixture of Gaussians policy for the non-concave action-valus function Q(a) 100 a a². The left plot shows the action-value function and mixture of Gaussians policy. In the middle plot show the evolution of the policy under a standard policy gradient update. both with samples from the policy and the change in the means of each mixture component. On the left show the same evolution for WPO. WPO converges faster, is more stable around the optimum, and converges to both optima if the policy is initialized symmetrically.
that A 4,0 A,8.
For the mean, we then have normal distribution are then respectively E Q(aa and Eir, Voo]. For this choice of parameterization. (a = wa)aa, which exactly coincides with WPO. SAC is similar to SVG(O), and extends this by including an entropy bonus to encourage exploration, as well a mechanism based on double Q-learning (van Hasselt, 2010) to avoid value overestimations, Such extensions could be combined with WPO well.
4 F F VaQ(a) log N Vay
= = VaQ(a)
Vou
o²
= E. [a.((aVool)
This is very similar to the DPG update: (u) ex- cept that we take the gradient of the action value at μs, rather than sampling т. While this similarity holds for the normal distribution, is not necessarily true in general.
Note that the natural WPO update is independent of pa- rameterization, while SVG(0) not, so in general the two updates may still differ. For instance, suppose we have a task that requires non non-negative actions, and we decide use an exponential policy with components aş exp(a;/ /B))///// //, where 3, are learnable scale parameters for each action di- mension. The Fisher for this distribution is diagonal with components 1/5, on the diagonal. Then, the natural WPO update with respect to 3 is
For the variance we have:
A = E. V.Qlav log Voo)
Ex
Ex V.Q(a) V log = Voo
2
0²
Ex vaQ(a)V. "-"Voo Voo
0²
F F- (V.Q(a) V. V log (a))
a E o VV.((aVVVoo
diag (3²) E, (.Q(V)Q(.VV....... V log a)
diag (5²) E, [V.Q(a) (a/ 1/3)]
This update is little more complicated, but on inspection it also quite intuitive. The variance increases when (a - 4) and V. (()) have the same sign. So, when we sample action, then we increase the variance if the gradient of Q with respect to that action points even further away from the mean. If the gradient points back towards the mean, we will instead decrease the variance. Some special cases are illustrated in Figure 2.
diag (3a) Ex (aQ(a)diag (1/39)
E,, [V.Q(a)]
We can reparameterize the policy for SVG(O) with a stan- dard exponential 7 with density p(n = x) = exp(-x) and then define a = 3 on. Then, the SVG(O) update is
Stochastic extensions of DPG such as SVG(O) (Heess al., 2015) and soft actor-critic (SAC) (Haarnoja et al., 2018) use the reparameterization trick and define gradient E,,IV (s,, a. ( 7)), where π denotes the action as deterministic function of the state, as well as noise term n. For instance, we can define π (s, 77) = (ss) + (s)) o 7). where n ~ (0, and o denotes the elementwise prod-
Ey V V.Q((aV = E,, V VVQ (1) = Eg (VoQ(a)diag (a/3)]
where we used 1 a/3, by definition. This is not generally the same as the WPO update.
Not only does the natural WPO update coincide with DPG augmented with the reparameterization trick in the Gaussian
uct. The mean and variance updates for the single-variate
5

---

Wasserstein Policy Optimization
Humanoid CMU
Quadruped Run
Point Easy
Cartpole Two poles
Stand
1000
Episodic Return
800
600
400
200
Cartpole Three poles
Dog Run
Humanoid Run
Humanoid Stand
1000
Episodic Return
800
600
400
200
0.0
0.5
1.0
1.5
0.0
0.5
1.5
0.0
0.5
1.5
0.0
0.5
1.5
Actor Steps
1e7
Actor Steps
1e7
Actor Steps
1e7
Actor Steps
1e7
WPO
MPO
DDPG
SAC
Figure 4. Results from selected DeepMind Control Suite tasks. Full results are in Fig, 7 in the appendix.
case, also equals the standard policy gradient. Specifically,
(=Q(s, ((s) does not depend the action, and be- cause the action-value gradient is then the same for each sampled action the WPO update for the mean of the policy will have zero variance. In contrast, in this case the standard policy gradient update will have non-zero variance. This observation is consistent with Fig. I: the variance in the WPO updates will be low when locally all the gradients point roughly in the same direction.
A = E. (aQ((a Vonl = E. [vQ((a)
J Q(a)(((V) Q(a))a=(ssaa) Von
(as V V.E,(Q(a)) 0)
=Exa(QQ(a) V. [a)]] Vau E. (Q(a)) log T Vo4 (as V. log nog==- = v log m)
-
4.2. Mixture of Gaussian Case
where the last line is just the expected policy gradient up- date for the parameters of the mean. A similar deriva- tion, with the same result, can be done for the variance; this can be found in Appendix A.3. we add those con- tributions together, we get the standard policy gradient E, [Q(s, (V log (s, Vou Va log (s,, (oo)) E, [Q(s, V log a).
To understand the qualitative differences between different updates, we will have move beyond the case of Gaussian policies. We consider the case of one-dimensional mix- ture of Gaussians policy T TTa) = Σ. P.N μ4; (.,a,)) where Σ: P: are the mixture weights. Because the assign- ment of a sample to a mixture component is discrete latent variable, the reparametcrization trick cannot be used exactly (though can be approximated through relaxations such as the concrete/Gumbel-softee 2017; 7; Jang et al., 2017)) and WPO and classic policy gradient and not DPG/SVG(0), In this case, the dynamics of learning are too complex for closed-form results as in the previous section, so we con- sider an illustrative numerical example, with the action- value function Q(a) = T 4 + a², which has maxima at 555 We initialize the policy with two components with P; = 0.5, a, = 10 and 14; = +1. We used batch size of 1024 and learning rate of 0.003. For WPO, rather
These equivalences can be extended to multivariate normal distributions quite straightforwardly. This counterintuitive result suggests that there essentially only one correct way to update a Gaussian policy, and that the major differences between approaches will only become clear when going beyond simple normal distributions over actions.
distribution (Maddison et al.,
exclusively consider
Importantly, even if the expected natural WPO update coin- cides with the expected policy gradient update, the sampled updates could still have dramatically different variance. For instance, consider an action-valui function that (locally) linear in the actions, such that (s, a) = w(s)Ta. a. Then
6

---

Wasserstein Policy Optimization
than approximate the true Fisher information matrix, we rescale the gradients for Pt. 14 and o; by of. a choice which is justified by the form of the FIM for a single Gaussian. We show results in Fig. 3. Despite being the same in ex- pectation for the Gaussian case, policy gradient and WPO clearly have qualitatively different learning dynamics in the tixture-of-Gaussian case. WPO converges faster, more stable around the minimum, and finds both local maxima (so long as the mixture components metrically). Notably, early in optimization, the variance of the policy actually increases for WPO, when the means of the mixture components are in the region with positive curvature, consistent with the intuition in Fig. 2.
the fusion task from Tracey al. (2024). While a variety of forms of KL regularization are used in continuous control, we closely follow the form used in MPO (Abdolmaleki et al., 20. We can express this as soft constraint, and modify the loss so that at each step we take a step towards solving
T
max Bar
nd (Earaalrd
initialized sym-
or express it as hard constraint
T Σ
max N.
5. Implementation
s.t. E BLND [KL. P (1s))()) <c
(8)
We make two modifications the update in Eq. 6 to make WPO into a practical method. First, the use of the full natu- ral gradient update is not practical for deep neural networks. It is tempting to simply drop it, but this could lead to serious numerical stability issues, as the Valoga (a) s) term in the update blows up as the policy converges to deterministic result. Note that this is different from classic policy gra- dient methods, where natural gradient descent accelerates learning but is not strictly necessary (Kakade, 2001).
which can be implemented by treating the a in the soft constraint as a Lagrange multiplier and performing a dual optimization step. Here denotes a previous state of the policy, e.g., a target network. In either case, the gradient of the KL penalty is computed conventionally - only the gradient of the reward uses the approximate Wasserstein gradient flow and variance rescaling. While the KL diver- gence could be replaced by the Wasserstein distance, which would be more mathematically consistent and has been ex- plored elsewhere as a regularizer (Richemond & Maginnis, 2018; Zhang et al., 2018; Pacchiano al., 2020), we find the KL. divergence works well in practice and leave it to future work to explore alternatives. We also note that while the KL. penalty slows the convergence of WPO to a deterministic policy, it does not prevent it, as we show in Fig. 6.
We could in theory use approximate second order methods such as KFAC (Martens & Grosse, 2015), but we have found that a much simpler approximation works well in practice. First, while the WPO update can be applied to arbitrary stochastic policies, we focus on normally-distribure cies To ( is) AV (s)) Es(s) where the covariance is constrained to be diagonal with elements (s). The normal distribution with diagonal covariance has a diagonal Fisher information matrix with 1 for the mean element μ4 and 2 o? for the standard deviation oz
poli-
6. Experiments
To evaluate the effectiveness of WPO, we evaluate it on the DeepMind Control Suite (Tassa et al., 2018; Tunyasuvu- nakool et al., 2020), a set of tasks in MuJoCo (Todorov et al., 2012). These tasks vary from one-dimensional actions, like swinging a pendulum, up to a 56-DoF humanoid. We addi- tionally consider magnetic control of tokamak plasma in simulation, problem originally tackled by MPO in Degrave et al. (2022). On Control Suite, compare WPO against both conceptually related and state-of-the-art algorithms which can be used in the same setting: Deep Deterministic Policy Gradient (DDPG; Lillicrap, 2015), Soft-Actor Critic (SAC: Haarnoja et al., 1)1,, and Maximum a Posteriori Policy Optimization (MPO; Abdolmaleki et al., 2018).
Thus, rather than try to approximate the full Fisher infor- mation matrix, we may redefine the gradient of the log like- lihood such that à iye, lol Σ) of Bjg
a logN(allu, )
2 = foo 8- Iog// 2). While this ignores the contribution of the gradients of μ and E with respect to 0 in the Fisher information matrix, it provides qualitatively correct scaling that cancels out the tendency of the likelihood gradient to blow up as Σ - O. A similar heuristic was suggested in the original REINFORCE paper (Williams, 1992).
3 and 30, log
Secondly, we regularize the policy with a penalty on the KL divergence between the current and past policy. Regular- ization to prevent the policy from taking excessively large steps is standard practice in deep reinforcement learning (e.g., Schulman et al., 2015a; 2017). We similarly found that without regularization, WPO will prematurely collapse onto a deterministic solution and fail to learn, especially on
Our training setup similar to other distributed RL systems (Hoffman et al., 2020): we run 4 actors in parallel to gener- ate training data for the Control Suite tasks, and 1000 actors for the tokamak task. For WPO, the policy update uses sequences of states from the replay buffer, which may come from an old policy, but the actions are resampled from the current policy, making the algorithm effectively off-policy
7

---

Wasserstein Policy Optimization
Single Base Task
Combined Environments
Combined Environments
1000
150
Epicodic Raturn
Epicodic Return
Epicodic Raturn
200
600
400
50
200
0.0
0.2
0.4
0.6
0.8
1.0
0.0
0.2
0.4
0.6
0.8
1.0
0.0
0.2
0.4
0.6
0.8
1.0
Actor Steps
1eB
Actor Steps
1eB
Actor Steps
1e8
WPO
MPO
DDPG
SAC
Figure 5. Plots of reward from various agents on combined Humanoid Stand environments. Left to right: 3 and 5 replicated environments (21, 65 and 105 action dimension). Solid line denotes the mean and shaded region highlights the minimum and maximum over 5 seeds. As the number of replicas grows, WPO is able to Iearn faster than other methods by a larger margin.
for states but on-policy for actions. MPO is implemented similarly. We used the soft KL penalty in Eq. 7 for Control Suite tasks, as we found the hard KL penalty did not no- ticeably improve results, but used the hard KL penalty for the fusion task for consisteney with previously published results. Separate KL penalties with different weights were put on the mean and variance of the policy, as described in Song et al. (2019). We found that the penalty on the mean had little effect on stability and mainly slowed convergence, while the penalty on the variance noticeably helped stability. Training hyperparameters are listed in Sec. B and an outline of the full training loop is given in Alg. I in the appendix. For each RL algorithm, the same hyperparameters were used for each control suite environment.
off on a number of tasks. This is even noticed for relatively low dimensional tasks (particularly with sparse reward). WPO robustly takes off and is in the same range as the best performing method across nearly all tasks. Through hyper- parameter tuning, we noticed that SAC is particularly sen- sitive to the weighting of its entropy objective. This meant that finding generally hyperparameters across all tasks led to difficult trade-offs of stability and performance. We note that WPO and MPO demonstrated greater out-of-the-box generalisation across Control Suite. On the Humanoid CMU domain, one of the highest dimensional tasks in the Control Suite, neither SAC nor DDPG took off at all, while WPO made progress on all tasks and, in the case of the Walk task, initially learned faster than MPO. On Dog, another high-dimensional domain, WPO converged to roughly the same final reward, but often took longer. Dog is the only domain where the state space is larger than the observation space, suggesting that WPO may have difficulty in partially observed settings. We experimented with several choices of nonlinearity, squashing function (see Sec. C.3 and Fig. 12), KL regularization weight, and critic bootstrapping update for WPO. This is still less than the years of experimentation which has gone into many other popular continuous control algorithms (Huang et al., 2022). This shows that, while DDPG or SAC may struggle to learn on certain tasks, WPO almost always makes some learning progress, and is often comparable to state-of-the-art methods for these tasks, even without significant tuning.
For the critic update we use a standard n-step TD update with target network:
ÖTD= Σ rut + 7"h ((S Qo an) (9)
where the target value V($++n) is approximated by sam- pling multiple actions from π. This target value could be the mean of the target action value network, the maximum, or something in between. In MPO, the softmax over samples is theoretically optimal, so we use that for Control Suite. We use the maximum for WPO on all Control Suite tasks, as that worked well on the hardest domains, but use the mean for both MPO and WPO on fusion tasks, the performance is significantly better.
6.2. Combined Tasks
6.1. Deep Control Suite
Methods that use action-value gradients may perform well in sigh-dimensional action spaces, but no Control Suite task goes beyond few dozen dimensions. To evaluate the effec- tiveness of WPO in higher dimensional action spaces, we construct tasks that consist of controlling many replicas of Control Suite environment simultancously with a single centralized agent. Specifically, the action and observation vectors are concatenated, and rewards are combined with a
Figure 4 shows results from subset of the DeepMind Con- trol Suite. We selected tasks which show the full range of WPO's performance from excelling to struggling. Learning curves for all Control Suite tasks are in Fig. 7 in the ap- pendix. While no single algorithm uniformly outperformed all others on all tasks, is notable that DDPG and SAC con- verged to lower rewards and occasionally struggled to take
8

---

Wasserstein Policy Optimization
0.19
0.14
0.27
0.38
0.34
WPO (Hard KL Constraint) MPO IHand Constraint)
0.43
0.50 0.50
5.75
1,00
1.25
3.75
2.50
3.50
0.52
0:75
1.00
1.25
1.0
1.75
2.50
Actor Steps
260
Actor Steps
260
Figure 6. Reward and policy average standard deviation evolution throughout training on the fusion task discussed Section 6.2
Smoot hMin n operator (see Eq. 31 in the Appendix), biasing the reward towards lower performing replicas to encourage learning across all replicas. We use the same agent hyperpa- rameters as we used in the standard control suite benchmark no matter the number of replicas. We selected Humanoid Stand as the base task due to its relatively high dimensional action space (21) and moderate difficultly.
MPO. Interestingly, we see a notable difference in the adap- tation of the policy variance between the two algorithms. WPO evolves in the direction of a deterministic policy as training evolves, as one would expect for a fully observed environment, while MPO maintains approximately constant policy variance average. We see similar results in B.3 the original shape_701 _70 66 task, where MPO and WPO achieve similar rewards with different policy adaptation be- havior. These results show WPO as a viable alternative to MPO in this complex real-world task.
The results in figure 5 demonstrate that as the number of replicas grows, WPO continues to be able to learn across the environments. In the singleton case, all methods are qualitatively similar except SAC which converges to lower reward than other learning algorithms. As the number of replicas grows, WPO takes off earlier in training than MPO, which takes off earlier than DDPG, with SAC being the slowest to take off even in the singleton case. This trend becomes more pronounced for larger numbers of replicas, even if the asymptotic performance similar for all meth- ods. This suggests that for tasks with hundreds of action dimensions, WPO may be able learn faster than other methods.
7. Discussion
We derived novel policy gradient method from the the- ory of optimal transport: Wasserstein Policy Optimization. The resulting gradient has an elegantly simple form, closely resembling DPG, but can be applied to learn policies with arbitrary distributions over actions. On the DeepMind Con- trol Suite, WPO quite robust, performing comparably to state-of-the-art methods on most tasks, despite little tun- ing. Results magnetic confinement fusion domains show that WPO can be applied to diverse tasks beyond simu- lated robotics. Promising initial results suggest it can learn quickly on tasks with over 100 action dimensions. We hope these results can help unlock performant algorithms for very larger control problems, and inspire researchers develop new, challenging nigh-dimensional benchmarks in continu- ous control with hundreds of action dimensions or more.
6.3. Fusion
Degrave et al. (2022) used MPO discover policies to con- trol the magnetic coils of the TCV tokamak, toroidal fusion experiment (Duval et al., 2024). The resulting policies suc- cessfully ran at 10kHz to hold the plasma stable. Follow-on work (Tracey et al., 2024) explored methods to improve the accuracy of RL-derived magnetic control policies. We run WPO on variant of the shape_701 66 task considered therein, modified reward higher control accuracy (details in Sec. B.3). This task has 93 continuous measurements, 19-dimensional continuous action space and lasts 10,000 environment steps, simulating Is experiment on TCV.
While we chose particular instantiation of WPO for the experiments, the WPO update is general and can be the foun- dation for many different implementations. We have only begun to scratch the surface of how WPO can be applied, and we are hopeful that extensions such as non- non-Gaussian policies or more advanced critic updates (Bellemare et al., 2017) could fully exploit the advantages of this new method. This could greatly increase the performance and scope of problems to which WPO can be applied.
We compare the performance of WPO (with cube-root squashing as described in Sec. C.3) alongside MPO in Fig. 6. Both MPO and WPO used the default network and settings as in Tracey et al. (2024), including the hard KL regulariza- tion in Eq. 8. WPO achieved slightly higher reward than

---

Wasserstein Policy Optimization
Acknowledgements
Benamou, J.-D. and Brenier, Y. A computational fluid chanics solution to the longgekantoroviccl mass transfer problem. Numerische Mathematik, 84(3):375-393, 2000.
We would like to thank Maria Bauza, Nimrod Gileadi, Ab- bas Abdolmaleki for assistance running experiments, and Tom Erez, Jonas Buchli, Alireza Makhzani, Guy Lever, Leonard Hasenclever and Kirill Neklyudov for helpful dis- cussions.
Carpanese, F. Development of fice-boundary equilibrium and transport solvers for simulation and real-time interpre- tation of tokamak experiments. Technical report, EPFL, 2021.
Impact Statement
Castro, P. S., Kastner, T., Panangaden, P., and Rowland, M.
Mico: Improved representations via sampling-base state similarity for markov decision processes, 2022. URL https://arxiv.e /2tp://...g.abb/......oo 08229.
The present work can be useful for any reinforcement learn- ing domain with continuous actions, such as robotics or industrial control. As this fundamental algorithmic ad- vance, it does not favor any particular socially beneficial or dangerous application, but instead could enable progress across all applications. Other policy optimization methods like PPO are widely used in alignment of generative AI models, so WPO could conceivably have applications in alignment in domains with continuous action spaces, which could help to avert undesirable behavior.
Chen, R. T., Rubanova, Y., Bettencourt, J., and Duvenaud, D. K. Neural ordinary differential equations. Advances in neural information processing systems, 31, 2018.
Chewi, S., Le Gouic, T., Lu, C., Maunu, T., and Rigollet, P. Svgd kernelized wasserstein gradient flow of the chi-squared divergence. Advances in Neural Information Processing Systems, 33:2098-2211 2109, 2020.
References
Clevert, D.-A., Unterthiner, T., and Hochreiter, S. Fast and accurate deep network learning by exponential linear units (elus). arxiv 2015. International Conference on Learning Representations, 2016.
Abdolmaleki, A., Springenberg, T., Tassa, Y., Munos, R., Heess, N., and Riedmiller, M. Maximum a posteri- ori policy optimisation. In International Conference Learning Representations, 2018.
Degrave, J., Felici, F., Buchli, J., Neunert, M., Tracey, B., Carpanese, F., Ewalds, T., Hafner, R., Abdolmaleki, A., de Las Casas, D., et al. Magnetic control of tokamak plasmas through deep reinforcement learning. Nature, 602(7897):414 2022.
Abdullah, M. A., Ren, H., Ammar, H. B., Milenkovic, V., Luo. R., Zhang, M., and Wang. J. Wasserstein robust reinforcement learning, 2019. URL https://axxii btg/abs/1907.133
Duval, B., Abdolmaleki, A., Agostini, M., Ajay, C., Alberti, S., Alessi, E., Anastasiou, G., Andrèbe, Y., Apruzzese, G., Auriemma, F., et al. Experimental research on the
Agarwal, A., Kakade, S. M., Lee, J. D., and Mahajan, G. On the theory of policy gradient methods: Optimality, approximation, and distribution shift. Journal of Machine Learning Research, 22(98): 2021.
TCV tokamak. Nuclear Fusion, 64(1 (M(11))112023) oct 2024. doi: 10. 00.10888/7411-432//a88833 doi. 10t.://1...111/11111111 / adB361.
URL https://p////. dx.
Ambrosio, L., Gigli, N., and Savaré, G. Gradient flows: in metric spaces and in the space of probability measures. Springer Science & Business Media, 2008
Elfwing. S., Uchibe, E., and Doya, K. Signoid-weighted linear units for neural network function approximation in reinforcement learning. Neural Networks, 107:3-11, 2018.
Barth-Maron, G., Hoffman, M. W., Budden, D., Dabney, W., Horgan, D., Tb, D., Muldal, A., Heess, N., and Lil- licrap, T. Distributed distributional deterministic policy gradients. International Conference on Learning Repre- sentations, 2018
Farebrother, and Castro, S. Cale: Continuous arcade learning environment. arXiv preprint arXin:2410.23810. 2024.
Barto, A. G., Sutton, R. S., and Anderson, C. W. Neuronlike adaptive elements that can solve difficult learning con- trol problems. IEEE transactions on systems, man, and cybernetics, (5)8334844, 1983.
Ferns, N., Panangaden, P., and Precup, D. Metrics for finite markov decision processes. In Uncertainty in Artificial Intelligence, 2004.
Haarnoja, T., Zhou, A., Abbeel, P., and Levine, S. Soft actor-critic: Off-policy maximum entropy deep reinforce- ment learning with a stochastic actor. In International conference on machine learning, pp. 1861-1870. PMLR, 2018.
Bellemare, M. G., Dabney, W., and Munos, R. A distribu- tional perspective on reinforcement learning. In Interna- tional Conference on Machine Learning, pp. 458. PMLR, 2017.
10

---

Wasserstein Policy Optimization
Haarnoja, T., Moran, B., Lever, G., Huang, S. H., Tirumala, D., Humplik, J., Wulfmeier, M., Tunyasuvunakool, S., Siegel, N. Y., Hafner, R., et al. Learning agile soccer skills for bipedal robot with deep reinforcement learn- ing. Science Robotics, :8))eadi802,, 2024.
Metelli, A. M., Likmeta, A., and Restelli, M. Propagat- ing uncertainty in reinforcement learning via wasserstein barycenters. In Advances in Neural Information Process- ing Systems, volume 32. Curran Associates, Inc., 2019.
Moskovitz, T., Arbel, M., Huszar, F., and Gretton, A. Ef- ficient Wasserstein natural gradients for reinforcement learning. arXiv preprint arXiv:2010.05380 2020.
He, S., Jiang, Y. Zhang, H., Shao, J., and Ji, X. Wasser- stein unsupervised reinforcement learning, 2021. URL https: //arxiv...vvv///a///a/eaaaaaiav org/ abs/ 2110.1111100 07940.
Neklyudov, K., Nys, J., Thiede, L.,, Carrasquilla, J., Liu, Q., Welling, M., and Makhzani, A. Wasserstein quantum Monte Carlo: a novel approach for solving the quantum many-body schrödinger equation. Advances in Neural Information Processing Systems, 36, 2023.
Heess, N., Wayne, G., Silver, D., Lillicrap, T., Erez, T., and Tassa, Y. Learning continuous control policies by stochastic value gradients. Advances in neural informa- tion processing systems, 28, 2015.
Pacchiano, A., Parker-Holder, J., Tang, Y., Choromanski, K., Choromanska, A., and Jordan, M. Learning to score behaviors for guided policy optimization. In Interna- tional Conference on Machine Learning, pp. 7445-7454. PMLR, 2020.
Hoffman, M. W., Shahriari, B., Aslanides, J., Barth-Maron,
G., Momchev, N., Sinopalnikov, D., Starczyk, P., Ramos, S., Raichuk, A., Vincent, D., al. Acme: A research framework for distributed reinforcement learning. arXiv preprint ar arXiv:2006.009 2020.
Pascanu, R. and Bengio, Y. Revisiting natural gradient for deep networks, arXiv preprint arX1v:1301.3584, 2013.
Howard, R. A. Dynamic Programming and Markov Pro- cesses. 1960.
Prokhorov, D. and Wunsch, D. C. Adaptive critic designs. IEEE transactions on Neural Networks, 8(5):997-1007, 1997.
Huang, S., Dossa, R. F. J., Raffin, A., Kanervisto, A., and Wang, W. The 37 implementation details of proximal policy optimization. The ICLR Blog Track 2023, 2022.
Richemond, P. H. and Maginnis, B. Diffusing poli- cies Towards Wasserstein policy gradient flows, 2018. URL id Id-rk3mjYRp--
Jang, E., Gu, S., and Poole, B. Categorical reparameteriza- tion with Gumbel-softuax. International Conference on Learning Representations, 2017.
forrm?
Kakade, S. M. A natural policy gradient. Advances in neural information processing systems, 14, 2001.
Schulman, J., Levine, S., Moritz, P., Jordan, M., and Abbeel, P. Trust region policy optimization. International Con- ference on Machine Learning, 2015a.
Likmeta, A., Sacco, M., Metelli, A. M., and Restelli, M. Wasserstein actor-critic: Directed exploration via optimism for continuous-actions control, 2023. URL https:: //ariiv///xvv...// httg//bs/23300 02378.
Schulman, J., Moritz, P., Levine. S., Jordan, M., and Abbeel, P. High-dimensional continuous control using generalized advantage estimation. arXiv preprint 150600550 2015b.
Lillicrap, T. Continuous control with deep reinforcement learning. arXiv preprini (5t::::500009777 2015.
Schulman, J., Wolski, F., Dhariwal, P., Radford, A., and
Klimov, O. Proximal policy optimization algorithms. arXiv preprint arXiv: 2017.
Maddison, C. J., Mnih, A., and Teh, Y. W. The Concrete Distribution: A Continuous Relaxation of Discrete Ran- dom Variables. International Conference on Learning Representations, 2017.
Silver, D.,, Lever, G., Heess, N., Degris, T., Wierstra, D., and Riedmiller, M. Deterministic policy gradient algorithms. In International Conference on Machine Learning, pp. 387-395. PMLR, 2014.
Martens, J. and Grosse, R. Optimizing neural networks with Kronecker-factored approximate curvature. In Interna- tional Conference on Machine Learning, pp. PMLR, 2015.
Song, H. F., Abdolmaleki. A., Springenberg, J. T.,. Clark, A., Soyer, H., Rae, J. W., Noury, S., Ahuja, A., Liu, S., Tirumala, D., et al. V-MPO: On-policy maximum a posteriori policy optimization for discrete and continuous control. arXiv preprint ArXiv:1909.1220 12238, 2019.
Messaoud, S., Mokeddem, B., Xue, Z., Pang, L., An, B., Chen, H., and Chawla, S. AC: Energy-Base Rein- forcement Learning with Stein Soft Actor Critic. Interna- tional Conference on Learning Representations, 2024.
Sutton, R. S. and Barto, A. G. Reinforcement learning: An introduction. MIT press, 2018.

---

Wasserstein Policy Optimization
Sutton, R. S., McAllester, D., Singh, S., and Mansour, Y. Policy gradient methods for reinforcement learning with function approximation. Advances in Neural Information Processing Systems, 12, 1999.
Tassa, Y., Doron, Y., Muldal, A., Erez, T., Li, Y., Casas, D. d. I.., Budden, D., Abdolmaleki, A., Merel, J., Lefrancq. A., al. Deepmind control suite. arXiv preprint 18X1v0/801.00690, 2018.
Todorov, E., Erez, and Tassa, Y. Mujoco: A physics engine for model-based control. In 2012 IEEE/RSJ inter- national conference intelligent robots and systems, pp. 0026-5033. IEEE, 2012.
Tracey, B. D., Michi, A., Chervonyi, Y., Davies, Padu- raru, C., Lazic, N., Felici, F,, Ewalds, T., Donner, C., Galperti, C., et al. Towards practical reinforcement learn- ing for tokamak magnetic control. Fusion Engineering and Design, 200: 14161, 2024.
Tunyasuvunakool, S., Muldal, A., Doron, Y., Liu, S., Bohez, S., Merel, J., Erez, T., Lillicrap, T., Heess, N., and Tassa, Y. dm control: Software and tasks for continuous control. Software Impacts, 6:100022, 2020.
van Hasselt, H. Double qelearning In Lafferty, J., Williams, C., Shawe-Taylor, J., Zemel, R., and Culotta, A. (eds.), Advances in Neural Information Processing Systems, vol- ume 23, Curran Associates, Inc., 2010.
van Hasselt, H. and Wiering, M. Reinforcement learning in continuous action spaces. In 2007 IEEE International Symposium Approximate Dynamic Programming and Reinforcement Learning, pp. 272-279. IEEE, 2007.
Werbos, Beyond regression: New tools for prediction and analysis in the behavioral sciences. PhD thesis, Commit- tee on Applied Mathematics, Harvard University, Cam- bridge, MA, 1974.
Williams, J. Simple statistical gradient-following algo- rithms for connectionist reinforcement learning. Machine learning, 8:229-25 1992.
Zhang, R., Chen, C., Li, C., and Carin, Policy optimiza- tion as Wasserstein gradient flows. In International Con- ference on Machine Learning, pp. 5737-5746. PMLR, 2018.
12

---

Wasserstein Policy Optimization
A. Extended Derivations
A.1. Functional Derivative of the Loss
To prove that WPO a true policy gradient method, and not simply a policy improvement method, we need to prove that the functional derivative of the policy return with respect to the policy has the form in Eq. 4. While this is known result the tabular case (see Agarwal et al. (2021), Eq. 7), we give a derivation in the case of continuous action and state spaces for completeness here. Suppose we perturb policy x (a s) by an amount óx (al (s)dt, such that f daō (als) 0's. Then in the dt 0 limit the change to the expected long-term reward is
lim de-+0
JIh + ôTdi)
[[T = [ =0
f darls dn
(s, (s)a))
dt
to solve for
551=] 8n
we plug in this perturbation into the definition of the objective function:
T
J = E,
T'm
(10)
1=0
7" f dreatiP(su)
Σ t=
II K==0
x maalss (sv)
II F=00
(11)
f ddrasrio (To:))
(12)
where т is shorthand for full-length trajectories So: Sr.ar. To: is shorthand for truncated trajectories So, ao,..., St, a, and p" (Toze) denotes the probability of trajectory, combining environment transitions and policy steps.
Plugging in π + π + ōrdt into Eq 11 and expanding, the perturbed value is then:
J[T + Srdt]]
[J
lim T
[[T]]
(13)
dt
t-1
8n (ar (yr)
dToare
(sp))
II
x (ar (s))
II d'=0
P(s (p)
(14)
π [ay (sc)
6'=0
8'=0
/ droaTeW
8t |ar [Se)
Σ
(15)
(ay (8e)
0
F==0
x' [ ddooeri"" (Tax)
8m (ar/ (sw) (ar) T T (sv)
f dretirpp (Te))
Σ
(16)
{ ddrip (T)
8T (sr)
(17)
π T(ar((ss)
Ó (sy) drp" (T)
Σ
y'ra
(18)
T(a (.
V=O"
t=t'
T
= Σ
droa
n ca (sp))
Q" (se, as))
(19)
r=0
(als)
yd (s: =s)
(20)
P=2
JJ
1
Q" a)
Pr (s, s)
Q" as)
(21)
1 1-y
13

---

Wasserstein Policy Optimization
where (s s) the marginal probability of the state time / under the policy x and d (1-7) E, P'Pr =
is the discounted state occupancy function. Thus the action-value function Q" is almost the functional derivative with an additional discounted occupancy over the states.
65
but
8m
A.2. Projection Onto a Parametric Function Space
Here we give a more complete version of the derivation in Sec. 2.3
@Tp(a da a
F19-
V Vologro(als)
(22)
v Velogmo (
(23)
Vologmo(als)
da
(24)
Vorglas)V: Valomm(assaammmasaaa
(25)
V (a (as
Vono(als Valerr(BB5..m/m/33350 a)da (26)
V Vo Vs."
ajda
(27)
V v a To a
da
(28)
= E [ [T
a =EA- V V,V,
(29)
(s)
(30)
where the step Eq. 26 comes from integration by parts, and we assume that Va (a|s) goes goes to infinity, The last step can be derived in reverse by expanding ValogTo To Vana and applying the product rule.
A.3. Variance Update for SVG(O) WPO
AO
Elaa - VoVa())
Veo
[ [(s,0)((-
Voo
VBo /. - Q(s..)
Vag (x((s,,)) V) +1]]
V Ex [Qss a) [- (a n)o 1)]
Q(s, a) (a
(0-1))(-10) (-o + +0)
log
E,(Q(ss a) V, log (s,, Vooj
We note that the last equation is the standard expected policy gradient with respect to the gradients flowing through the variance term. (If we are updating both mean and variance, as typically, we just add the contributions with respect to the mean get the full policy gradient update.)
14

---

Wasserstein Policy Optimization
Algorithm WPO with Replay and n-step TD critic learning for multi dimensional Gaussian Policy
Require: Initialize actor aa s), critic (s, a), target actor Ty (als), and target critic Qa (s, a) with parameters w. 0 <- 0, and vi +- respectively. Initialize replay buffer D)
1: for each episode do
2: Initialize state 80
3: for = 0 to T do
4: Select action a, ~ To(a)s;) (Ng Ess))
5: Execute action a, and observe reward r, and next state 8:4-
6:
Store transition (st, ae, Tt. S(+1) in replay buffer D
7:
if length(D)) 2 batch_size then
8:
Sample mini-batch of transitions (sj, agj T.. 81+1) from D
9:
Sample n-step transitions (8j+*, aj+k, Tj+*) ) for k: = l,..., n or until the end of the episode is reached.
10:
Calculate n-step TD target (using target critic):
Gj:ttn = Σbn ? ryth + N >, Qu (8)+n: dn+,,) (or greedy)
where j'nn, T8 ( -N -N(ug(8)ta) (8()))
12:
Update critic:
13:
V (Qw a,) Gj:t+n²²
14:
Update actor (with action sampling):
15:
for i = to N (Number of action samples) do
16:
Sample action a' To(alss) . (sy) E(s;))
17:
Calculate policy gradient for the sampled action:
9: Vo log To (a 8,) Qw a.)) ) WPO Update
where V = EV, and V EV
19:
end for
Update actor using the average gradient:
21:
0 E 9 [VoDkulTellTa)
22:
Update target networks: if mod PERIOD = 0 then
23:
Update target networks (hard update): O- w;
25:
26:
end if
27:
end if
28:
t 81+1
29:
end for
30:
end for
15

---

Wasserstein Policy Optimization
B. Experiment Hyperparameters
All algorithms, WPO, MPO, DDPG and SAC are implemented the same codebase.
B.1. Deep.Mind Control Suite
We use a single set of hyperparameters for each learning algorithm across all control suite environments.
use a longstanding implementation MPO which accordingly has well optimised hyperparameters. For WPO, DDPG and SAC, we sought hyperparmeters which worked well across all control suite tasks. Drawing some values from the wider literature continuous control, we performed hyperparemeter sweep. We arrived parameters which performed well across the suite of tasks.
Our focus is on the comparison of algorithms rather than the search for the best hyperparameters.
We note that improved performance can be attained though environment-specinn hyperparameter tuning. We found this to be particularly true for Soft-Actor Critic. Maintaining hyperparameter settings across environments helps us investigate variation due learning algorithm differences.
Table Common Hyperparameters
Hyperparameter
Value
Actor Network Hidden Layer Sizes
(256, 256, 128)
Critic Network Hidden Layer Sizes
(512, 512, 256)
Optimizer
ADAM
Actor Learning Rate
3 x 10-4
Critic Learning Rate
3 x 10-4
Activation Function
ELU (Clevert et al., 2016)
Discount Factor (?)
0.99
Target Network Update Period
Every 100 updates
Samples-per-insert (SPI)
Batch Size
256
(for TD(n) loss calculation)
(Except SAC, which used n 3)
Replay Buffer Size
2 x 106
Trajectory Length for Network Update
10
Table WPO Hyperparameters
Hyperparameter
Value
Table DDPG Hyperparameters
Hyperparameter
Value
Oa
log(2)
10000
(Fixed) Policy Variance (a)
0.3
Actions Sampled per Update
30
Table 4. MPO Hyperparameterss
Table 5. SAC Hyperparameterrs
Hyperparameter
Value
Hyperparameter
Value
ag (initial)
100
Initial a
0.0001
oz (initial)
1000
Minimum a
1 x 10-8
Initial Log Temperature (log
10
Final Policy Network Activation Function
tanh
0.1
Polyak Average Coefficient (T)
0.005
5 x 10-3
Target Entropy
-|A]
E
10-6
Maximum Policy Variance
exp(4)
Actions Sampled per Update
30
Minimum Policy Variance
exp(-10)
16

---

Wasserstein Policy Optimization
Acrobot Swingup
Acrobot Swingup sporse
Ball in cup Catch
Caripole Ba ance
Cartpole Balance sparse
Cartpole Swingup
1,000
Eoisodic Episedic
Cartpole
Cartpole Tir
Cartpole Three
Cheetah Run
Dog Walk
Swingup
Stand
1000
Finger Spin
Finger easy
Finger
Fetch
Episadic Retuin
Hopper Stand
Hopper Hop
Humanoid
Humanoid Walk
Upright
Episodic Return
1000
500
Humnnnaid
Humanoid Hun pure state
Humanold CMU Stand
Humanoid CMU Walk
Humannid CMU Run
Manipulator Bring ball
Run
1000
Episodic
500
Man pulator Bring peg
Manipulator Insert ball
Manipulator Insert peg
Pendulum Swingup
Point mass Easy
Point mass Hard
1000
Quadruped Walk
Quadruped Run
Quadruced Escape
Quadruped Fetch
Reacher Easy
Reacher Hard
1000
Stacker
Stackar Stack
Swimmer Swimmers
Swimmer Swimmer15
like Stand
Walker Walk
1000
Episedic
Episodic
500
WPO
DDPG
MPO
SAC
Actor Steps
Figure Results on all DeepMind Control Suite tasks for WPO, MPO, DDPG and SAC. The bold lines denote average returns over 10 seeds on evaluation episodes. The shaded region spans minimum and maximum returns.
17

---

Wasserstein Policy Optimization
Figure 8. Schematic of the TCV tokamak, including cutout showing 2-D profile of plasma contour
B.2. Combined Tasks
In order to combine the rewards across component environments in the combined tasks we apply a SmoothMin operation. This is the application of SmoothMax operation with negative value of a.
Σ i= 22
SmoothMax(21.ooo a)
(31)
ii=l ea0)
The SmoothMin operation is non-linear. To make the hyperparameter easier to tune and maintain across differing choices of component environment, we scale the rewards such that they are in the unit interval by dividing by the highest reward observed from any agent in a singleton environment.
B.3. Fusion Task
Fig 8 depicts the TCV tokamak, including the control coils, vacuum-vesst cross section and a notional plasma inside the domain. The TCV tokamak has 19 control coils, composed of 16 poloidal field coils, ohmic coils, and "fast" coil. The measurements observed by the agent for these experiments consist of 38 magnetic field probes, 34 magnetic flux loops, and 19 coil current measurements, one for each loop. We simulate the evolution of the tokamak using the FGE Trad-Schafranov simulator (Carpanese, 2021). For more details of the training setup, please (Degrave al., 2022) and (Tracey et al., 2024).
The task shown in Section 6.2 is based on the shape shape_701- 7016606666 65 task from (Tracey et al., 2024). The primary goal of the task is for the plasma last closed flux surface to align with the control points, and also for the current in the plasma to be maintained at a reference value (not depicted). There are some additional reward terms to regularize the policy and help hardware transfer, specifically penalizing out-of-bounds voltage commands, keeping the ohmic coil currents close together, and keeping the poloidal coil currents away from zero. The results shown in Fig. 9 use a version of the task with tighter demands on the shape and plasma current accuracy. These changes are similar to those discussed by Tracey et (2024) with regards to reward shaping. We also show in Fig. 10 results from the original task in (Tracey et al., 2024). The specifics of the reward components and combiner is shown in Table 6, please see Sec. 3 of (Tracey et al., 2024) and in particular Eq.2 and Eq. 3 for how these terms combine into a scalar reward. In both figures, MPO and WPO are run with the same set of hyperparameters, must notably with a hard constraint on the KL regularization with c,. = 5e and 6, = le - 7. Note that the KL constraint on the variance, 6g. is set to the same value for WPO and MPO. So while the value is tighter than typically seen, not responsible for the difference in policy convergence rates between the two algorithms.
18

---

Wasserstein Policy Optimization
References
Flux
Figure 9. The reference control points and example flux field for the shape_701 66 task. The last elosed flux surface target points are shown in blue the left, with the limit point location in cyan. An example magnetic flux field matching these control points in shown on the left, with the limit location magenta and the LCFS in black
0.0
Average
0.18
Retur
0.37
035
20
WPO (Itard KL Constraint) MPO (Hard KL Constraint)
034
0.02
6.25
3.50
0.75
1.00
1.25
1.53
3.75
2.20
0.32
0.25
5.50
0.75
1.00
3.25
1.52
1.75
7.50
Actor Steps
1ed
Actor Steps
1st
Figure Return and policy average standard standard-deviation evolution throughout training on the shape_70166 task from (Tracey al., 2024)
C. Additional Results and Ablations
C.1. Complete Deep Control Suite Results
In Fig. 7, we show results on all 49 Control Suite tasks (excluding the LQR domain, which suffered from numerical issues in simulation) for WPO and baseline algorithms, providing a more complete picture of the results in Fig 4. It can be seen that WPO is competitive with MPO on nearly all tasks, while there are numerous tasks for which DDPG and SAC struggle, especially those on the Cartpole and Point mass domains for DDPG and Acrobot and Humanoid tasks for SAC. Notably, even on those tasks where DDPG outperforms other methods, WPO still converges eventually, while many tasks for with WPO performs well, DDPG does not reliably converge. We note that the performance of WPO on some domains is sensitive to modeling choices. We explore these choices below.
C.2. MPO Hyperparameter Comparison
To ensure that we are making fair comparison between WPO and the strongest possible baseline on DeepMind Control Suite, we investigated several different hyperparameter settings for MPO, with results shown in Fig. C.2. In addition to the hyperparameters in Table 4, we also ran MPO with the same hyperparameters as WPO - reelacing the hard KL constraint with a soft constraint and relaxing the KL regularization on the mean. We additionally tried adding the variance rescaling trick described in Sec. As can be seen in the figure, these modifications did not significantly affect the performance of
19

---

Wasserstein Policy Optimization
Acrobot
Carppole Balance
Cartpole sparse
Cartpole Swingup
Swingup
Swingup sparse
1000
Reward
500
Cartpole Swingup sparse
Cartpole Two poles
Cartpole Three poles
Cheetah Run
Dog Stand
Dog Walk
1000
Reward
500
Dog
Dog Run
Dog Fetch
Finger Spin
Finger Turn easy
Finger Turn hard
Trot
1000
Reward
500
Rish
Rish Swim
Hopper Stand
Hopper Hop
Humanold Stand
Humanold Walk
Upright
1000
Reward
Humanoid
Humanaid
Humanaid CMU Stand
Humanold Walk
Humanoid CMU
Manipulator ball
1000
Reward
Manipulator Bring peg
Manipulator ball
Point
peg
Swingup
Easy
Reward
500
Quadruged
Quadruped Run
Quadruped Escape
Quadruped Fetch
Reacher Easy
Reacher
Walk
1000
Reward
500
Stacker
Stacker Stack 4
Swimmer Swimmer6
Swimmer Swimmer15
Walker Stand
Walloer Walk
Stack
1000
Reward
500
Walker
Run
1000
Reward
MPO best hypers
MPO WPO hypers
MPO WPO hypers, variance rescaling
0
2
Actor Steps 1e7
Figure 11. Results on all DeepMind Control Suite tasks for MPO with different hyporparameters: optimal settings tuned for Control Suite, settings matching WPO, settings matching WPO with variance rescaling. The performance MPO not significantly impacted by these choices.
20

---

Wasserstein Policy Optimization
Component
Precise
Standard
Good
Bad
Good
Bad
LCFS Distance (m)
0.001
0.01
0.005
0.05
Plasma Current (A)
100
2000
500
30000
Poloidal away from 0 (A)
100
50
100
50
Normalized Voltages in Bounds
0
1
0
Ohmic Coils Close (A)
50
1050
50
1050
Combiner
-0.5
Table 6, Fusion task reward components and scaling parameters
MPO. With the WPO hyperparameters, the Humanoid CMU and Hopper tasks were slightly slower to learn for some tasks, but adding the variance rescaling largely restored the performance. This gives us high confidence that our baseline is as strong as possible, and that any difference in performance cannot attributed to arbitrary hyperparameter differences.
C.3. Squashing Functions and Network Activations for WPO
There are various possible sources numerical instability unique to WPO which we investigate. One that, for certain tasks where the action-value function changes rapidly, the gradient VaQ* (s, a) might blow up, making the gradients unstable. Fortunately, there is fix for this which still principled approach approximating Wasserstein gradient flows. In Neklyudov et al. (2023), was shown that much of the theory in Sec. 2.2 and 2.3 can be extended to the c-Wassserstein distance:
/ p(a, b)c(a - b)dadb
2 ( (1)
inf pEI()o.)))
(32)
where e is some convex function. In this case, the appropriate PDE for the c-Wasserstein gradient flow to minimize J becomes
dm 0t
= Va TVc" (-Va
5J 8T
(33)
where o* the convex conjugate This means that applying a nonlinear squashing function the velocity field can still result in dynamics that minimize the functional of interest.
If we follow the same derivation of the parametric update in Sec. 2.3 and A.2 for the c-Wasserstein distance, the derivation is largely unchanged, but the term VaQ² a) is replaced with Vc (V.Q" This means applying nonlinear function to the gradient of the action-value function still results in a principled update, which may be useful for numerical stability in cases where the action-valle function changes rapidly. To investigate the impact that this has on the performance of WPO on the DeepMind Control Suite, we choose Ve" (a) 4/3, where the cube root applied elementwise. This a natural choice of squashing function it smooth, odd, and doesn't saturate.
Another possible source of numerical instability comes from the choice of neural network activation. The exponential linear unit (ELU) (Clevert et al., 201 6) has a discontinuity in the second derivative, which may cause issues with WPO as the second derivative appears in the update. We try changing both the actor and critic networks to use sigmoid linear units (SiLU) (Elfwing et al., 2018). We present the results of both changing the network nonlinearity and adding a cube root squashing function to the action value gradient in Fig. 12
On most tasks, the choice of nonlinearity and squashing function does not make significant difference. On a number of tasks such as Cartpole Swingup sparse and tasks on the Dog domain, both changing the nonlinearity and adding squashing function seem improve performance, but the Humanoid CMU domain they potentially harm performance. Therefore, it's difficult to conclude that any of these choices decisively improve performance, but it is important to be aware that some of the results presented in the main paper can be sensitive to these choices.
21

---

Wasserstein Policy Optimization
Acrobet Swingup
Acrahot Swingup
Rall cup Catch
Garipole Ba
Cartpole
Balance sparse
Swingup
Cartpole Swingup sperse
Cartpole poles
Cartpoll poles
Dog
Reward
500
Dog
Dog Run
Dog Fetch
Finger Spin
Finger Turn easy
Finger hard
Trot
1000
Rewerd
500
Rish
Fish Swim
Hopper Stand
Hopper Hop
Humaneld Stand
Humanoid Walk
Upright
1000
500
Humanoid Run
Humanoid Run pure state
Humanoid CMU Stand
Humanoid CMU Walk
Humanald CMU Run
Manipulator Bring ball
1000
Manipule punipultoo
Manipulator Insert bell
Manipulator Insert
Pendulum Swingup
Point mass
Point Hard
Bring peg
1000
Quadruped
Quadruped
Quadruned Escape
Quadruped
Easy
Stacker
Stackar Stack
Swimmer Swimmers
Swimmer Swimmer15
We liee
Walker Walk
1000
Reward
500
Walkeer
Run
1000
Cube root, SiLU Cube root, ELU
500
squashing, SiLU
squashing, ELU
Actor Steps le7
Figure 12, Results on all DeepMind Control Suite tasks for WPO with different activation functions (ELU vs. SiLU) and with and without applying cube root clementwise the aarion-valus gradients. The bold lines denote average returns over 3 secds on evaluation episodes. The shaded region spans minimum and maximum returns.
22


Reviews:
[{"review_id": "hv3ZKMhSfR", "reviewer": "Reviewer_A8Ro", "summary": "The paper introduces a novel policy gradient update using Wasserstein Gradient Flows called Wasserstein Policy Optimization. While classical policy gradient update works with stochastic policy, it does not take gradient through the action-value space, while deterministic policy gradients are able to take gradients through the action value function but are designed for deterministic policies. Using the notion of the vector flows derived from Wasserstein-2 distance and second order Fisher approximation of the KL divergence, the authors derive a policy gradient update that's best of both worlds. They derive a tractable practical algorithm for the special case of Gaussian policies.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your helpful comments. We appreciate the overall positive evaluation of the paper, and wanted to address the two points raised in “Claims and Evidence”:\n\n1. To go from the idealized WPO update to the practical algorithm presented in the paper, two approximations were made: the Fisher information matrix was approximated by the FIM for a diagonal Gaussian distribution, and an entropy regularization term was added.\n\nFirst, there is no reason that the “bare” WPO update (i.e. without the Fisher term) couldn’t be used for general distributions over actions. As noted in the paper, the variance in the update would grow as the policy becomes more deterministic, but with gradient clipping this might be manageable. Beyond that, there are many classes of distributions for which the exact Fisher information matrix is known (e.g. exponential families) and, going even beyond that, numerous practical approximations to natural gradient descent exist which could be applied here, e.g. KFAC (Grosse and Martens 2015). Ultimately, there are going to be a spectrum of approaches with the exact update at one end and the “bare” update at the other which can be explored in future work.\n\nSecond, regularization of the policy update through penalties or trust region methods is a standard technique, and the approach used here could be extended to more complex distributions over action spaces so long as an appropriate way to estimate the KL divergence between action distributions exists. For instance, a sampling-based approximation could be used for generic distributions, or an upper bound or other approximation to the KL could be used in cases where a closed form solution does not exist. In practice, the exact measure of divergence between action distributions is probably not that important, as long as it is simple to compute and is well behaved.\n\n2. Understanding the precise cause of WPO underperforming other algorithms in certain environments is challenging. Note that we only tried tuning a small number of parameters, primarily the weights of the KL regularization, while other methods have benefited from years of extensive hyperparameter tuning on many of these environments to achieve optimal performance. Note also that, while WPO is not always the top performing algorithm, it does not fail dramatically on any environments in the way that, say, DDPG does on Cartpole - Swingup Sparse, which illustrates the overall robustness of WPO. We agree that it would be quite valuable to investigate the cases where WPO underperforms in more detail, but given the large number of moving parts and possible causes, we feel this is best left to future work. Furthermore, while Figs. 3 and 6 show that WPO is competitive over the DM control suite, our main performance claims are highlighted in Figs. 4 and 5 that show that, as conjectured, WPO especially shines in domains with relatively high-dimensional action spaces.\n\nWe also appreciate the suggestion to compare against Zhang et al (2018). The algorithm implemented in Zhang (2018) is a combination of a DDPG-like gradient with the reparameterization trick, and thus closely resembles SVG(0) and SAC, albeit with a slightly different regularization. The significance of WPO lies mainly in its generality, going beyond methods like SVG(0), SAC and Zhang (2018) that rely on reparameterization. Additionally, we believe that by including SAC and DDPG in the experiments already, we have good coverage of methods with the same basic form of the policy gradient update. If we investigate the effect of regularization in more detail in future work we will be sure to include an analysis of Zhang (2018). Once again, we are pleased that you liked the paper, and appreciate all of your helpful comments.", "claims_and_evidence": "The authors claim to use Wasserstein gradient flows to derive policy gradient updates that have best of both worlds: use stochastic policies and gradients through the action value function. They derive the gradient update that has this property. But what concerning are:\n\n(1) Its not always simple to derive a practical algorithm for the gradient update. For a special case of policy distribution: normal, the authors derived a tractable approximate algorithm for these updates but how generalizable is this?\n\n(2) While WPO performs robustly well, there are environments where WPO fails. It would be interesting to see why and in what situations WPO fails.", "methods_and_evaluation_criteria": "The proposed method and evaluation criteria though make sense. Though the authors didn't compare against prior Wasserstein based policy gradient methods like [1]. Comparisons against these methods will justify if their gradient updates are the ideal way of using Wasserstein Gradient Flows. \n\n[1]: Zhang, Ruiyi, et al. \"Policy optimization as wasserstein gradient flows.\" International Conference on machine learning. PMLR, 2018.", "theoretical_claims": "The theoretical claims made are justified and proven.", "experimental_designs_or_analyses": "(1) Some baselines like the one mentioned above (prior attempts to use Wasserstein Gradient Flows) should be compared against.", "relation_to_broader_scientific_literature": "Policy gradient methods updates are used in a majority of RL algorithms. There have been two commonly used updates: classical and DPG. This work proposes a novel policy gradient update rule using Wasserstein gradient flows to have best of both worlds. But it is not clear if this method scales to other policy distributions (i.e. practical algorithms designed for them).", "essential_references_not_discussed": "To my knowledge most of the related work has been referred to.", "other_strengths_and_weaknesses": "Nil", "other_comments_or_suggestions": "Nil", "questions_for_authors": "Please address the concerns raised.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Qw17ll1nla", "reviewer": "Reviewer_mryg", "summary": "This paper introduces Wasserstein Policy Optimization (WPO) for continuous-action reinforcement learning. By viewing policy optimization as a Wasserstein gradient flow in the space of distributions, the authors derive a closed-form update that: i) Uses gradients of action values w.r.t. actions (like deterministic policy gradients). ii) Works for any stochastic policy (no reparameterization needed, unlike SVG or SAC). iii) Can be turned into a simple actor-critic algorithm with minimal additional structure. The paper provides theoretical derivations, then extends WPO with practical features (e.g., KL regularization, a Fisher information rescaling) for deep RL. Experiments on the DeepMind Control Suite (including high-dimensional tasks like Humanoid) and a real-world-inspired fusion-control task show that WPO is robust, sometimes outperforming or matching strong baselines (MPO, SAC, DDPG). The authors highlight WPO’s favorable scaling in large action spaces and its ability to converge stably to high-performing solutions.\n\n## update after rebuttal\n\nI thank the authors for their detailed response and additional experimental results. I continue to find the theoretical formulation of Wasserstein Policy Optimization novel and compelling. I also appreciate the effort made to re-run experiments with more seeds and to clarify the challenges involved in evaluating on larger benchmarks.\n\nThat said, some of my original concerns remain. The updated results indicate noticeable variance across seeds in several tasks, and some performance drops (e.g., Humanoid CMU Walk, Manipulator: Bring Ball) raise questions about the method’s robustness. Additionally, while I understand the limitations during the rebuttal phase, the lack of evaluation on more realistic high-dimensional environments makes it difficult to fully assess the claimed scalability of WPO.\n\nOverall, I believe the paper presents an interesting direction, and I hope to see future versions strengthen the empirical foundation. I am maintaining my score.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer for their thoughtful comments and favorable review. The reviewers made a few suggestions for improving the paper which we will address:\n1. We have re-run experiments on the high-variance environments with more seeds to smooth out the learning curves and will add the new results to the paper. The new results look qualitatively similar to those in the original submission, with some amount of noise.\n2. We appreciate the suggestion for more challenging high-dimensional tasks to investigate such as Bi-DexHands. The state-of-the-art RL methods for bidextrous manipulation often involve many additional steps beyond blank-slate actor-critic methods, such as imitation learning from human data (e.g. ALOHA Unleashed, https://arxiv.org/pdf/2410.13126). This adds an additional layer of complexity for comparison against baselines that goes beyond the scope of what we have the space and time to investigate in this paper, however we will be sure to investigate these sorts of tasks going forward.\n\nWith regards to the lack of uniformly superior performance that the reviewer points out, we refer to the rebuttal to reviewer A8Ro for a longer discussion. We hope that this addresses all of the issues the reviewer has and thank them again.", "claims_and_evidence": "The paper’s main claims—i.e., that WPO (1) combines properties of deterministic and standard policy gradients, (2) applies to arbitrary continuous action distributions, and (3) matches or outperforms strong baselines in large action spaces—are backed by both theoretical derivations and empirical results on benchmark (DeepMind Control Suite) and real-world-like (fusion) tasks. The paper does not appear to make unsupported or overstated claims", "methods_and_evaluation_criteria": "The paper evaluates WPO in three ways. (1) the DeepMind Control Suite, a widely used benchmark for continuous control, enabling clear comparisons with established methods. (2) a Magnetic Confinement Fusion domain, which tests real-world relevance and benchmarks WPO against strong baselines like MPO. (3) Combined Tasks, The paper combines multiple copies of a Control Suite environment to form high-dimensional action spaces, revealing WPO’s potential for scaling to large numbers of actions. These three perspectives offer a well-rounded assessment of WPO’s capabilities.", "theoretical_claims": "There are no apparent flaws in the theoretical arguments as presented.", "experimental_designs_or_analyses": "**Experiments in DeepMind Control Suite.** From Figure 3, we can see that there is a high variance environments, e.g. Cartpoles, Don Run, Manipulator Bring ball and Humanoid Stand. It should run more seeds to reduce the variance.\n\n**Experiments in Combinded Tasks.** By replicating a single environment (e.g., Humanoid - Stand) multiple times to expand the action space, the authors leverage SmoothMin to aggregate rewards. This approach effectively tests scalability to higher-dimensional controls, but it still falls short of simulating the diversity and complexity of real-world scenarios. It would be good if test WPO in Bi-DexHands tasks with high-dimensional control.", "relation_to_broader_scientific_literature": "The authors propose Wasserstein Policy Optimization (WPO), presenting it as an extension of classic gradient-based methods (e.g., REINFORCE, actor-critic) by exploiting action-value gradients similarly to DPG, yet without restricting the policy class to deterministic forms. Grounded in optimal transport, WPO’s update reflects a continuous steepest descent in policy space, differing from approaches like MPO that exponentiate Q-values rather than relying on a direct gradient. This formulation aligns stochastic and deterministic policy gradients, offering a unified viewpoint and potentially stronger scalability in high-dimensional action spaces.", "essential_references_not_discussed": "No", "other_strengths_and_weaknesses": "**Strengths:**\n\n- Principled Theoretical Foundation\n- Combines Advantages of Deterministic and Stochastic Policy Gradients: This allows WPO to leverage value information for efficient updates without the limitations of deterministic policies or the reparameterization trick.\n\n**Weaknesses:**\n\n- High variance in training: From Figure 4, WPO has higher variance compared with MPO and DDPG.\n- Not Universally Superior: While WPO performs competitively on many tasks, it does not consistently outperform state-of-the-art algorithms like SAC and MPO across all environments in the DeepMind Control Suite", "other_comments_or_suggestions": "Do more experiments to show WPO’s validness.", "questions_for_authors": "WPO combines the advantages of deterministic and stochastic policy gradients, but why it is not universally superior than the baselines?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "SrLRSmRjrF", "reviewer": "Reviewer_L61D", "summary": "This paper naturally and clearly derives a natural gradient version of the stochastic policy extension of deterministic policy gradient from the perspective of Wasserstein gradient flow, proposes a practical implementation method, and conducts detailed tests on DeepMind Control Suite and scientific control tasks, achieving comparable results.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the comments and helpful suggestions.\n1. With regards to the state-occupancy term $d^\\pi(s)$ in the definition of the gradient of the value function, we agree that including a correction in the off-policy case could potentially improve performance, but want to emphasize that virtually no methods in off-policy RL actually use a correction of this form in practice. Therefore, our algorithm conforms to standard practice and we are confident that the experimental comparisons in the paper are fair. One reason that this term is almost always excluded is that estimating it in continuous state spaces is extremely challenging. It’s also worth highlighting that it has been argued in previous work that excluding this term actually can improve the efficiency of the updates, and it has been shown to be theoretically a sound thing to do (see, e.g., https://arxiv.org/pdf/2302.11381). We will discuss this explicitly in the paper. \n2. We have fixed the typo in the equation on line 685\n3. Unfortunately, we do not control the rights to the fusion simulator used in our experiments, and so have no control over open-sourcing it. However, the organization which does own the rights intends to release an open-source version of the simulation code in the long term. We will do our best to encourage them to release this as soon as they can.\n4. The use of the reparameterization trick becomes much more complicated with distributions with discrete latent variables, like mixtures of Gaussians. The exact SVG(0) update, and the efficiency of this, will be sensitive to the choice of reparameterization. It is difficult to find a practical problem for which mixture-of-Gaussian action distributions are necessary, since Gaussian action distributions work well in practice for most environments investigated here. However, we will extend Section 4 to include a toy example of a mixture of Gaussians policy on a complex loss landscape showing the qualitative difference in dynamics between WPO and standard policy gradient, and Gaussian vs. mixture-of-Gaussian action distributions.\n5. The discussion of the variance of the linear action-value function case in Section 4 can be improved. First, we will clarify that the WPO gradient is only constant when the variance of the policy is fixed (e.g. only the mean is updated). Second, we will clarify that it is only standard policy gradient for which the variance is finite, while SVG(0) with the appropriate parameterization is identical to WPO, as you noted. Third, we will highlight that the agreement between WPO and SVG(0) is only true for specific choices of parameterization for SVG(0), and provide a concrete example where this identity is broken. For instance, replacing a Gaussian action distribution with an exponential distribution π(a) = exp(a/β)/β (with the natural reparameterization a = β * η) leads to different updates for WPO and SVG(0), even in expectation: the expected natural WPO update with respect to the scale parameters β will then simply be E[ ∇_a Q(a) ], while the SVG(0) update will be E[ ∇_a Q(a) diag(a/β) ], with η a standard exponential with p(η = x) = exp(-x). (If we consider the gradient with respect to the network parameters, both these updates just get multiplied with the Jacobian ∇_θ β, as per standard backprop.)\n6. We agree that using WPO in combination with diffusion models is a promising future avenue to investigate. Thank you for the suggestion.", "claims_and_evidence": "yes.", "methods_and_evaluation_criteria": "yes.", "theoretical_claims": "I have checked most of the proofs but cannot ensure that they are entirely accurate.", "experimental_designs_or_analyses": "The authors conduct extensive experiments on the widely used DMC and show that the proposed method is comparable to the commonly used methods.", "relation_to_broader_scientific_literature": "This paper proposes a new reinforcement learning policy update method, and the paper explains the motivation well from the perspective of Wasserstein gradient.", "essential_references_not_discussed": "Enough", "other_strengths_and_weaknesses": "The overall writing of the paper is clear, the overall motivation is well supported, and the relationship with previous work is fully discussed. My main concern about the paper is that the experimental performance results are relatively comparable with baseline methods mainly focusing on the comparison with the MPO algorithm, and the selected comparison methods lack popular algorithms such as PPO and TD3. The paper omits the $d^\\pi(s)$ term when deriving the formula, and I think it is better to add it. In line 358 of the left column, the author says that off-policy state data can be used, but this will cause $d^\\pi(s)$ to be outdated, making the gradient direction not the direction in which $\\mathcal{J}$ decreases the fastest.", "other_comments_or_suggestions": "In line 685, there is a missing right bracket before $\\nabla_a Q^\\pi(s,a)$. \n\nI really appreciate that authors conducted comparative experiments on the magnetic control problem, which is extremely valuable. Since the simulator is closed-source and few groups used it, it is difficult for other RL researchers to reproduce the experimental results. If authors have plans to make the experimental code open source, it will yield more benefits.", "questions_for_authors": "Many current works like Zhang et al. (2018) choose the reparameterization trick because the method is convenient for sampling actions. The authors emphasize in the paper that the main difference between WPO and other methods is whether to use parameterization. Can authors give some more practical examples or scenarios to illustrate the necessity of using some difficult-to-reparameterize policies?\n\nIn line 229 of the right column, the authors mentioned that for the sampled case, the variance of WPO is 0 when Q is locally linear in the actions. I am a little confused about this because for the method using reparameterization, this property is still enjoyed in this case due to $\\frac{\\partial Q}{\\partial \\theta}=\\frac{\\partial Q}{\\partial a}\\frac{\\partial a}{\\partial \\theta}=w(s)\\frac{\\partial a}{\\partial \\theta}$.\n\nHave the authors considered the application of this method to the diffusion model as a policy representation? The diffusion model can be considered to directly model $\\nabla_a \\log \\pi$. I think this may be interesting.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "1uA08xHbcy", "reviewer": "Reviewer_PGXN", "summary": "The paper applies Wasserstein gradient flows to reduce the parameter space, and this way obtain a closed-form update rule. It lifts the necessity of using the reparameterization trick in stochastic policy learning. The method merges the strengths of the policy gradient approaches that work on sampled evaluations of the action-value and deterministic policy gradients that work on a single steepest-ascent direction of the action-value. It can generate a vector field towards an increasing direction.", "questions": "", "limitations": "", "rating": 1, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their comments and hope to convince them of the merits of our approach.\n1. With regards to the performance of WPO relative to other baseline methods: while Figs. 3 and 6 show that WPO is competitive over the DM control suite, our main performance claims are highlighted in Figs. 4 and 5 that show that, as conjectured, WPO especially shines in domains with relatively high-dimensional action spaces. We would also dispute the characterization that “the alternative approaches consistently outperform the proposed WPO approach”. As we state in the paper, no single approach consistently outperforms the others across all tasks. Even on DeepMind Control Suite, there are tasks for which WPO performs the best, such as Cartpole - Three Poles, and many others where it matches state-of-the art. We also think the reviewer may misunderstand the main claim of the paper. We are *not* claiming that “training with a Wasserstein gradient flow facilitates learning” in all cases. We are claiming that WPO is novel, fills in a much needed gap in the space of actor-critic algorithms (how to generalize DPG-like updates to arbitrary stochastic policies without reparameterization) and, with only minimal effort, can be adapted into a performant deep RL algorithm.\n2. In response to the issues raised in “Experimental Designs and Analyses”: we wanted to show that the WPO gradient (Eq. 6) can be dropped into existing deep RL methods with minimal modification, and to do a fair comparison we wanted the baselines to be trained in as similar a fashion as possible. Our results use the same hyperparameters per algorithm across all tasks. Many deep RL algorithms are highly sensitive to small details in how they are trained. Thus we may exactly match the absolute best performance of each baseline method. However, where possible, we made sure to compare against numbers in the literature to make sure our comparison is fair - for instance, we both ran the open-source Dopamine implementation of SAC ourselves and added baseline numbers from the literature.\nSince submission, we implemented SAC in the same framework that was used for the other experiments, and will include those results in the final camera-ready paper. We are also running DDPG with 1-step updates instead of 5-step updates for comparison and can include those results in the appendix, although they do not look appreciably better than what we have already included. In general, while it may be possible to tune the hyperparameters to get better baseline performance, it may also be possible to tune WPO hyperparameters to get better performance, so we believe the comparison provided in the paper is quite fair.\n3. Re:TD3 and other possible baseline methods, we are primarily interested in the WPO update itself (Eq. 6), and compare like-for-like with algorithms like MPO. The algorithm we arrive at in this paper which we call “WPO” is one of an enormous number of possible algorithms that use the update in Eq 6. TD3 is an extension of DDPG that has some additional algorithmic ideas like target policy smoothing that can also be combined with other algorithms. It would be lovely to see future work decomposing these separable ideas and combining them in novel ways, to understand the individual contributions of all these ideas separately, rather than just comparing complete agents with somewhat-arbitrary combinations. Maybe WPO with TD3-like tricks will outperform TD3, maybe not, but that is out of scope of this particular paper.\n4. When you say “the proposed solution is straightforward in the sense that it does not yield downstream scientific challenges that would contribute to the advancement of the field” - we do not understand if this is a criticism or a compliment. A solution which is straightforward and does not create future challenges is a good thing, is it not?\n5. “Why is ∇_θπ(s_t) a Jacobian but not the gradient?” - ∇_θπ(s_t) is only a Jacobian in the case of DPG, because in this case π(s_t) is a deterministic vector-valued mapping from state vector to action vector. In the case of WPO or standard policy gradient, π(a_t|s_t) is a scalar-valued probability density and thus ∇_alogπ(a_t|s_t) is a vector of the same shape as π(s_t) in DPG.\n6. The DMC results are based on proprioceptive control, not pixel-based control.\n\nWe hope that this resolves any lingering issues with the paper. We want to emphasize again that our goal was not to show that we have the right combination of tricks and hyperparameter settings to achieve optimal performance across all possible tasks. We wanted to show that an open question in continuous control - how to train arbitrary stochastic policies with updates that exploit the action-value gradient - has a clear and elegant answer using approximations to Wasserstein gradient flows, and make a first foray into translating that into a practical deep RL algorithm. We hope you will agree that that is a significant contribution worth publishing at ICML.", "claims_and_evidence": "I am having hard time to conclude from Figures 3 and 4 that the proposed method improves the state of the art. The alternative approaches consistently outperform the proposed WPO approach. This is a rather strong evidence against the central claim of the paper: Training with a Wasserstein gradient flow facilitates learning.", "methods_and_evaluation_criteria": "The adopted methodology to build a deep actor-critic algorithm, the built experiment, pipeline, and the used performance scores make sense. There are issues about the choice of the baselines as well as the implementation details of the proposed algorithm. I give more details under Experimental Designs or Analyses.", "theoretical_claims": "I checked especially the functional derivative proof in Appendix A.1 in detail and it looks in order.", "experimental_designs_or_analyses": "The chosen baselines make sense. However, the DMC results do not match the ones I know from the literature and from my own experience. For instance, in Humanoid Walk, I am able to reach much higher performances with vanilla SAC in much less actor steps than what is reported. I suspect that this is because the authors train their networks with a single critic. There are also some additional hyperparameter choices which are uncommon, such as using 5-step temporal difference and using hard target network updates in intervals instead of Polyak averaging. All in all, I am not able to conclude that the reported results are informative about the match of the claims to evidence.", "relation_to_broader_scientific_literature": "The paper takes the Maximum A-Posterior Policy Optimization (MPO)  and DDPG as the prior work on which it brings an improvement. While I agree that MPO is an appropriate representative of stochastic policy gradient approaches, I do not think DDPG represents the state of the art. Its TD3 variant works significantly better. I would expect discussion about how the significant problem and the proposed method comes together with the established state of the art that uses DPG on an ensemble of critics, which essentially addresses quite many of the issues targeted by the WPO solution. For example the REDQ method is known to bring significant reduction to the gradient estimator variance while also correcting the overestimation bias. I would definitely expect a comparison against it to draw a conclusion about the value of the solution. The authors may see the following detailed benchmarking study to have an impression about what is possible with the existing approaches:\n\nNauman et al., Overestimation, Overfitting, and Plasticity in Actor-Critic: the Bitter Lesson of Reinforcement Learning, ICML, 2024", "essential_references_not_discussed": "The paper does cite the most important papers with respect to the chosen problem formulation. With this regard, I do not have concerns. However, I will reiterate that problem formulation has issues.", "other_strengths_and_weaknesses": "While the idea of using Wasserstein Gradient Flows to create a sweet spot between PPO-like and DDPG-like algorithms is an interesting idea, the proposed solution is straightforward in the sense that it does not yield downstream scientific challenges that would contribute to the advancement of the field. The implementation of the idea also has severe flaws I pointed out above.", "other_comments_or_suggestions": "Post-rebuttal update:\n\nHaving completed the rebuttal and discussions, I am still missing the concrete problem for which a solution is sought. The paper describes the problem as:\n\n\"Adding in stochasticity can be difficult to tune, and extensions that learn the variance (Heess et al., 2015; Haarnoja et al., 2018) rely on the reparameterization trick, which limits the class of policy distributions that can be used.\"\n\nMy grade stems for a serious doubt about the significance of this problem and the claim that the solution addresses it. The proposed solution does work, however the experiment setup does not represent the state of the art in the sense that the decade-old techniques to mitigate the harmful effects of the Bellman target are not used. The outcome is then a set of half-working models, one new, the other two are baselines. The new one is working tiny bit better than others. The questions would such a performance difference be visible when all the commonplace training stabilization techniques were used. \n\nThe authors clarify in their response is that the results are shown in the proprioceptive setup. I have first-hand experience to achieve much better results in much less environment interactions using multiple methods all of which are well-known in the literature. Under these circumstances I cannot conclude that the paper makes a concrete scientific contribution. Hence I maintain my score.", "questions_for_authors": "Why is $\\nabla_\\theta \\pi(s_t)$ a Jacobian but not the gradient?\n\nAre the DMC results obtained on the visual (pixel-based) or proprioceptive control setting? I interpret them as the latter in my evaluation above. I can update it based on the answer here.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: The paper proposes a novel actor critic algorithm, Wasserstein Policy Optimization, that uses Wasserstein gradient flows to optimize stochastic policies using the gradient of the state-action value function without relying on reparameterization tricks.
However some reviewers do raise the issue of weak empirical evaluation and baselines performing below par.
Decision: Accept (poster)