PB²: Preference Space Exploration via Population-Based Methods in Preference-Based Reinforcement Learning

Official implementation of "PB²: Preference Space Exploration via Population-Based Methods in Preference-Based Reinforcement Learning" (Reinforcement Learning Conference, RLC 2026).

Paper: https://arxiv.org/abs/2506.13741

Overview

Preference-based reinforcement learning (PbRL) learns behaviors from human feedback, without a hand-designed reward function. In practice, existing PbRL methods explore the preference space poorly. They tend to converge early to policies that capture only a narrow slice of what the human wants, and they keep generating similar queries that carry little new information. This also makes the feedback noisier, because people are inconsistent when they have to compare trajectories that look alike.

PB² trains a small population of agents with an explicit diversity mechanism. The population covers more of the preference space while staying aligned with what the human is asking for. This helps in two ways: reward learning improves, and the queries shown to the evaluator are visibly different from each other, so they are easier to judge.

Method overview

PB² keeps a population (3 agents by default):

• an anchor policy π_ref that maximizes only the current learned reward model r_φ. It provides a performance baseline R_ref.
• diverse policies π_i that maximize r_φ + λ · log q_ψ(i | s), that is, the learned reward plus a diversity bonus from a discriminator q_ψ that predicts which agent produced a given state. The bonus is applied only while an agent's return stays within α · R_ref (performance-constrained diversity).

A single reward model r_φ is shared across the population and trained from pairwise preference feedback (Bradley-Terry). Right after each reward-model update the diversity bonus is switched off for a short while, so the agents can first re-adapt to the new preferences. See Algorithm 1 in the paper for the full procedure.

Installation

Tested with Python 3.10 and CUDA 12.1.

conda create -n mjc python=3.10 -y
conda activate mjc
bash setup_env.sh

setup_env.sh installs the pinned dependencies and the vendored custom_dmc2gym package. The bundled rlkit and custom_dmcontrol are imported from the repo root, so they need no install. See requirements.txt for the pip list. Meta-World is optional and is not needed for the paper's experiments.

Quick Start

Train PB² on walker_walk and read the results:

# seed 1, lambda=0.25 (diversity), disc_lr=1e-5, epsilon=0.1 (noisy teacher)
python train_PB2.py env=walker_walk population_size=3 \
    agent.params.beta_init=0.25 disc.lr=1e-5 threshold_ratio=0.9 \
    max_feedback=100 teacher_eps_equal=0.1 seed=1 wandb=false device=cuda

# outputs (returns, logs) are written under ./exp/<env>/...

Runs default to wandb=false. Set wandb=true wandb_project=<p> wandb_entity=<e> to log to Weights & Biases.

Running experiments

Each method has a launch script per environment under run/. PB² scripts take positional arguments:

bash run/<env>/grun_PB2.sh  <seed>  <lambda>  <disc_lr>  <eps_equal>  [<on_policy>]  [<copy_agent>]

• lambda (λ): diversity coefficient. Use 0.25 for locomotion and 0.5 for navigation.
• disc_lr: discriminator learning rate (1e-5).
• eps_equal (ε): similarity threshold for simulated-teacher inconsistency. 0 is a perfect oracle, while 0.05 or 0.1 are more realistic.
• on_policy, copy_agent: on-policy discriminator training and anchor-to-explorer inheritance (both true in the paper).

# PB² and baselines on walker_walk
bash run/walker-walk/grun_PB2.sh    1 0.25 1e-5 0.1 true true   # PB2 (ours)
bash run/walker-walk/grun_QPA.sh    1 0.1                       # QPA
bash run/walker-walk/grun_PEBBLE.sh 1 0.1                       # PEBBLE
bash run/walker-walk/grun_RIME.sh   1 0.1                       # RIME
bash run/walker-walk/grun_TS.sh     1 0.25 1e-5 0.1 true true   # Thompson Sampling (naive DPS)
bash run/walker-walk/rune.sh                                    # RUNE

# navigation (low-feedback)
bash run/grid/grun_PB2.sh       1 0.5 1e-5 0.05                 # 2D Navigation
bash run/point-maze/grun_PB2.sh 1 0.5 1e-5 0.05                 # PointMaze

Available environments: walker_walk, walker_run, cheetah_run, quadruped_walk, ContGridWorld (2D Navigation), PointMaze. Per-method Hydra configs live in config/ (train_PB2.yaml, train_QPA.yaml, train_PEBBLE.yaml, train_RIME.yaml, train_PEBBLE_explore.yaml). TS reuses train_PB2.yaml.

Python API

The training entrypoints are Hydra-driven (train_PB2.py builds everything from config/). If you want to extend the method, the core components can also be constructed directly:

from discriminator import Discriminator
from reward_model_pb2 import PopulationRewardModel
from replay_buffer_diverse import ReplayBuffer

pop_size = 3

# Per-agent replay buffers (performance-constrained via threshold_ratio = alpha)
buffers = [ReplayBuffer(obs_shape, act_shape, capacity, device,
                        max_episode_len=T, threshold_ratio=0.9)
           for _ in range(pop_size)]

# Discriminator q_psi(i | s): predicts which agent generated a given state
disc = Discriminator(state_dim=obs_dim, num_latents=pop_size,
                     hidden_size=256, learning_rate=1e-5,
                     layernorm=False, device=device)

# Shared population reward model r_phi, trained from preference feedback
reward_model = PopulationRewardModel(
    obs_dim, act_dim, device=device,
    ensemble_size=1, size_segment=50, pop_size=pop_size,
    lr=3e-4, mb_size=10, replay_buffer=buffers[0], disc=disc,
)

Policies are SAC agents (agent/sac_diverse.py), one per population member. The diversity bonus λ · log q_ψ(i | s) is added to r_φ during their SAC updates. See train_PB2.py for the full training loop.

Repository structure

train_PB2.py             PB² (main method): anchor + diverse population + discriminator
train_QPA.py             QPA baseline           (Nakamoto et al., 2023)
train_PEBBLE.py          PEBBLE baseline        (Lee et al., 2021)
train_RIME.py            RIME baseline          (Cheng et al., 2024)
train_TS.py              Thompson Sampling / naive DPS baseline (shares PB² infra)
train_PEBBLE_explore.py  RUNE baseline          (uncertainty-driven exploration)

reward_model_pb2.py      population reward model (PB² / TS)
reward_model.py          reward model (QPA / PEBBLE / RUNE)
reward_model_RIME.py     reward model (RIME)
discriminator.py         agent-classification discriminator (diversity bonus)
replay_buffer*.py        per-agent replay buffers
agent/                   SAC variants (sac, sac_diverse, sac_RIME, sac_explore)
hooks/                   Hydra hooks: population reward-model training, logging
config/                  Hydra configs (one per method; agent sub-configs in config/agent/)
run/                     launch scripts grouped by environment
utils.py, dmc.py         environment creation and helpers

custom_dmc2gym/          vendored dm_control to gym adapter (pip install -e)
custom_dmcontrol/        vendored DeepMind Control Suite
rlkit/                   vendored RL utilities (grid-world / box-env wrappers)

Results

The full results and ablations are in the paper. The main takeaways:

• Robustness to noisy feedback (DMControl). As the similarity threshold ε grows from 0 to 0.1, PB² stays robust while the single-agent and naive posterior-sampling baselines drop off. On walker_walk at ε = 0.1, PB² reaches a return of about 750, against roughly 400 for QPA and 350 for PEBBLE.
• Feedback efficiency (navigation). With very little feedback, PB² beats QPA by up to about 50% on 2D Navigation (N from 4 to 8) and by about 20 to 30% on most PointMaze budgets (N from 12 to 20).
• Diversity is what does the work. The Thompson Sampling baseline reuses PB²'s population but swaps the explicit diversity bonus for a vanilla neural ensemble, and it ends up behaving like the purely exploitative QPA. In other words, standard reward ensembles do not capture the reward posterior well on their own.

Numbers are averaged over 10 seeds (DMControl) and 5 seeds (navigation). See the paper for the tables and figures.

Citation

@article{driss2026pb2,
  title   = {{PB\textsuperscript{2}}: Preference Space Exploration via Population-Based Methods in Preference-Based Reinforcement Learning},
  author  = {Driss, Brahim and Davey, Alex and Akrour, Riad},
  journal = {Reinforcement Learning Journal},
  year    = {2026}
}

Preprint: arXiv:2506.13741.

Acknowledgements

This code builds on the official QPA implementation, which itself extends B-Pref / PEBBLE. We thank the authors of those works. The baseline implementations (PEBBLE, RUNE, QPA, RIME) follow their original codebases. Released under the MIT License (see LICENSE).

Contact

For questions, open a GitHub issue or email brahim.driss [at] inria [dot] fr.