✓✓ Beats tuned baseline 2026

Entropy-Regularized Wasserstein Actor

Implementation & benchmark of arXiv:2608.07433 — Wasserstein Policy Gradient for Entropy-Regularized Linear-Quadratic Control

Usefulness7/10

Difficulty5/10

Novelty7/10

Optimization Stability Sample efficiency Accuracy Rl Optimizer Loss Optimal transport Dynamical systems Probability Control theory

Source paper: Wasserstein Policy Gradient for Entropy-Regularized Linear-Quadratic Control arXiv:2608.07433 ⓘ · analyzed Sep 1, 2026

AI-generated research hypothesis, automatically tested. Not peer-reviewed.

Idea description

Replace the usual parameter-space actor update with an action-space transport update. For every visited state, move sampled actions along a critic-improving velocity field while adding the entropy velocity, then fit the transported action cloud back to the actor's Gaussian mean and covariance. This preserves the paper's key idea that policy improvement is a Wasserstein flow over conditional action laws while remaining implementable for neural actors.

Formulas

$$\pi_{K,\Sigma}(\cdot\mid x)=\mathcal{N}(-Kx,\Sigma),\qquad \Sigma\succ0.$$

$$\partial_t\rho_t(a\mid x)+\nabla_a\!\cdot\big(\rho_t(a\mid x)v_t(a,x)\big)=0.$$

$$v_t(a,x)=\nabla_a Q_\phi(x,a)-\tau\nabla_a\log\rho_t(a\mid x).$$

$$\rho_t=\mathcal{N}(m_\theta(x),\Sigma_\theta(x))\quad\Longrightarrow\quad v_t(a,x)=\nabla_aQ_\phi(x,a)+\tau\Sigma_\theta(x)^{-1}\big(a-m_\theta(x)\big).$$

Mathematical statement

The paper's optimal policy class is the state-conditional Gaussian family $\pi_{K,\Sigma}(\cdot\mid x)=\mathcal{N}(-Kx,\Sigma)$, where $x$ is the state, $K$ is the linear feedback gain, and $\Sigma\succ0$ is the action covariance. Its policy-space tangent vectors are represented by action-space continuity equations, $\partial_t\rho_t+\nabla_a\cdot(\rho_t v_t)=0$, where $a$ is the action and $v_t(a,x)$ is the transport velocity. For the entropy-regularized actor objective $\mathcal{J}(\pi)=\mathbb{E}[Q_\phi(x,a)]+\tau\mathbb{E}[\mathcal{H}(\pi(\cdot\mid x))]$, with critic $Q_\phi$, temperature $\tau\ge0$, and differential entropy $\mathcal{H}(\rho)=-\int\rho\log\rho\,da$, the Wasserstein ascent velocity is $v(a,x)=\nabla_aQ_\phi(x,a)-\tau\nabla_a\log\rho(a\mid x)$. For a Gaussian actor $\rho=\mathcal{N}(m_\theta(x),\Sigma_\theta(x))$, $\nabla_a\log\rho(a\mid x)=-\Sigma_\theta(x)^{-1}(a-m_\theta(x))$, giving $v=\nabla_aQ_\phi(x,a)+\tau\Sigma_\theta(x)^{-1}(a-m_\theta(x))$. The first term moves actions toward higher critic value and the second expands low-entropy distributions. The paper proves that this transport is tangent to the Gaussian linear-policy family in LQ systems and produces globally well-posed, exponentially convergent gain/covariance dynamics; the neural adaptation tests whether this local stability benefit survives for nonlinear actors and critics.

Implementation notes

Integrate this at the actor update in an off-policy continuous-control algorithm such as SAC, TD3, or a vanilla Gaussian actor-critic. The actor outputs a mean $m_\theta(x)$ and a Cholesky factor $L_\theta(x)$, with covariance $\Sigma_\theta(x)=L_\theta(x)L_\theta(x)^\top+\epsilon I$. For each replay state $x_i$, draw $M$ actions $a_{ij}=m_\theta(x_i)+L_\theta(x_i)z_{ij}$ with $z_{ij}\sim\mathcal{N}(0,I)$. Compute the action-space velocity $v_{ij}=\nabla_aQ_\phi(x_i,a_{ij})+\tau\Sigma_\theta(x_i)^{-1}(a_{ij}-m_\theta(x_i))$, using automatic differentiation for $\nabla_aQ_\phi$. Form transported targets $\tilde a_{ij}=a_{ij}+h v_{ij}$, where $h$ is a small transport step. Fit the actor to these targets by minimizing $\mathcal{L}_{\mathrm{transport}}=\frac{1}{BM}\sum_{i,j}\|m_\theta(x_i)+L_\theta(x_i)z_{ij}-\tilde a_{ij}\|_2^2$, while retaining a covariance-floor penalty if needed. Use one or two gradient steps on this regression loss and stop gradients through the transported targets. Pseudocode: sample replay states; sample Gaussian actions; compute $g=\nabla_aQ$; compute $s=\Sigma^{-1}(a-m)$; set $v=g+\tau s$; set $a_{\mathrm{target}}=a+h v$; update $\theta\leftarrow\theta-\eta_\theta\nabla_\theta\mathcal{L}_{\mathrm{transport}}$. Clip $\|v\|$ and choose $h\|\nabla_aQ\|$ below a fixed fraction of the action range. The paper supplies the continuity-equation/Wasserstein interpretation and Gaussian invariant structure in LQ systems; the neural implementation estimates the critic gradient, discounted state occupancy through replay sampling, and covariance inverses numerically. First compare against SAC on Pendulum-v1 and HalfCheetah-v4 using identical critics, replay buffers, environment interactions, and gradient counts. Measure return versus environment samples, actor-gradient variance, critic-target variance, covariance collapse, and sensitivity to actor learning rate. Success means faster return improvement, fewer entropy-collapse failures, or lower actor-gradient variance at matched final return.

Verification

Beats tuned baseline

Stage 1 · Toy mechanism gate: Passed ✓

Stage 2 · Mechanism transferred to benchmark: Not reproduced ✗

Stage 2 · Practical benchmark result: Significant win ✓

Methodology: 8 paired seeds, permutation test (p<0.05) against a learning-rate-tuned baseline. How verification works

How to read this: The implementation beat the baseline, but the claimed mechanism did not reproduce at benchmark scale. The gain may come from another effect; treat the cause as unresolved.

Stage 1 — Mechanism check agent confidence 8/10

Built an entropy-regularized Wasserstein actor toy implementation with critic-gradient transport, entropy velocity, stop-gradient transported targets, and a standard reparameterized actor baseline. The quantitative mechanism checks passed: instability appeared between hλ=1.99 and 2.01, stationary variance matched τ/λ to approximately 1e-15, and observed mean/variance rates matched λ and 2λ within about 0.4%. In the eight-run toy comparison, transport had mean error 1.82e-5 versus 0.00423 and variance error 0.000127 versus 0.00306 for the baseline; this is a promising controlled signal, not evidence of an RL win.

Agent confidence: 8/10

Baseline: 8 runs: mean absolute mean error 0.0042277; mean absolute variance error 0.0030616; final objective loss -0.1178
Idea: 8 runs: mean absolute mean error 0.00001823; mean absolute variance error 0.0001271; transport regression loss 5.50e-10. Stability observed at hλ=1.99 and divergence at hλ=2.01; variance equilibrium errors <=1.6e-15; rate errors <=0.4%.

Limitations: Only a one-dimensional quadratic-critic Gaussian toy setting was tested. No Pendulum-v1, HalfCheetah-v4, replay-buffer training, learned critic, multidimensional covariance, actor-gradient variance, wall-clock, or environment sample-efficiency benchmark was run; the baseline and transport losses are also not directly comparable because they optimize different objectives.

How to run: python3 experiment.py

Stage 2 — Benchmark vs. tuned baseline (arXiv:2608.07433) Sep 2, 2026

Beats the tuned baseline by 22.9% (p=0.0081, wins 8 / 8 paired seeds; metric: lower is better).

Benchmark: Dynamics forecasting
Model: rnn_small
Paired seeds: 8
Baseline mean: 0.0007
Idea mean: 0.0005
Effect (Δ): -0.0002 (−22.9%; negative = idea better)
Wins: 8 / 8 paired seeds
p-value: 0.0081 (permutation test, 20 000 shuffles)
Smallest detectable effect: n/a
Mechanism: Not confirmed ✗
Practical effect: Helps
Baseline tuning: swept over 3 configs

Limitations:

Only supervised controlled-pendulum dynamics was tested; no online SAC/TD3, replay-buffer interaction, HalfCheetah, multidimensional covariance, variance metrics, or wall-clock study. Mechanism signature correlation was 0.12919 and confirmed=false.

How to run: python3 stage2_bench.py

Verdict computed by deterministic test code from paired-seed statistics — not by the language model.

Stage 2 — Benchmark Sep 1, 2026 latest Worked ✓

Agent confidence: 8/10

Idea better (significant): transport achieved lower held-out MSE on all 8 paired seeds, with permutation p=0.0081. The trained-model mechanism signature was not confirmed, so this is a task-level win rather than confirmation of transport alignment.

Baseline: mean MSE 0.0007032144
Delta mean: -0.0001607343
Idea: mean MSE 0.0005424801
P value: 0.0081

Benchmark result