refactor(algorithms): centralize algorithm dispatch and capability validation

docs/.vitepress/config.mts

@@ -247,6 +247,7 @@ export default defineConfig({
                 { text: 'Installation', link: '/en/guide/installation' },
                 { text: 'Quick Start', link: '/en/guide/quick-start' },
                 { text: 'Customize Training', link: '/en/guide/customize-training' },
+                { text: 'Adding an Algorithm', link: '/en/guide/adding-an-algorithm' },
                 { text: 'SFT Training', link: '/en/guide/sft-training' },
                 { text: 'MTP Training', link: '/en/guide/mtp-rl-training' },
                 { text: 'PPO Training', link: '/en/guide/ppo-training' },
@@ -363,6 +364,7 @@ export default defineConfig({
                 { text: '安装', link: '/zh/guide/installation' },
                 { text: '快速上手', link: '/zh/guide/quick-start' },
                 { text: '自定义训练', link: '/zh/guide/customize-training' },
+                { text: '接入新算法', link: '/zh/guide/adding-an-algorithm' },
                 { text: 'SFT 训练', link: '/zh/guide/sft-training' },
                 { text: 'MTP 训练', link: '/zh/guide/mtp-rl-training' },
                 { text: 'PPO 训练', link: '/zh/guide/ppo-training' },

docs/en/examples/algorithms.md

@@ -293,23 +293,45 @@ where $r_t = \exp(-\text{KL}_t)$ and $\text{KL}_t = \log\pi_{\theta_\text{old}}(
 
 $$M_2 = \frac{1}{|\mathcal{H}|} \sum_{t \in \mathcal{H}} (\log r_t)^2$$
 
+Relax solves this statistic independently on each Megatron microbatch's local
+tokens. Logging preserves the existing aggregation: with
+`--calculate-per-token-loss`, `train/ppo_kl_m2_before` and
+`train/ppo_kl_m2_after` are weighted by the microbatch's loss-token count;
+otherwise, each microbatch scalar is summed unchanged and the framework
+divides by the sample count. The latter is not a sample-weighted mean. Neither
+mode pools all harmful tokens in the global batch. These logged diagnostics
+do not feed back into the loss.
+
 - If $M_2 \le$ `kl2_budget`: no clipping, all tokens are kept;
-- Otherwise, solve for a trust-region radius $\tau$ by water-filling so the capped second moment returns exactly to budget, i.e. $\sum_{t\in\mathcal{H}} \min\!\left((\log r_t)^2,\ \tau^2\right) = |\mathcal{H}| \cdot \text{kl2\_budget}$, yielding the clip band $[e^{-\tau},\ e^{\tau}]$.
+- Otherwise, use water-filling to select the previous observed breakpoint. This gives a conservative trust-region radius $\tau$ whose capped sum does not exceed $|\mathcal{H}| \cdot \text{kl2\_budget}$, yielding the clip band $[e^{-\tau},\ e^{\tau}]$.
 
-The final clipping margin is $\varepsilon = \max(\text{adaptive value},\ \text{miniclip})$, guaranteeing it is never tighter than GRPO. The policy loss reuses the PPO-Clip pessimistic form from the GRPO section, with only the clip bounds solved adaptively.
+The final clipping margin is $\varepsilon = \max(\text{adaptive value},\ \text{miniclip})$, so it is never tighter than the configured floor. Choose the floors to match the GRPO margins when that is the desired lower bound. `ppo_kl_m2_after` is the legacy solver diagnostic for the selected breakpoint before that floor is applied; in the first-breakpoint edge case it retains the uncapped local mean. It is not a recomputation after the final policy clip. The policy loss reuses the PPO-Clip pessimistic form from the GRPO section, with only the clip bounds solved adaptively.
 
 ### Key Parameters
 
 | Parameter | Default | Recommended | Description |
 |-----------|---------|-------------|-------------|
 | `--advantage-estimator m2po` | — | — | Enable M2PO |
 | `--m2po-kl2-budget` | `0.01` | `0.01`–`0.04` | Second-moment budget per harmful token. Smaller = tighter/more-frequent clipping, larger = more off-policy tolerance (the paper uses `0.04`) |
-| `--m2po-miniclip-low` | `0.3` | `0.2` | Lower clip-margin floor (ratio lower bound is no less than `1 - miniclip_low`) |
+| `--m2po-miniclip-low` | `0.3` | `0.2` | Lower-side clip-margin floor (the margin is at least `miniclip_low`) |
 | `--m2po-miniclip-high` | `0.5` | `0.28` | Upper clip-margin floor |
-| `--use-tis` | off | on | Token Importance Sampling — recommended to enable with M2PO |
+| `--use-rollout-logprobs` | off | on for stale async data | Use the behavior-policy log probabilities that generated each token as M2PO's old policy |
+| `--use-tis` | off | off with rollout log probs | TIS is a post-loss correction and does not drive M2PO's adaptive threshold; it is mutually exclusive with `--use-rollout-logprobs` |
 
 > M2PO derives its clip bounds adaptively, so it does **not** use `--eps-clip` / `--eps-clip-high`.
 
+::: warning Existing implementation behavior
+The registry refactor preserves M2PO's existing computation: reward processing
+passes through raw sample rewards, and the threshold solver does not filter
+tokens by the loss mask or aggregate its statistic across CP ranks. Its
+breakpoint comparisons still use host scalars. Registration does not establish
+equivalence to the paper. Because the clipping statistics are local, M2PO
+declares `supports_context_parallel=False`: startup requires
+`--context-parallel-size 1` and dynamic context parallelism disabled. This
+validation leaves the algorithm's reward, solver and metric calculations
+unchanged.
+:::
+
 ### When to Use
 
 M2PO's benefit grows with how off-policy the training data is, so reach for it **first** in these scenarios:
@@ -320,6 +342,13 @@ M2PO's benefit grows with how off-policy the training data is, so reach for it *
 
 Conversely, under strictly on-policy synchronous training (`--max-staleness 0` with per-step weight sync), M2PO's gain over GRPO is limited — start from GRPO as a baseline there.
 
+M2PO is valid in true-on-policy mode, but then its importance ratio is exactly
+one and adaptive clipping does not engage. To evaluate its stale-data behavior,
+make sure the run supplies old-policy log probabilities rather than
+auto-enabling `--true-on-policy-mode`. For fully-async rollout staleness, pass
+`--use-rollout-logprobs`: `--use-tis` is applied only after M2PO has already
+chosen its clip bounds.
+
 ### Quick Start
 
 Use any existing GRPO training script and replace `GRPO_ARGS` with `M2PO_ARGS`:
@@ -330,7 +359,6 @@ M2PO_ARGS=(
    --m2po-kl2-budget 0.01
    --m2po-miniclip-low 0.2
    --m2po-miniclip-high 0.28
-   --use-tis
 )
 ```
 
@@ -347,7 +375,7 @@ M2PO_ARGS=(
 | **CISPO** | Group-relative reward | Stop-gradient coefficient | Recommended KL loss |
 | **GSPO** | Group-relative reward | PPO-Clip + sequence-level KL | Sequence-level ratio |
 | **SAPO** | Group-relative reward | Sigmoid gate | Temperature-controlled |
-| **M2PO** | Group-relative reward | Adaptive second-moment clip | Optional KL loss (favor for large-staleness / off-policy) |
+| **M2PO** | Raw sample reward broadcast to tokens | Adaptive second-moment clip | Optional KL loss (favor for large-staleness / off-policy) |
 | **RLOO** | Leave-one-out baseline | Unclipped REINFORCE | Optional KL loss (same as GRPO) |
 
 ## Next Steps

docs/en/guide/adding-an-algorithm.md

@@ -0,0 +1,197 @@
+# Adding an Algorithm
+
+Algorithms plug into Relax through the registry under `relax/algorithms/`. An
+algorithm name is no longer scattered across `if/elif` chains — it is described
+by one `AlgorithmSpec`, and each stage looks up what it needs.
+
+## Registry Layout
+
+```
+relax/algorithms/
+├── spec.py        AlgorithmSpec definition + the ALGORITHM_SPECS registry
+├── rewards.py     reward normalization strategies + REWARD_NORMALIZERS
+├── advantages.py  advantage estimators + ADVANTAGE_FNS
+└── policy.py      policy loss adapters + POLICY_LOSS_FNS
+```
+
+Three hard constraints:
+
+1. **No heavy top-level imports under `relax/algorithms/`** — not `megatron`,
+   `ray`, `transfer_queue`, `tensordict`, `relax.components` or
+   `relax.backends`. The registry is imported by argument parsing and by both
+   worker processes; one heavy import drags the whole training stack into
+   `--help` and into a CPU-only CI runner. Import inside the function when you
+   genuinely need one.
+2. **Spec fields hold string identifiers, not callables.** The advantage
+   computation runs in the Ray Serve `Advantages` process while the policy loss
+   runs in the Megatron worker, and those two import different module subsets.
+   Only the algorithm name crosses the process boundary; each side resolves it
+   against its own table.
+3. **Do not hand-edit the `ALGOS` role table.** It is derived from the registry,
+   so a new algorithm gets the standard RL role set automatically.
+
+## How Much Does Adding One Cost
+
+Honestly: **not "one dict entry".**
+
+| Situation | Files to touch |
+|---|---|
+| Reuses existing reward normalization / advantage / policy loss, just combined differently | 1 (`spec.py`) |
+| Needs new maths (a new advantage formula, say) | 2-3 (`spec.py` plus the implementation module) |
+| Also needs new command-line options | 4-6 (the above, plus the option and its validation in `arguments.py`, plus an example and docs) |
+
+What the registry removes is one algorithm name being interpreted in six
+scattered if/elif chains — not the cost of adding an algorithm. An algorithm
+that needs both new maths and new options lands in the last row.
+
+The `ALGOS` role table is the one part that genuinely costs nothing: it derives
+itself from the registry.
+
+## Steps
+
+### 1. Add a spec entry
+
+Edit `ALGORITHM_SPECS` in `relax/algorithms/spec.py`:
+
+```python
+"my_algo": AlgorithmSpec(
+    name="my_algo",
+    reward_normalizer="group_mean_std",   # reuse an existing one, or see step 2
+    requires_complete_reward_groups=True,
+    advantage_fn="grpo_broadcast",
+    policy_loss_fn="ppo_clip",
+),
+```
+
+If your algorithm is identical to an existing one at some stage, reuse that
+identifier. GRPO, GSPO, SAPO, CISPO, M2PO and RLOO all broadcast the scalar
+reward at the advantage layer, so they share `"grpo_broadcast"`. Their reward
+preprocessing differs: M2PO keeps `reward_normalizer="none"` to preserve its
+existing behavior.
+
+Capability fields:
+
+| Field | Effect |
+|-------|--------|
+| `requires_complete_reward_groups` | Preserve complete prompt groups during debug subsampling when reward processing relies on group-level statistics; currently consumed by debug-data selection |
+| `kl_level` | `"token"` or `"sequence"` (GSPO constrains the sequence) |
+| `needs_full_log_probs` | Whether the loss needs CP-gathered full log probs |
+| `supports_context_parallel` | Whether the algorithm's advantage and policy paths support CP-sharded responses; defaults to `True`. `False` rejects static CP sizes other than 1 and enabled dynamic CP at startup |
+| `policy_scalar_metric_names` | Names, in return order, for extra scalar diagnostics produced by the policy adapter |
+| `advantage_normalization` | What `--normalize-advantages` does: `"whiten"` (masked whitening) or `"token_global"` (REINFORCE++'s global token-level normalization, which also switches on the mask-safe loss reducer) |
+| `needs_critic` | Whether a critic service is required; drives `args.use_critic` |
+| `requires_normalize_advantages` | Demand `--normalize-advantages` |
+| `forbids_normalize_advantages` | Reject `--normalize-advantages` (the estimator keeps the advantage's scale on purpose) |
+| `requires_rewards_normalization` | Reject `--disable-rewards-normalization` |
+| `min_group_size` | Floor on `--n-samples-per-prompt` |
+| `forbids_reward_side_kl` | Demand `--kl-coef 0`; there is nowhere to put a reward-side KL term (`--use-kl-loss` is unaffected) |
+| `requires_global_token_loss` | Demand `--calculate-per-token-loss`; the per-sample token-mean reducer would reweight responses by `1 / response_length` |
+| `requires_on_policy_updates` | Rejects five knobs at once: `--fully-async` / `--hybrid`, `--max-staleness != 0`, `--num-steps-per-rollout != 1`, `rollout_batch_size * n_samples != global_batch_size`, and `--partial-rollout` / `--use-dynamic-global-batch-size`. For objectives with no importance-ratio correction |
+
+M2PO and both REINFORCE++ variants declare `supports_context_parallel=False`.
+This describes the whole algorithm, including advantage computation and policy
+loss. It is independent of `requires_complete_reward_groups`, which concerns
+samples sharing a prompt rather than tokens within a response. M2PO keeps
+`requires_complete_reward_groups=False` because its reward stage performs no
+group normalization.
+
+The `validate_*` functions in `relax/utils/arguments.py` consume the startup
+constraint fields. Runtime consumers read the remaining capabilities: reward
+dispatch uses `reward_normalizer`, debug subsampling uses
+`requires_complete_reward_groups`, and the policy path uses the KL,
+normalization, full-log-probability and scalar-metric declarations.
+Declaring an existing capability is enough; do not add an algorithm-name `if`
+to those consumers. A genuinely new enum value or implementation still needs
+one generic handler for that value.
+
+### 2. Write pure functions for genuinely new maths
+
+Only needed when your algorithm differs from every existing one at that stage.
+
+**Reward normalization** (`relax/algorithms/rewards.py`), signature
+`fn(args, samples, raw_rewards) -> list[float]`:
+
+```python
+def normalize_my_strategy(args, samples, raw_rewards):
+    positions_by_group = group_positions(samples, args.n_samples_per_prompt)
+    ...
+    return normalized  # one scalar per sample
+
+REWARD_NORMALIZERS["my_strategy"] = normalize_my_strategy
+```
+
+The output must be **one scalar per sample**. That constraint is what keeps the
+TransferQueue schema fixed — an algorithm reading several reward components
+collapses them to a scalar here.
+
+**Advantage estimator** (`relax/algorithms/advantages.py`), signature
+`fn(args, *, rewards, kl, loss_masks, response_lengths, total_lengths, values)`
+returning `(advantages, returns)`, both `list[Tensor]`:
+
+```python
+def advantage_my_algo(args, *, rewards, kl, **_unused):
+    ...
+    return advantages, returns
+
+ADVANTAGE_FNS["my_algo"] = advantage_my_algo
+```
+
+**Policy loss** (`relax/algorithms/policy.py`), signature
+`fn(args, *, log_probs, ppo_kl, advantages) -> (pg_loss,
+pg_clipfrac, *scalar_metrics)`. The underlying kernels take different argument
+lists; the adapter normalizes them. Most adapters return only the first two
+values. If yours returns scalar diagnostics, declare their names in
+`policy_scalar_metric_names` in the same order. Each diagnostic must contain
+exactly one value. Scalar logging preserves the existing M2PO convention:
+sample mode contributes each microbatch scalar unchanged before the framework
+divides by the sample count; token mode first multiplies it by the microbatch's
+token count. This is not a sample-weighted mean of microbatch diagnostics.
+Real diagnostics are normalized to float32 so one adapter cannot promote the
+distributed logging vector; complex values are rejected.
+
+### 3. Write unit tests
+
+Tests under `tests/algorithms/` need only torch — no megatron, ray or
+transfer_queue:
+
+```bash
+pytest tests/algorithms/ -v
+```
+
+Cover at least:
+
+- Registration and dispatch: the name is in `ALGORITHM_SPECS`, capability fields
+  match expectations, an unregistered name raises.
+- Numerics: hand-compute a small case as the reference. Do not use all-zero or
+  all-equal rewards — every formula returns 0 on those, so the test proves
+  nothing.
+- Degenerate cases: a group where all rewards are equal, boundary values of
+  `n_samples_per_prompt`, missing fields, non-numeric input.
+- **When changing an existing algorithm**: freeze the old implementation into
+  the test file as a reference and compare bit-for-bit
+  (`view(torch.int32).equal`). Do not use `allclose` — its default tolerance is
+  wide enough to swallow the difference between a biased and an unbiased
+  standard deviation. `tests/algorithms/test_reward_normalizers.py` is a
+  worked example.
+
+### 4. Add an example and documentation
+
+- `examples/<algo>/`: a launch script, plus a custom reward function if needed.
+- `docs/{zh,en}/examples/algorithms.md`: how it works, the parameter table, a
+  quick start, and **known deviations** — write down where the implementation
+  differs from the paper rather than leaving users to discover it.
+
+## Arguments
+
+Algorithm-specific options go in `add_algo_arguments` in
+`relax/utils/arguments.py`. The `--advantage-estimator` choices come from
+`list_algorithm_names()`, so registering is enough; there is no name list to
+maintain.
+
+Put cross-argument validation in `validate_algorithm_args`, and prefer
+expressing it through a spec field over comparing algorithm names — the latter
+is exactly what this registry exists to remove.
+
+## References
+
+- [Algorithm Reference](../examples/algorithms.md)

docs/en/guide/configuration.md

@@ -299,7 +299,7 @@ bash scripts/training/text/run-qwen3-4B-fp16-8xgpu.sh \
 
 | Parameter | Type | Default | Options | Description |
 |-----------|------|---------|---------|-------------|
-| `--advantage-estimator` | str | grpo | `grpo`, `gspo`, `reinforce_plus_plus`, `reinforce_plus_plus_baseline`, `ppo`, `sapo`, `cispo` | Advantage estimator. OPD is independent of this choice; enable it with `--use-opd` and its KL/loss coefficient |
+| `--advantage-estimator` | str | grpo | generated from `ALGORITHM_SPECS` in `relax/algorithms/spec.py`; currently `grpo`, `gspo`, `sapo`, `cispo`, `rloo`, `ppo`, `reinforce_plus_plus`, `reinforce_plus_plus_baseline` | Advantage estimator. `--help` is authoritative: the choices are read from the registry, so a new algorithm appears there without this table being edited. OPD is independent of this choice; enable it with `--use-opd` and its KL/loss coefficient |
 | `--normalize-advantages` | flag | False | - | Whether to normalize advantages |
 | `--disable-grpo-std-normalization` | flag | - | - | Disable GRPO standard deviation normalization (from [Dr.GRPO](https://arxiv.org/pdf/2503.20783)) |
 | `--disable-rewards-normalization` | flag | - | - | Disable reward normalization |