[No.29] register, verify and document REINFORCE++ / REINFORCE++-baseline

[howtomakeaname]

What

Register the reinforce_plus_plus / reinforce_plus_plus_baseline algorithm variants in ALGOS, add numerical and distributed tests, a design doc, and colocate recipes.

Why

Closes #177.

The advantage functions for REINFORCE++ / REINFORCE++-baseline already existed in ppo_utils.py, but the two variants were never registered in ALGOS, so the Controller rejected the key (Algorithm key 'reinforce_plus_plus' not registered in ALGOS) and they could not run end-to-end. This makes them runnable and verifiable, and pins the formulas, normalization dimensions, mask and reduction semantics in a design doc and tests to resolve the "naming and baseline convention are not unified" gap flagged by task #29.

How

• Registry (relax/core/registry.py): register both variants reusing the GRPO topology (no critic); factor the duplicated GRPO-family topologies into a shared _GRPO_TOPOLOGY constant (single source of truth — registry dispatch instead of if/elif). Behavior-equivalent for existing algorithms.
• Algorithm convention (per review, pinned to the paper arXiv:2501.03262):
  • reinforce_plus_plus (§3.1): per-token k1-style KL penalty (--kl-loss-type k1) folded into the discounted return via --kl-coef; global advantage whitening via --normalize-advantages. No group baseline.
  • reinforce_plus_plus_baseline (§3.2): group-mean baseline (upstream in post_process_rewards, no std division) + global advantage whitening + separate k2 KL loss (--use-kl-loss --kl-loss-type k2 --kl-loss-coef) — the KL penalty is NOT folded into the advantage. get_reinforce_plus_plus_baseline_advantages no longer subtracts -kl_coef * kl (signature drops kl_coef; call sites in loss.py / advantages.py updated); arguments.py now rejects --kl-coef != 0 for the baseline estimator so the "silently no KL" configuration (v1 recipes) cannot recur.
• Tests:
  • tests/utils/training/test_ppo_utils_reinforce.py — element-wise parity vs independent plain-torch reference implementations for the returns, baseline advantages, and the shared compute_policy_loss. Covers variable-length responses, all-zero rewards, single-sample batches, gamma discount, mask not polluting prompt/padding, fully-masked rejection, and the paper convention that per-token KL does NOT enter the baseline advantage.
  • tests/backends/megatron/test_reinforce_pp_cp_parity.py — CP gather/compute/slice wiring of the returns (zig-zag chunking is cp_utils' responsibility, out of scope), baseline CP locality, and DP partition invariance of masked advantage whitening (fake mpu + mocked collectives, single process, per test_ppo_gae_parity.py convention).
  • tests/core/test_registry_reinforce.py — registration, topology reuse, no critic, process_role dispatch.
• Docs (docs/algorithms/reinforce_plus_plus.md): formulas, normalization dimensions, mask and reduction semantics, variant comparison vs GRPO/GSPO/SAPO, CP/DP behavior, and the known async-path whitening limitation. §3.1/§3.2/§4/§5 updated to the paper convention (k1 folded penalty for reinforce_plus_plus; separate k2 KL loss for the baseline variant).
• Recipes: examples/algorithms/run-qwen3-0.6B-1xgpu-reinforce-pp.sh (--kl-coef 0.001 --kl-loss-type k1) and ...-reinforce-pp-baseline.sh (--use-kl-loss --kl-loss-coef 0.001 --kl-loss-type k2), switching --advantage-estimator and adding --normalize-advantages; header step math fixed to 100 × 4 × 8 / 16 = 200.
• Docstrings: make the group-baseline convention explicit (subtracted upstream in post_process_rewards, group-mean without std unlike GRPO; whitening applied only in the sync path).

Testing

pytest tests/utils/training/test_ppo_utils_reinforce.py \
       tests/backends/megatron/test_reinforce_pp_cp_parity.py \
       tests/core/test_registry_reinforce.py -v

• 20 new tests pass; 14 regression tests (test_ppo_utils_grpo.py, test_registry_sft.py) pass. ruff check / ruff format --check / docformatter clean (CI's pinned ruff 0.15.9).

• End-to-end training comparison with non-zero KL per the paper convention (Qwen3-0.6B + GSM8K, same budget, sync colocate, 3 seeds {1234, 42, 7}, 40 steps each, mean±std over stable region steps 5–39):

  algorithm	raw_reward (last)	pg_loss	grad_norm	ppo_kl	kl_loss	clipfrac
  GRPO	0.9271±0.0147	+0.0537±0.0076	1.129±0.052	6.20e-04	1.05e-03 (raw; coef 0)	0.0007
  REINFORCE++	0.9479±0.0295	-0.0943±0.0017	1.736±0.015	6.07e-04	folded into return (Δ returns−reward ≈ −0.001)	0.0014
  REINFORCE++-baseline	0.9688±0.0000	-0.0022±0.0033	1.457±0.053	6.48e-04	1.00e-03 (k2 loss, coef 0.001)	0.0014

  All three stable (9/9 runs 40/40 steps, 0 NaN/Inf, 0 dropped samples); pg_loss / grad_norm / KL metrics are clearly distinct across algorithms, confirming --advantage-estimator actually switches the algorithm. Step time 12.4–13.3 s (samples/s 1.21–1.29, GBS 16) — budget-fair; GRPO reproduces the v1 numbers under the same seed/config (no regression from the changes). All 9 logs and 6 curve figures are included in the task-questionnaire (wps) submission bundle.

• Reproducibility materials (reviewer request): complete per-algorithm/per-seed commands, stats interval (steps 5–39), step time, and the 9 raw logs are listed in the training report (deliverables/training-report.md); logs archive: redai-test/runs/log-{grpo,reinforcepp,reinforcepp-baseline}-seed{1234,42,7}.txt (also bundled as a gist for the review thread). No NaN/Inf, no dropped samples, effective batch/seq-len unchanged — evidence extracted from the logs and tabulated in the report.

Type of Change

• New feature (non-breaking change that adds functionality)
• Documentation update
• Breaking change / Bug fix / Refactoring / Performance / CI

Notes

• No new CLI args; reuses existing --advantage-estimator / --normalize-advantages.
• Advantage whitening is applied only in the sync path; the async path does not whiten (pre-existing behavior shared by all algorithms, documented as a known limitation). Recommend sync colocate mode for REINFORCE++.
• Training comparison covers 3 seeds (1234/42/7) × 3 algorithms × 40 steps with non-zero KL; per-seed std is small (reward ≤ 0.05). Raw logs, run scripts, and the metric extraction script are listed in the report (§4) and bundled as a gist for the review thread.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

[RexFlux]

实验尚不完整，暂不通过，主要有两个问题需要补一下：

1. 两个新增 recipe 都配置了：
   --use-kl-loss
   --kl-loss-coef 0.00
   --kl-loss-type low_var_kl

这实际上等于没有启用 KL。与此同时，文档里又写了普通版本把 KL 放进 token return、baseline 版本把 KL 放进 advantage，所以当前 recipe、文档和算法定义对不上，现有训练结果也没有真正验证文档描述的完整算法。

麻烦先明确这里到底是遵循原论文，还是采用 Relax 自己定义的 convention。如果遵循论文，baseline 版本应该是 group mean + global advantage normalization + 独立 k2 KL loss；如果采用 Relax 自己的定义，也请把实现、
文档、recipe 和测试统一起来，并用非零 KL 系数重新跑实验。

2. Task 29 要求提供和现有算法相同预算下的 reward、loss、KL、吞吐对比，并满足可复现要求。目前 PR 里只有汇总表，吞吐只笼统写了 11～12 秒，原始日志和曲线也没有附上，只写了 “available on request”。建议补充每个算
   法/seed 的完整命令、统计区间、step time 或 samples/s、原始日志或实验链接，并说明没有 NaN/Inf、丢样本或有效 batch 变化。

另外有个小问题：脚本开头写的是 100 × 4 × 4 / 16 = 100 steps，但默认 N_SAMPLES=8，实际应该是 200 steps，也请顺手修正。

其他部分没什么大问题：两个 estimator 注册正确，复用 GRPO 的无 critic topology 也合理；新增和相关回归测试本地共 34 条通过，GitHub CI 也都是绿的。把上面的算法口径和实验材料补齐后，再看是否可以合入。

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

[howtomakeaname]

@RexFlux 感谢审阅，已根据审阅修复，辛苦再Review下～

1. 根据论文（arXiv:2501.03262）统一了实现、文档、recipe 和测试，并用非零 KL 系数重跑（3 算法 × 3 seed × 40 步，同预算）。REINFORCE++ 用 k1-style 惩罚逐 token 折入折扣回报（--kl-coef 0.001 --kl-loss-type k1）；baseline 按论文 §3.2 改为 group-mean + 全局归一化 + 独立 k2 KL loss（--use-kl-loss --kl-loss-type k2 --kl-loss-coef 0.001，KL 不再折入 advantage），并加了 --kl-coef == 0 校验，避免参数写了但没生效。证据：REINFORCE++ 的 returns−raw_reward 三 seed 为 −0.0009 ~ −0.0012（≈−βΣk1，v1 的 kl_coef=0 版差值为 0）；baseline 的 kl_loss 稳定在 0.9 ~ 1.1e-3。

2. 已补齐实验材料，每个算法/seed 的完整命令在运行脚本里（runs/run-*.sh，SEED 注入 --seed），统计区间 step 5–39、3 seed mean±std；step time 12.4 ~ 13.3 s、samples/s 1.21 ~ 1.29（GBS 16 折算）。9 份原始日志、训练曲线和报告已随WPS任务问卷提交的压缩包附上（logs/ 下有原始日志，报告内有每算法/seed 的完整命令，可复算全部指标）。护栏：9 个 run 全部 40/40 步完成、无 NaN/Inf、无丢样本、有效序列长度稳定（1102 ~ 1275）。

同时，步数注释已修正为 100 × 4 × 8 / 16 = 200 steps，注册、拓扑复用、测试未改动，CI（pre-commit / lint / tests）通过。

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.