fix(orchestrator): gate group-relative algorithms to a single trainable agent

docs/algorithms.md

@@ -285,6 +285,8 @@ The per-token training signal is set by `algo.type` and the [algorithm](#the-alg
 
 The default advantage is per-group reward minus per-group baseline (DR-GRPO without std normalization). For each prompt's group of `group_size` rollouts, every token in rollout $i$ receives advantage $s_i - \bar{s}$ where $\bar{s}$ is the group mean.
 
+The group-relative algorithms (`grpo`, `max_rl`, `echo`) require the group's trainable traces to come from **one** agent — the baseline assumes exchangeable attempts by a single agent. Multi-agent envs work as long as exactly one agent is trainable (e.g. a trainable solver rewarded by a frozen agentic judge); a group spanning several trainable agent names is refused at scoring time. Credit assignment across multiple trainable agents is not something prime-rl prescribes — implement your own algorithm for it (see [Authoring an Algorithm](#authoring-an-algorithm)).
+
 This is intentionally simple — it does the right thing for most envs. Write a named algorithm class when you need group-aware shaping that depends on trajectory metadata (sub-agent rollouts, relative-rank shaping, …) — see [Authoring an Algorithm](#authoring-an-algorithm).
 
 A **length penalty** (`length_penalty` on the `grpo`-family algorithms) can be layered on top to discourage rambling. The `linear` penalty subtracts a single `pass_rate`-scaled penalty from each reward before the GRPO baseline, combining output tokens (`num_output_tokens_weight`), input / context tokens (`num_input_tokens_weight`), and turns (`num_turns_weight`) — each normalized by the group's own max for that quantity, with `num_input_tokens_weight` and `num_turns_weight` defaulting to `0.1`.

src/prime_rl/orchestrator/algo/base.py

@@ -78,6 +78,22 @@ async def connect_frozen_pool(
     return pool
 
 
+def assert_single_trainable_agent(group: list[Rollout]) -> None:
+    """Group-relative credit assumes the group is exchangeable attempts by one
+    agent; a cohort spanning several trainable agents would mix different
+    agents' reward scales into one baseline. What the right cross-agent credit
+    assignment is is an open question — refuse instead of guessing."""
+    names = {rollout.agent_name for rollout in group}
+    if len(names) > 1:
+        raise ValueError(
+            f"group for env '{group[0].env_name}' spans trainable agents "
+            f"{sorted(str(name) for name in names)} — group-relative advantages assume a "
+            "single trainable agent. Mark the other agents untrainable in the env's "
+            "setup() (agents.<name>.trainable = False), or implement an algorithm that "
+            "assigns credit across agents."
+        )
+
+
 class Algorithm:
     """Base class for one env's training algorithm — the runtime of the
     algorithm config's per-token training signal (its sibling :class:`Sampler`

src/prime_rl/orchestrator/algo/grpo.py

@@ -5,7 +5,7 @@
 import torch
 
 from prime_rl.configs.algorithm import GRPOAlgoConfig
-from prime_rl.orchestrator.algo.base import Algorithm
+from prime_rl.orchestrator.algo.base import Algorithm, assert_single_trainable_agent
 
 if TYPE_CHECKING:
     from prime_rl.orchestrator.types import Rollout
@@ -22,6 +22,7 @@ def __init__(self, config: GRPOAlgoConfig, policy_pool: InferencePool):
         self.length_penalty = config.length_penalty
 
     async def score_group(self, group: list[Rollout]) -> None:
+        assert_single_trainable_agent(group)
         rewards = torch.tensor([rollout.reward for rollout in group], dtype=torch.float32)
         length_penalty = self.length_penalty
         if length_penalty is None:

src/prime_rl/orchestrator/algo/max_rl.py

@@ -4,7 +4,7 @@
 
 import torch
 
-from prime_rl.orchestrator.algo.base import Algorithm
+from prime_rl.orchestrator.algo.base import Algorithm, assert_single_trainable_agent
 
 if TYPE_CHECKING:
     from prime_rl.orchestrator.types import Rollout
@@ -24,6 +24,7 @@ class MaxRLAlgorithm(Algorithm):
     drops it, matching the paper's no-success convention)."""
 
     async def score_group(self, group: list[Rollout]) -> None:
+        assert_single_trainable_agent(group)
         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

tests/unit/orchestrator/test_advantage.py

@@ -183,6 +183,22 @@ def test_max_rl_mean_normalized():
     assert _max_rl(_make_group(rewards=[1.0, 1.0])) == pytest.approx([0.0, 0.0])
 
 
+def test_group_algos_refuse_multiple_trainable_agents():
+    """One baseline can't span agents: a group mixing trainable agent names is
+    refused; a same-agent multi-trace cohort stays exchangeable and passes."""
+
+    def _agent_rollout(name: str, reward: float) -> Rollout:
+        rollout = _build_rollout(reward, sampled_lengths=[2])
+        rollout.agent = vf.AgentInfo(model="policy", name=name)
+        return rollout
+
+    for drive in (_grpo, _max_rl):
+        with pytest.raises(ValueError, match="single trainable agent"):
+            drive([_agent_rollout("solver", 1.0), _agent_rollout("opponent", 0.0)])
+    advs = _grpo([_agent_rollout("solver", 1.0), _agent_rollout("solver", 0.0)])
+    assert sum(advs) == pytest.approx(0.0, abs=1e-6)
+
+
 # --------------------------------------------------------------------------
 # GRPO linear length penalty: pass_rate-scaled penalty before the baseline.
 # --------------------------------------------------------------------------