Policy Dual Averaging (PDA)

Policy Dual Averaging (PDA) is an on-policy reinforcement learning algorithm with theoretical guarantees and competitive empirical performance. This package is a lightweight PyTorch implementation of PDA and actor-accelerated PDA for problems with discrete and continuous action space. The environment wrappers and on-policy training loops are adapted from CleanRL and Tianshou. We use Gymnasium for environment simulation.

The algorithm is based on paper Actor-Accelerated Policy Dual Averaging for Reinforcement Learning in Continuous Action Spaces, and Policy Optimization over General State and Action Spaces.

PDA

On-policy RL methods (e.g., TRPO, PPO) under the Policy Mirror Descent (PMD) framework regularize the new policy toward the previous iterate. PDA instead regularizes (through Bregman divergence $D(\cdot,\cdot)$ ) toward a fixed prox-center policy $\pi_0$ and accumulates all past advantages $\psi^{\pi_t}$ with weights $\beta_t$:

$$ \pi_{k+1}(s) = \arg\min_{a \in \mathcal{U}} \sum_{t=0}^{k} \beta_t , \psi^{\pi_t}(s, a) + \lambda_k , D\bigl(\pi_0(s), a\bigr). $$

Acceleration: PDA's policy is defined implicitly. Every time an action $\pi_{k+1}(s)$ is needed (at every environment step), we have to solve the optimization subproblem over the action space $\mathcal{U}$, which is theoretically clean but computationally challenging for deep RL at scale. Actor-accelerated PDA addresses this by training a neural network actor $\hat\pi_\theta$ to approximate the minimizer of the cumulative regularized objective: $\hat\pi_{k+1}(s;\theta^\pi) \approx \pi_{k+1}(s)$.

Benchmarks

Comparison with on-policy methods: Actor-accelerated PDA consistently outperforms on-policy baselines (PPO, TRPO, NPG) on MuJoCo-v4 and OR-Gymnasium. Hyperparameters are kept fixed across all environments. Benchmarks are conducted using our Tianshou (V1.2.0) implementation using Tianshou tuned hyperparameters. Benchmark results table records average episodic return over the final 5 epochs, 10 seeds × 10 evaluations per seed, 1M (or 3M) environment steps.

Comparison with off-policy methods: PDA is an on-policy algorithm yet remains competitive with SAC and TD3 on many tasks, with 10–40× faster wall-clock time.

Algorithm runtime on MuJoCo-v4: Average runtime (5 parallel runs) for 1M environment steps of training and testing. GPU is used for off-policy methods (SAC and TD3) for acceleration due to larger neural network sizes. Times are reported in seconds (mean ± standard deviation). Hardware: 14700K, RTX 4070.

Installation

Install with:

pip install pdarl
pip install "pdarl[mujoco]" # MuJoCo envs
pip install "pdarl[track]" # Weights & Biases
pip install "pdarl[all]" # MuJoCo + W&B

For other envs, see Gymnasium for installation details.

We included CleanRL style single-file RL script for PDA_ACT and PDA_DSC in folder cleanrl_pda. The files can be copied and ran directly when dependencies are installed.

Quickstart

Copy the config folder or create your own and train actor-accelerated PDA (default) with a config file:

pdarl-run --config config/pda_act.yaml

Use the config files for other PDA variants and PMD:

pdarl-run --config config/pda_opt_rgd.yaml
pdarl-run --config config/pda_opt_acfgm.yaml
pdarl-run --config config/pda_bcd.yaml
pdarl-run --config config/pda_dsc.yaml
pdarl-run --config config/pmd_act.yaml

Customization

To customize the agents and training process, customize agents like PDA_ACT and Trainer class directly. See pdarl/trainer.py for more details.

import numpy as np
 
from pdarl.utils.args import load_args_from_yaml
from pdarl.trainer import Trainer
from pdarl import PDA_ACT

args = load_args_from_yaml("config/pda_act.yaml")

class CustomPDA(PDA_ACT):
    # Override agent for custom behavior
    def __init__(self, obs_dim, act_dim, args, device):
        super().__init__(obs_dim, act_dim, args, device)
        print("Custom PDA initialized!")
    
    def compute_act(self, obs):
        return super().compute_act(obs)

class CustomTrainer(Trainer):
    # Override trainer setup to use CustomPDA
    def _setup_agent(self):
        obs_space = self.env.env.observation_space
        act_space = self.env.env.action_space
        obs_dim = int(np.array(obs_space.shape).prod())
        act_dim = int(np.prod(act_space.shape))
        self.agent = CustomPDA(obs_dim, act_dim, self.args, self.device).to(self.device)
        print("Custom Trainer agent changed to CustomPDA!")

trainer = CustomTrainer(args)
trainer.setup()
trainer.run()

Logging

When logging: true (default), TensorBoard scalars are written under runs/. View them with:

tensorboard --logdir runs

Set track: true in a config file to log to Weights & Biases (wandb_project_name, wandb_entity).

Agents

Key configs

See pdarl/utils/args.py and the files under config for the full schema.

To cite this repository

@misc{gao2026actorpda,
  title={Actor-Accelerated Policy Dual Averaging for Reinforcement Learning in Continuous Action Spaces},
  author={Gao, Ji and Ju, Caleb and Lan, Guanghui and Tong, Zhaohui},
  year={2026},
  eprint={2603.10199},
  archivePrefix={arXiv},
  primaryClass={cs.LG},
  url={https://arxiv.org/abs/2603.10199}
}

@misc{ju2022policyoptimizationgeneralstate,
  title={Policy Optimization over General State and Action Spaces},
  author={Ju, Caleb and Lan, Guanghui},
  year={2022},
  eprint={2211.16715},
  archivePrefix={arXiv},
  primaryClass={cs.LG},
  url={https://arxiv.org/abs/2211.16715}
}

References

@article{huang2022cleanrl,
  author  = {Shengyi Huang and Rousslan Fernand Julien Dossa and Chang Ye and Jeff Braga and Dipam Chakraborty and Kinal Mehta and João G.M. Araújo},
  title   = {CleanRL: High-quality Single-file Implementations of Deep Reinforcement Learning Algorithms},
  journal = {Journal of Machine Learning Research},
  year    = {2022},
  volume  = {23},
  number  = {274},
  pages   = {1--18},
  url     = {http://jmlr.org/papers/v23/21-1342.html}
}

@article{weng2022tianshou,
  title   = {Tianshou: A Highly Modularized Deep Reinforcement Learning Library},
  author  = {Weng, Jiayu and Chen, Hang and Yan, Dong and You, Kaiwen and Zhang, Chen and Gong, Yunchang and Zhu, Rongqing and Li, Sijin and Zhou, Yiwen and Liu, Qian and Song, Yueyang},
  journal = {Journal of Machine Learning Research},
  year    = {2022},
  volume  = {23},
  number  = {267},
  pages   = {1--6},
  url     = {http://jmlr.org/papers/v23/21-1127.html}
}

@misc{towers2024gymnasium,
  author = {Towers, Mark and Kwiatkowski, Ariel and Terry, Jordan and Balis, John U and De Cola, Gianluca and Deleu, Tristan and Goul{\~a}o, Manuel and Kallinteris, Andreas and Krimmel, Markus and KG, Arjun and others},
  title  = {Gymnasium: A Standard Interface for Reinforcement Learning Environments},
  year   = {2024},
  eprint = {arXiv:2407.17032}
}

@misc{HubbsOR-Gym,
    author={Christian D. Hubbs and Hector D. Perez and Owais Sarwar and Nikolaos V. Sahinidis and Ignacio E. Grossmann and John M. Wassick},
    title={OR-Gym: A Reinforcement Learning Library for Operations Research Problems},
    year={2020},
    Eprint={arXiv:2008.06319}
}