Behaviour Policy Optimization (BPO)

The repository provide the official implementation of Behaviour Policy Optimization (BPO), an off-policy extension of the classic Proximal Policy Optimization (PPO) algorithm with provable variance reduction. The full details can be found in:

Behaviour Policy Optimization: Provably Lower Variance Return Estimates for Off-Policy Reinforcement Learning Alexander W. Goodall, Edwin Hamel-De le Court, Francesco Belardinelli arXiv: https://arxiv.org/abs/2511.10843

Disclaimer: This implementation is simplified from the full version used to collect the results provided in the paper.

Software stack

Training runs the actor-critic on CPU or CUDA GPUs with JAX, uses TensorFlow and TensorFlow Probability for action distributions and TensorBoard, and steps through the environments with Gymnasium.

Quick Start

Installation

Prerequisites

Python 3.12 is required.

Installation with conda

• Install conda, e.g., via anaconda.
• Clone the repo:

git clone https://github.com/sacktock/BPO.git
cd BPO

• Create a conda virtual environment:

conda env create --name bpo --file conda-environment.yaml
conda activate bpo

• Install the project and its dependencies:

pip install -e .

pip install -e . installs the dependencies declared in pyproject.toml. requirements.txt lists the same base dependencies and is kept in lockstep with that file.

Installation with uv

From the repository root, uv sync installs the base dependency set (CPU JAX):

uv sync

Enabling GPU acceleration with JAX (optional)

BPO uses JAX for the actor-critic. CPU is enough to smoke test the environments and run a short training job.

macOS and Windows: we recommend JAX with CPU. No further action is required if you correctly followed the earlier steps.

pip

Linux x86_64/aarch64 and Windows WSL2 x86_64: We strongly recommend using Ubuntu 22.04 or similar. jax and jaxlib 0.4.35 should already be installed. You need to reinstall JAX for your CUDA driver. Do not use the -U option here as it may break earlier dependencies.

pip install "jax[cuda12]==0.4.35"

CUDA 13 does not ship with this stack. JAX 0.4.35 only publishes a CUDA 12 extra, so CUDA 13 needs a JAX and jaxlib upgrade, which is a later change.

uv

uv sync --group cuda12

Follow-up: native JAX distributions

TensorFlow Probability pins this repository to TensorFlow 2.18, NumPy 1.26, and JAX 0.4.35. Moving to a newer JAX 0.7 stack, with a native JAX distributions implementation in place of TensorFlow Probability, is a follow-up. It does not ship today, so results stay reproducible on the stack used for the paper.

Extras

MuJoCo (used by Ant, HalfCheetah, Hopper, and Walker2d) is installed through the gymnasium[mujoco] extra. On a headless Linux machine (CI job), also install OSMesa and select it, or those environments fail while importing the renderer:

sudo apt-get install -y libosmesa6 libgl1 libglib2.0-0
export MUJOCO_GL=osmesa PYOPENGL_PLATFORM=osmesa

How to run BPO

Training is configured from bpo/configs.yaml. Named blocks are merged left to right onto defaults. Later configs override earlier ones.

After pip install -e . (conda or pip), the bpo console script is on PATH. With uv, pass the same arguments to uv run bpo.

bpo --configs cartpole_ppo

uv run bpo --configs cartpole_ppo

Any key in the loaded config can be overridden from the command line (for example --env.env_id ant, --run.seed 0).

Running experiments

The basic pattern is:

# conda or pip
bpo --configs <base_config> [<override_config> ...] [flag overrides...]

# uv
uv run bpo --configs <base_config> [<override_config> ...] [flag overrides...]

• --configs can take one or more names, each corresponding to a top-level key in bpo/configs.yaml (e.g., cartpole_ppo, mujoco_ppo_gsde, ppo_bpo_zero, ...).
• After config merging, any remaining CLI arguments are parsed as overrides (e.g., --env.env_id ant, --run.seed 0, etc.).

A minimal PPO run (CartPole preset):

# conda or pip
bpo --configs cartpole_ppo

# uv
uv run bpo --configs cartpole_ppo

This uses the cartpole_ppo block (timesteps, env, PPO hyperparameters, etc.). A MuJoCo PPO run (Ant baseline preset):

# conda or pip
bpo --configs mujoco_ppo --env.env_id ant --run.seed 0 --run.logdir runs/mujoco/ant/ppo_seed_0

# uv
uv run bpo --configs mujoco_ppo --env.env_id ant --run.seed 0 --run.logdir runs/mujoco/ant/ppo_seed_0

(mujoco_ppo sets MuJoCo-style PPO defaults; --env.env_id can be changed among ant, half_cheetah, hopper, walker_2d.)

Enabling BPO

Option A -> flip the flag directly

# conda or pip
bpo --configs mujoco_ppo_gsde --bpo True --env.env_id ant --run.seed 0

# uv
uv run bpo --configs mujoco_ppo_gsde --bpo True --env.env_id ant --run.seed 0

Option B -> include a BPO config block For example, ppo_bpo_zero is a small "add-on" config that sets bpo: True and applies BPO-specific settings (e.g., symlog_targets, polyak_tau, and a zero-norm-final Q-head).

# conda or pip
bpo --configs mujoco_ppo_gsde ppo_bpo_zero --env.env_id ant --run.seed 0

# uv
uv run bpo --configs mujoco_ppo_gsde ppo_bpo_zero --env.env_id ant --run.seed 0

Switching BPO hyperparameter settings

The repo includes several ready-made BPO variants that mainly differ in importance-weight clipping (clip_rho, clip_c) and whether trajectory clipping is enabled (clip_traj).

Common presets:

• ppo_bpo_zero -> clip_rho=1.5, clip_c=1.5, plus "zero final norm" Q-head
• ppo_bpo_zero1.0 -> clip_rho=1.0, clip_c=1.0
• ppo_bpo_zero1.0_1.5 -> clip_rho=1.5, clip_c=1.0
• ppo_bpo_zero1.0_1.4 -> clip_rho=1.4, clip_c=1.0
• ppo_bpo_zero1.0_traj -> clip_rho=1.0, clip_c=1.0, clip_traj=True
• ppo_bpo_zero1.5_traj -> clip_rho=1.5, clip_c=1.5, clip_traj=True

So, to run the "rho/c = 1.0/1.0" setting you can do:

# conda or pip
bpo --configs mujoco_ppo_gsde ppo_bpo_zero1.0 --env.env_id ant --run.seed 0

# uv
uv run bpo --configs mujoco_ppo_gsde ppo_bpo_zero1.0 --env.env_id ant --run.seed 0

Logging and Miscellaneous

TensorBoard logging is supported by our implementation and is enabled in the following example:

# conda or pip
bpo \
  --configs mujoco_ppo_gsde ppo_bpo_zero \
  --env.env_id ant \
  --tensorboard True \
  --logdir "runs/mujoco/ant/ppo_gsde_fqe_zero_seed_0" \
  --seed 0

# uv
uv run bpo \
  --configs mujoco_ppo_gsde ppo_bpo_zero \
  --env.env_id ant \
  --tensorboard True \
  --logdir "runs/mujoco/ant/ppo_gsde_fqe_zero_seed_0" \
  --seed 0

The TensorBoard logs can be accessed via the command line:

tensorboard --logdir runs/mujoco/ant/

The verbose flag controls how much diagnostic information is logged during training. This affects both console output and the set of metrics written to TensorBoard / logs. There are three verbosity levels, with default set to 0.

The verbosity level can be changed via a command-line override:

--run.verbose <level>

Similarly, you can also override any nested field directly from the command line (after configs are merged). For example, to keep ppo_bpo_zero but change the clipping thresholds:

# conda or pip
bpo --configs mujoco_ppo_gsde ppo_bpo_zero \
  --env.env_id ant \
  --ppo.clip_rho 1.0 --ppo.clip_c 1.0 \
  --run.seed 0

# uv
uv run bpo --configs mujoco_ppo_gsde ppo_bpo_zero \
  --env.env_id ant \
  --ppo.clip_rho 1.0 --ppo.clip_c 1.0 \
  --run.seed 0

Tests

With uv:

uv sync --group dev
uv run pytest -v

With conda or pip, install the project, then pytest, and run from the repository root:

pip install -e .
pip install "pytest>=9"
pytest -v

The test suite is small and checks that the scientific stack imports, that each environment can be constructed and stepped, and that a tiny PPO or BPO training job can start.

Dependency sync

requirements.txt must match [project.dependencies] in pyproject.toml. CI runs scripts/check_requirements_sync.py when either file changes. Locally:

python scripts/check_requirements_sync.py

License

Our implementation of BPO is released under the MIT License.