Wasserstein Barycenter Soft Actor-Critic

• WBSAC is an off-policy reinforcement learning algorithm that augments SAC with a directed exploration strategy through a Wasserstein barycenter blending pessimistic and optimistic policies.
• It dynamically interpolates between conservative value estimation and uncertainty-driven exploration, enhancing sample efficiency in sparse-reward environments.
• Empirical evaluations on MuJoCo tasks confirm WBSAC achieves faster learning and improved state coverage compared to standard SAC and DARC.

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 (Shahrooei et al., 11 Jun 2025).

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 $Q_{\theta_2}$0 is constructed as the 2-Wasserstein barycenter (with Euclidean cost) of the pessimistic and optimistic actors:

$Q_{\theta_2}$1

For Gaussian policies, this yields

$Q_{\theta_2}$2

Scheduling Interpolation Weights

Initially $Q_{\theta_2}$3 (purely pessimistic); $Q_{\theta_2}$4 is increased to 1 over training steps, e.g., using a schedule parameter $Q_{\theta_2}$5:

$Q_{\theta_2}$6

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:

$Q_{\theta_2}$7

where $Q_{\theta_2}$8 and $Q_{\theta_2}$9 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 $\pi_p(\cdot|s; \varphi)$0 and samples $\pi_p(\cdot|s; \varphi)$1 from the barycenter.
2. The environment transitions to $\pi_p(\cdot|s; \varphi)$2; experience $\pi_p(\cdot|s; \varphi)$3 is added to the replay buffer.
3. The interpolation weights $\pi_p(\cdot|s; \varphi)$4 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

$\pi_o(\cdot|s; \phi)$9

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(\cdot|s; \varphi)$5, $\pi_p(\cdot|s; \varphi)$6, and their barycenter $\pi_p(\cdot|s; \varphi)$7, the following holds for each state $\pi_p(\cdot|s; \varphi)$8:

$\pi_p(\cdot|s; \varphi)$9

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 (Shahrooei et al., 11 Jun 2025).

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 = $\pi_o(\cdot|s; \phi)$0), batch size 256, replay buffer size $\pi_o(\cdot|s; \phi)$1, discount $\pi_o(\cdot|s; \phi)$2, target smoothing $\pi_o(\cdot|s; \phi)$3, initial entropy coefficient $\pi_o(\cdot|s; \phi)$4, $\pi_o(\cdot|s; \phi)$5, $\pi_o(\cdot|s; \phi)$6. 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 $\pi_o(\cdot|s; \phi)$7 and $\pi_o(\cdot|s; \phi)$8 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 (Shahrooei et al., 11 Jun 2025).

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.