fix(orchestrator): use raw reward for singleton MaxRL groups

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. |

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"
 

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)

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.
 # --------------------------------------------------------------------------