diff --git a/docs/algorithms.md b/docs/algorithms.md index 4cdb83a08d..679716839e 100644 --- a/docs/algorithms.md +++ b/docs/algorithms.md @@ -66,7 +66,7 @@ type = "grpo" # the default | `type` | Sampling | Loss | What it is | |---|---|---|---| | `grpo` | policy | `rl` on actions | Standard group-relative RL. | -| `max_rl` | policy | `rl` on actions | MaxRL ([arXiv:2602.02710](https://arxiv.org/abs/2602.02710)): GRPO's centered reward normalized by the group **mean** instead of the standard deviation — the gradient is unbiased for the order-`group_size` truncation of the maximum-likelihood objective, upweighting hard examples like `1/p`. | +| `max_rl` | policy | `rl` on actions | MaxRL ([arXiv:2602.02710](https://arxiv.org/abs/2602.02710)): raw reward for a singleton group (REINFORCE), otherwise GRPO's centered reward normalized by the group **mean** instead of the standard deviation — the gradient is unbiased for the order-`group_size` truncation of the maximum-likelihood objective, upweighting hard examples like `1/p`. | | `opd` | policy | `ref_kl` on actions | On-policy distillation ([Thinking Machines](https://thinkingmachines.ai/blog/on-policy-distillation/)): the policy samples, per-token reverse KL against a reference model as the gradient signal. Needs a `teacher`. | | `sft` | *(the teacher)* | `ce` on actions | Hard distillation: a frozen model generates rollouts, the policy trains with CE on its tokens. Needs a frozen `sampling.source` (the teacher it samples from). | | `opsd` | policy | `ref_kl` on actions | SDFT ([arXiv:2601.19897](https://arxiv.org/abs/2601.19897)): the model is its own reference, conditioned on an expert demonstration. The teacher *is* the live policy (the paper's setting, no extra deployment) — no model to configure. | @@ -106,7 +106,7 @@ kwargs = { patterns = ["WARNING"] } def drop_warnings(rollout, *, patterns: list[str]) -> list[list[bool]]: ... ``` -Component compatibility is validated at config time: frozen-model sampling can only feed the `ce` loss component — the `rl` and `ref_kl` components need the live policy's own sampling logprobs for importance ratios — `opd` pointed at `"policy"` is rejected as degenerate (zero KL), `sft` without a frozen source is rejected (CE on the policy's own tokens is not a distillation target). A group-relative algorithm with `group_size = 1` produces all-zero advantages; the resulting empty batch is caught at runtime (the orchestrator warns and aborts after repeated zero-trainable batches), not at config time. +Component compatibility is validated at config time: frozen-model sampling can only feed the `ce` loss component — the `rl` and `ref_kl` components need the live policy's own sampling logprobs for importance ratios — `opd` pointed at `"policy"` is rejected as degenerate (zero KL), `sft` without a frozen source is rejected (CE on the policy's own tokens is not a distillation target). GRPO-style centering with `group_size = 1` produces all-zero advantages; the resulting empty batch is caught at runtime (the orchestrator warns and aborts after repeated zero-trainable batches), not at config time. MaxRL is the exception: its order-1 estimator uses raw reward to recover REINFORCE. ### Per-Env Algorithms @@ -132,7 +132,7 @@ At runtime, each env's resolved config builds two objects: a `Sampler` (`prime_r |---|---|---| | `grpo` | `GRPOAlgorithm` | `score_group`: group-norm credit (optional length penalty) | | `echo` | `EchoAlgorithm` | `score_rollout`: weighted ce on observation tokens; `score_group`: group-norm credit (inherited) | -| `max_rl` | `MaxRLAlgorithm` | `score_group`: mean-normalized group credit | +| `max_rl` | `MaxRLAlgorithm` | `score_group`: raw singleton reward or mean-normalized group credit | | `opd` | `OPDAlgorithm` | `score_rollout`: own-context prefill under the teacher | | `opsd` | `OPSDAlgorithm` | `score_rollout`: demo-conditioned prefill under the live policy | | `sft` | `SFTDistillAlgorithm` | `score_group`: group-norm credit (feeds filters) | @@ -270,7 +270,7 @@ The per-token training signal is set by `algo.type` and the [algorithm](#the-alg | Type | Component | Effect | |---|---|---| | `grpo` | `rl` | Group-norm: reward minus per-group baseline, optional length penalty. | -| `max_rl` | `rl` | Mean-normalized group credit (maximum-likelihood RL). | +| `max_rl` | `rl` | Raw singleton reward or mean-normalized group credit (maximum-likelihood RL). | | `echo` | `rl` + `ce` | Group-norm on action tokens, plus weighted CE on env-provided tokens selected by message role (each role's `alpha` is its ECHO λ), optionally narrowed by a user filter. | | `opd` | `ref_kl` | On-policy distillation: per-token reverse KL to a reference model (`model`, an inline frozen hosted model), evaluated in the trainer from shipped reference logprobs. No credit — rollouts keep `advantages = None` (advantage-based filters never fire) and ship no advantage stream; `group_size` only fans out sampling. | | `opsd` | `ref_kl` | SDFT: per-token reverse KL to a demo-conditioned reference. No credit — rollouts keep `advantages = None` (advantage-based filters never fire) and ship no advantage stream. | diff --git a/packages/prime-rl-configs/src/prime_rl/configs/algorithm.py b/packages/prime-rl-configs/src/prime_rl/configs/algorithm.py index b3282475af..6b591ff612 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/algorithm.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/algorithm.py @@ -219,15 +219,16 @@ class EchoAlgoConfig(GRPOAlgoConfig): class MaxRLAlgoConfig(BaseAlgoConfig): type: Literal["max_rl"] = "max_rl" - """MaxRL (arXiv:2602.02710): scalar advantage = (reward − group mean) / - group mean, consumed by the ``rl`` loss component. Normalizing by the - mean instead of GRPO's standard deviation makes the policy gradient - unbiased for the order-``group_size`` truncation of the maximum-likelihood - objective: low-pass-rate examples get ~1/p weight, and ``group_size`` is - the truncation order interpolating REINFORCE (1) → exact maximum - likelihood (∞). Designed for non-negative (canonically binary) rewards; - a group with mean reward 0 carries zero advantages everywhere (the - zero-advantage filter drops it, matching the paper's K=0 convention).""" + """MaxRL (arXiv:2602.02710): singleton groups use the raw reward, while + larger groups use scalar advantage = (reward − group mean) / group mean, + consumed by the ``rl`` loss component. Normalizing by the mean instead of + GRPO's standard deviation makes the policy gradient unbiased for the + order-``group_size`` truncation of the maximum-likelihood objective: + low-pass-rate examples get ~1/p weight, and ``group_size`` is the + truncation order interpolating REINFORCE (1) → exact maximum likelihood + (∞). Designed for non-negative (canonically binary) rewards; a group with + mean reward 0 carries zero advantages everywhere (the zero-advantage + filter drops it, matching the paper's K=0 convention).""" action_loss_type: ClassVar[ActionLossType] = "rl" diff --git a/src/prime_rl/orchestrator/algo/max_rl.py b/src/prime_rl/orchestrator/algo/max_rl.py index 9a3978108d..1e8e4bd7b6 100644 --- a/src/prime_rl/orchestrator/algo/max_rl.py +++ b/src/prime_rl/orchestrator/algo/max_rl.py @@ -19,13 +19,18 @@ class MaxRLAlgorithm(Algorithm): ~1/p weight; ``group_size`` interpolates REINFORCE at 1 → exact maximum likelihood as it grows). - Assumes non-negative (canonically binary) rewards; a group with mean reward - <= 0 carries no signal and gets zero advantages (the zero-advantage filter - drops it, matching the paper's no-success convention).""" + A singleton group uses its reward directly, recovering REINFORCE as the + paper requires. Assumes non-negative (canonically binary) rewards; larger + groups with mean reward <= 0 carry no signal and get zero advantages (the + zero-advantage filter drops them, matching the paper's no-success + convention).""" async def score_group(self, group: list[Rollout]) -> None: rewards = torch.tensor([rollout.reward for rollout in group], dtype=torch.float32) - mean = rewards.mean() - advantages = torch.zeros_like(rewards) if mean <= 0 else (rewards - mean) / mean + if len(group) == 1: + advantages = rewards + else: + mean = rewards.mean() + advantages = torch.zeros_like(rewards) if mean <= 0 else (rewards - mean) / mean for rollout, advantage in zip(group, advantages.tolist(), strict=True): rollout.assign_advantages(advantage) diff --git a/tests/unit/orchestrator/test_advantage.py b/tests/unit/orchestrator/test_advantage.py index e65a61e5c8..4adaed1aac 100644 --- a/tests/unit/orchestrator/test_advantage.py +++ b/tests/unit/orchestrator/test_advantage.py @@ -167,6 +167,11 @@ def test_max_rl_mean_normalized(): assert _max_rl(_make_group(rewards=[1.0, 1.0])) == pytest.approx([0.0, 0.0]) +def test_max_rl_singleton_uses_reward_as_reinforce_advantage(): + assert _max_rl(_make_group(rewards=[1.0])) == pytest.approx([1.0]) + assert _max_rl(_make_group(rewards=[0.0])) == pytest.approx([0.0]) + + # -------------------------------------------------------------------------- # GRPO linear length penalty: pass_rate-scaled penalty before the baseline. # --------------------------------------------------------------------------