A PyTorch implementation of Soft Actor-Critic (SAC) augmented with a dedicated exploration policy whose output is blended with the actor via a Gaussian Wasserstein barycenter. Tested on MuJoCo and Deepmind Control suit Tasks.
Standard SAC balances exploitation and entropy-regularized exploration through a single actor. This implementation adds a separate explorer network that explicitly maximizes Q-value uncertainty (the gap between the two critics), encouraging it to seek out under-explored regions of the state space.
At each environment step, the actor and explorer produce independent Gaussian distributions. These are merged into a single Wasserstein barycenter distribution, which is then used to sample the action. A scheduling parameter a linearly shifts the blend weight from the actor toward the explorer over the course of training, controlled by --transition_rate.
action ~ Barycenter( Actor(s), Explorer(s), weights=(1-a, a) )
Key components:
- Actor — Standard SAC actor, optimized with entropy regularization.
- Explorer — Separate policy trained to maximize
(Q1 + Q2)/2 + β * |Q1 - Q2| / 2, where β (--explore_beta) scales the uncertainty bonus. - Twin Critics — Two Q-networks with soft target updates (REDQ-style repeated updates supported).
- Auto-entropy tuning — Automatic adjustment of the entropy coefficient α.
pip install torch numpy gymnasium[mujoco]python agent.py --seed 1Arguments:
| Flag | Default | Description |
|---|---|---|
--seed / -r |
1 |
Base random seed for training |
--updates_per_step |
1 |
Gradient updates per environment step |
--redq_updates |
1 |
Critic update repetitions per learning call |
--warmup_steps |
0 |
Steps before learning begins |
--eval_episodes |
10 |
Episodes per evaluation rollout |
Evaluation results are saved to:
results/actor_eval_seed_<seed>.pkl
This is a Python list of average returns, one entry per evaluation checkpoint (every 5,000 steps).
| Parameter | Value |
|---|---|
| Replay buffer size | 1,000,000 |
| Batch size | 256 |
| Actor / Critic LR | 3e-4 |
| Discount γ | 0.99 |
| Target update τ | 0.005 |
| Reward scale | 20 |
| Explorer β | 1.5 |
| Transition rate | 2.5 |
Trained and evaluated on different v5 environments from Gymnasium MuJoCo. Training and evaluation seeds are kept disjoint — evaluation uses a fixed offset of +100,000 from the base training seed to prevent overlap.
This implementation is based on the following paper:
@article{shahrooei2025wbsac,
title = {Wasserstein Barycenter Soft Actor-Critic},
author = {Shahrooei, Zahra and Baheri, Ali},
journal = {arXiv preprint arXiv:2506.10167},
year = {2025},
url = {https://arxiv.org/abs/2506.10167}
}Shahrooei, Z., & Baheri, A. (2025). Wasserstein Barycenter Soft Actor-Critic. arXiv:2506.10167.