feat(algorithm): add PKPO pass@k advantage estimator (arXiv:2505.15201)

[jeffreysijuntan]

What

Adds adv_estimator: pkpo — a per-group reward transform that optimizes pass@k instead of mean reward (Pass@K Policy Optimization, arXiv:2505.15201).

How it works

For each problem's group of n rollouts, every sample is credited with

credit_i = (1/C(n,k)) * sum over k-subsets S containing i of max_{j in S} reward_j

the unbiased estimator of the pass@k / max@k gradient, then mean-centered to form the advantage. Uses the new @rllm.register_adv_estimator hook, so it composes with every backend and downstream loss.

Properties (all covered by tests):

• k=1 → standard mean-centered reward.
• Binary rewards → matches the paper's Theorem 2 (correct → k/n, incorrect → (k/n)·ρ(n−1,c,k−1)); e.g. n=4, c=2, k=2 gives ±1/12.
• Continuous rewards → handled directly (max@k), no thresholding.
• Tied rewards share credit equally; k is clamped to [1, n].

Config

algorithm:
  adv_estimator: pkpo
  pass_at_k: 8      # k to optimize; k=1 == standard pass@1

Changes

• rl_algo.py — calculate_pkpo_advantages_per_group + _comb helper
• advantage.py — register the pkpo estimator (reads algorithm.pass_at_k)
• config.py — PKPO enum member + pass_at_k field (read in from_config)
• base.yaml — pass_at_k key
• tests/unified_trainer/test_pkpo.py — numeric transform tests
• tests/unified_trainer/test_algorithm_config.py — config-resolution tests

Tests

test_pkpo.py + test_algorithm_config.py pass (16). Verified end-to-end through the registry + AlgorithmConfig (adv_estimator=pkpo, pass_at_k=2). test_config_sync.py errors locally only because verl isn't installed in this env (pre-existing, unrelated).

Base branch

Targets terminal-rl (builds on @rllm.register_adv_estimator from #742).

🤖 Generated with Claude Code