---
title: Wasserstein Barycenter Soft Actor-Critic
url: https://www.emergentmind.com/topics/wasserstein-barycenter-soft-actor-critic-wbsac
type: topic
---

# Wasserstein Barycenter Soft Actor-Critic

Wasserstein Barycenter Soft Actor-Critic (WBSAC) is an off-policy reinforcement learning algorithm for continuous control that augments the standard Soft Actor-Critic (SAC) framework with a principled directed exploration strategy. By combining conservative (pessimistic) and exploratory (optimistic) policies through their Wasserstein barycenter, WBSAC addresses the problem of poor sample efficiency in environments characterized by sparse rewards. The algorithm adaptively interpolates between exploitation and exploration by dynamically scheduling the weights of the pessimistic and optimistic actors throughout learning, thereby enhancing exploration without sacrificing value estimation stability [2506.10167].

## 1. Motivation and Context

Deep off-policy actor–critic methods—including DDPG, TD3, and SAC—have established state-of-the-art performance in continuous control. However, their exploration strategies typically depend on undirected noise or entropy bonuses, yielding inefficient exploration in sparse-reward tasks; the agent is rarely exposed to positive reward signals, causing wasted samples and stagnation. WBSAC leverages the “optimism in the face of uncertainty” principle by maintaining (i) a pessimistic actor that focuses on reliable value estimation according to the lower bound of the critic and (ii) an optimistic actor that targets high epistemic uncertainty (as quantified by critic disagreement). This dual-actor approach is explicitly blended via the Wasserstein barycenter, forming the exploration policy.

## 2. Mathematical Formulation

### Pessimistic and Optimistic Actors

Let $Q_{\theta_1}$ and $Q_{\theta_2}$ denote the critic networks, and let $\pi_p(\cdot|s; \varphi)$ and $\pi_o(\cdot|s; \phi)$ denote two Gaussian policy distributions (pessimistic and optimistic, respectively). Their objectives are:

- **Pessimistic Actor Loss:**
  $$
  \mathcal{L}_{\pi_p}(\varphi) = \mathbb{E}_{s \sim \mathcal{D},\, a \sim \pi_p} [\,\alpha \log \pi_p(a|s) - \min_{i=1,2} Q_{\theta_i}(s, a)\,].
  $$
  The entropy temperature $\alpha$ is tuned via
  $$
  \mathcal{L}_\alpha(\alpha) = \mathbb{E}_{s, a \sim \pi_p} [-\alpha (\log \pi_p(a|s) + \mathcal{H}_0)].
  $$

- **Optimistic Actor Loss:**
  $$
  \mathcal{L}_{\pi_o}(\phi) = \mathbb{E}_{s \sim \mathcal{D},\, a \sim \pi_o} [-(\mu_Q(s,a) + \beta_o \sigma_Q(s,a))]
  $$
  where
  $$
  \mu_Q = \tfrac{1}{2}(Q_{\theta_1} + Q_{\theta_2}), \quad
  \sigma_Q = \sqrt{\tfrac{1}{2} \sum_{i=1}^2 (Q_{\theta_i} - \mu_Q)^2}.
  $$
  The hyperparameter $\beta_o > 0$ adjusts the intensity of the uncertainty bonus.

### Wasserstein Barycenter Exploration Policy

The exploration policy $\pi_e$ is constructed as the 2-Wasserstein barycenter (with Euclidean cost) of the pessimistic and optimistic actors:
$$
\pi_e = \arg\min_{\pi} \{ \xi_p W_2^2(\pi, \pi_p) + \xi_o W_2^2(\pi, \pi_o) \}, \quad \xi_p + \xi_o = 1.
$$
For Gaussian policies, this yields
$$
\mu_e = \xi_p \mu_p + \xi_o \mu_o, \qquad
\Sigma_e = ( \xi_p \Sigma_p^{1/2} + \xi_o \Sigma_o^{1/2} )^2.
$$

### Scheduling Interpolation Weights

Initially $\xi_o = 0$ (purely pessimistic); $\xi_o$ is increased to 1 over training steps, e.g., using a schedule parameter $\lambda$:
$$
\xi_o(t) = 1 - \exp(-t/\lambda), \quad \xi_p(t) = 1 - \xi_o(t).
$$
This ensures early data collection is conservative and late-stage exploration is predominant.

## 3. Training and Optimization Procedure

### Critic (TD) Update

The critic is trained using the following loss:
$$
\mathcal{L}_Q(\theta_i) = \mathbb{E}_{(s, a, r, s') \sim \mathcal{D}} \left[
Q_{\theta_i}(s, a) - (r + \gamma [ \min_j Q_{\theta'_j}(s', a') - \alpha \log \pi_p(a'|s') ])
\right]^2,
$$
where $a' \sim \pi_p(\cdot|s')$ and $\theta'$ are target network parameters.

### Actor and Temperature Updates

The pessimistic and optimistic actors and temperature parameter are updated using their respective losses as previously formulated.

### Data Collection and Update Cycle

1. The agent observes state $s_t$ and samples $a_t \sim \pi_e(\cdot|s_t)$ from the barycenter.
2. The environment transitions to $(s_{t+1}, r_t)$; experience $(s_t, a_t, r_t, s_{t+1})$ is added to the replay buffer.
3. The interpolation weights $(\xi_p, \xi_o)$ are periodically updated per the chosen schedule.
4. For each training step, mini-batches are used to update in order: critic parameters, optimistic actor, pessimistic actor, temperature, and target networks.

### Complete Training Step Summary

```
Initialize θ₁, θ₂; φ, ϕ; α; replay buffer 𝒟

for each environment step:
    compute π_p(.|s), π_o(.|s), π_e(.|s) via barycenter
    a ∼ π_e(.|s), step env → (s', r), store in 𝒟
    update ξ_o by schedule
    for N gradient steps:
        sample batch from 𝒟
        update critics θ₁, θ₂ via ℒ_Q
        update optimistic actor ϕ via ℒ_{π_o}
        update pessimistic actor φ via ℒ_{π_p}
        update α via ℒ_α
        soft-update targets θ′
```

## 4. Theoretical Analysis

A central theoretical property of WBSAC is a guaranteed entropy lower bound for the barycenter exploration policy. For factorized Gaussian $\pi_p$, $\pi_o$, and their barycenter $\pi_e$, the following holds for each state $s$:
$$
H(\pi_e(\cdot|s)) \geq \xi_p H(\pi_p(\cdot|s)) + \xi_o H(\pi_o(\cdot|s)).
$$
Because Gaussian differential entropy is concave in the covariance, the barycenter preserves at least the weighted average entropy of the component policies. This property ensures that the exploration policy maintains sufficient stochasticity, which is critical for effective exploration in continuous spaces. This entropy bound is geometrically intrinsic to the Wasserstein barycenter construction [2506.10167].

## 5. Empirical Evaluation and Benchmarking

WBSAC was evaluated on five MuJoCo continuous control tasks (Ant-v5, HalfCheetah-v5, Walker2d-v5, Humanoid-v5, Hopper-v5) and the sparse-reward PointMaze Medium-v3 environment. The following hyperparameters were used consistently across MuJoCo experiments: two-layer networks of 256 ReLU units per layer, Adam optimizer (lr = $3 \times 10^{-4}$), batch size 256, replay buffer size $10^6$, discount $\gamma=0.99$, target smoothing $\tau=0.005$, initial entropy coefficient $\alpha=0.2$, $\beta_o=1.5$, $\lambda=10$. Results were averaged over five random seeds for one million environment steps, with evaluation every 5,000 steps and reporting the mean and standard deviation across the final ten evaluations.

Key results:

| Task             | WBSAC (mean ± std) | SAC (mean ± std) | DARC (mean ± std)  |
|------------------|--------------------|------------------|--------------------|
| Ant-v5           | 3408 ± 495         | 3525 ± 1113      | 2983 ± 311         |
| HalfCheetah-v5   | **6466 ± 1411**    | 5409 ± 2219      | 4234 ± 1768        |
| Walker2d-v5      | **4417 ± 703**     | 3939 ± 262       | 3761 ± 485         |
| Humanoid-v5      | **5179 ± 89**      | 4996 ± 186       | 3815 ± 1581        |
| Hopper-v5        | 1920 ± 938         | **2460 ± 486**   | 2422 ± 876         |

In four out of five MuJoCo tasks, WBSAC matches or exceeds both SAC and DARC, and typically learns faster (e.g., on HalfCheetah-v5, WBSAC peaks at approximately 200k environment steps versus 400k for SAC). On the sparse-reward PointMaze Medium-v3, both WBSAC and SAC achieve similar final rewards over 100k steps, but WBSAC yields approximately 20% higher state coverage, indicating improved exploration.

Sensitivity analysis with respect to $\lambda$ and $\beta_o$ demonstrates robust performance across a wide range of schedules. Ablation studies confirm that dynamic scheduling of the barycenter coefficients and the incorporation of the critic-disagreement term in the optimistic actor are crucial components for realized performance.

## 6. Summary and Connections

WBSAC provides a geometric mechanism for regulated exploration in deep reinforcement learning by interpolating between pessimistic and optimistic actors with a Wasserstein barycenter. The algorithm begins with a conservative strategy to stabilize value estimation and transitions to increasingly optimistic exploration. Formal entropy lower bounds guarantee that exploration remains sufficiently stochastic. Empirical results show that WBSAC improves sample efficiency and achieves or surpasses state-of-the-art performance on both dense- and sparse-reward continuous control benchmarks, without introducing overestimation bias [2506.10167].

A plausible implication is that Wasserstein barycentric policy interpolation may generalize to other policy classes or uncertainty sources, motivating further investigation into its applicability in broader reinforcement learning contexts.

Source: https://www.emergentmind.com/topics/wasserstein-barycenter-soft-actor-critic-wbsac