From b0983bc8bb59e4755d1b8516032259da76eab7a1 Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Thu, 30 Jul 2026 19:14:03 +0000 Subject: [PATCH 01/58] feat(orchestrator): episode- and agent-level rollout metrics Adapt to verifiers#2187 (Episode derived aggregates): episode.last_error rename, submodule bump. Mirror the episode/trace hierarchy in the wandb layout: count metrics read at the episode level (per-episode sums), and a new {scope}/{subset}/// level reports trace-level metrics per agent, averaging in-episode fan-outs first. Co-Authored-By: Claude Fable 5 --- deps/verifiers | 2 +- skills/training/monitor-run/SKILL.md | 7 ++- src/prime_rl/orchestrator/envs.py | 4 +- src/prime_rl/orchestrator/metrics.py | 80 +++++++++++++++++++++---- tests/unit/orchestrator/test_metrics.py | 39 +++++++++++- 5 files changed, 113 insertions(+), 19 deletions(-) diff --git a/deps/verifiers b/deps/verifiers index f646beb37e..82919fc115 160000 --- a/deps/verifiers +++ b/deps/verifiers @@ -1 +1 @@ -Subproject commit f646beb37eef51869f886f244456e3d07818e4d6 +Subproject commit 82919fc11512e447d2cb304ff0ed2ab0d888fb7c diff --git a/skills/training/monitor-run/SKILL.md b/skills/training/monitor-run/SKILL.md index 45d9abf92f..ae9224c430 100644 --- a/skills/training/monitor-run/SKILL.md +++ b/skills/training/monitor-run/SKILL.md @@ -89,13 +89,14 @@ grep -E "WARNING|ERROR" {output_dir}/logs/envs/{train,eval}/*.log All metrics print to the console log (and W&B when configured). -**Progress** — orchestrator log. Rollout metrics are keyed `{scope}/{subset}//`: `scope` is `train/agg` (all train envs) or `train/` (`eval/` for eval); `subset` is `all` (every rollout) or `effective` (post-filter). +**Progress** — orchestrator log. Rollout metrics are keyed `{scope}/{subset}//` (episode-level: the token/turn/branch counts sum an episode's traces) and `{scope}/{subset}///` (trace-level, per agent name — an in-episode fan-out like n solvers averages within the episode first): `scope` is `train/agg` (all train envs) or `train/` (`eval/` for eval); `subset` is `all` (every rollout) or `effective` (post-filter). Single-agent envs have one trace per episode, so both levels agree. | Metric | Description | |--------|-------------| | `train/agg/effective/reward/mean` | mean training reward (per env: `train//effective/reward/mean`) | -| `train/agg/effective/num_total_tokens/mean` | avg tokens per rollout (also `num_input_tokens`, `num_output_tokens`) | -| `train/agg/effective/num_turns/mean` | avg turns per rollout (multi-turn only) | +| `train/agg/effective/num_total_tokens/mean` | avg tokens per episode (also `num_input_tokens`, `num_output_tokens`) | +| `train/agg/effective/num_turns/mean` | avg turns per episode (multi-turn only) | +| `train//effective//num_turns/mean` | per-agent avg turns (also `reward`, token counts, `is_truncated`) | | `train/agg/effective/is_truncated/mean` | fraction truncated | | `train/agg/all/has_error/mean` | fraction errored (per-type under `train/agg/all/error/`; also `dispatcher/errored/{train,eval}`) | | `train//effective/metrics//mean` | env-specific metrics (e.g. pass rate) | diff --git a/src/prime_rl/orchestrator/envs.py b/src/prime_rl/orchestrator/envs.py index 3bffb3504d..311249c3a6 100644 --- a/src/prime_rl/orchestrator/envs.py +++ b/src/prime_rl/orchestrator/envs.py @@ -217,14 +217,14 @@ async def run( sampling=self._sampling(cache_salt), ) if not episode.traces: - error = episode.error + error = episode.last_error detail = f"{error.type}: {error.message}" if error is not None else "no traces and no error recorded" raise RuntimeError(f"env-rollout failed before any trace was minted — {detail}") rollouts = [ROLLOUT_TYPE.model_construct(**dict(wire)) for wire in episode.traces] for rollout in rollouts: rollout.episode_id = episode.id if not episode.ok and rollout.ok: - error = episode.error or vf.Error( + error = episode.last_error or vf.Error( type="EpisodeFailed", message="A sibling trace in this episode failed" ) rollout.errors = [*rollout.errors, error] diff --git a/src/prime_rl/orchestrator/metrics.py b/src/prime_rl/orchestrator/metrics.py index ce860cb284..ce245b534d 100644 --- a/src/prime_rl/orchestrator/metrics.py +++ b/src/prime_rl/orchestrator/metrics.py @@ -6,8 +6,14 @@ ``rollouts.metrics.num_input_tokens.mean()`` works — and assembles the full ``{prefix}/{subset}//`` wandb dict via ``.to_wandb(...)``. -No I/O, no pandas — plain Python over the ``vf.Trace`` properties each rollout exposes. Aggregation -is flat over the rollout list except the solve rates, which group by ``group_id``. +The wandb layout mirrors the episode/trace hierarchy: ``{prefix}/{subset}//`` reads at +the episode level (the count metrics sum an episode's traces, matching ``vf.Episode``'s aggregates) +and ``{prefix}/{subset}///`` at the trace level, grouped by agent name +(``vf.Episode.by_agent``). A single-agent env has one trace per episode, so both levels coincide. + +No I/O, no pandas — plain Python over the ``vf.Trace`` properties each rollout exposes. The count +metrics group by ``episode_id`` and the solve rates by ``group_id``; everything else is flat over +the rollout list. """ from __future__ import annotations @@ -152,34 +158,86 @@ def stats(self) -> dict[str, Stat]: } +class AgentMetrics(StatGroup): + """Trace-level metrics for one agent, one value per episode: a fan-out (n same-agent traces + in one episode, e.g. n solvers) collapses to its within-episode mean first, so + ``/num_turns/mean`` is the mean over episodes of the agent's mean turns.""" + + DISTRIBUTIONS = ("reward", "num_total_tokens", "num_input_tokens", "num_output_tokens", "num_turns", "num_branches") + RATES = ("is_truncated", "is_completed") + + def __init__(self, episodes: list[list[Rollout]]) -> None: + super().__init__([r for episode in episodes for r in episode]) + self.episodes = episodes + + def stats(self) -> dict[str, Stat]: + return { + name: Stat([sum(float(getattr(r, name)) for r in episode) / len(episode) for episode in self.episodes]) + for name in (*self.DISTRIBUTIONS, *self.RATES) + } + + def to_dict(self, prefix: str) -> dict[str, float]: + """Full ```` fan-out for the distributions; ``/mean`` only for the 0/1 rates.""" + stats = self.stats() + out: dict[str, float] = {} + for name in self.DISTRIBUTIONS: + out |= stats[name].to_dict(f"{prefix}/{name}") + for name in self.RATES: + out[f"{prefix}/{name}/mean"] = stats[name].mean() + return out + + class RolloutMetrics: """Metrics shared by train and eval over a rollout list. Distributional metrics are ``Stat``s - (mean/max/min); boolean metrics are ``Stat``s of 0/1 (use ``.mean()`` for the rate). ``to_wandb`` - assembles the full ``{prefix}/{subset}/...`` dict; ``TrainMetrics`` / ``EvalMetrics`` extend it.""" + (mean/max/min); boolean metrics are ``Stat``s of 0/1 (use ``.mean()`` for the rate). The count + metrics (tokens/turns/branches) are episode-level — one value per episode, summing its traces — + with the per-agent trace-level view under ``by_agent()``. ``to_wandb`` assembles the full + ``{prefix}/{subset}/...`` dict; ``TrainMetrics`` / ``EvalMetrics`` extend it.""" def __init__(self, rollouts: list[Rollout]) -> None: self.rollouts = rollouts - # Distributional metrics + def episodes(self) -> list[list[Rollout]]: + """The subset's rollouts grouped into their episodes. A rollout without an + ``episode_id`` (legacy envs, synthesized error markers) is its own episode.""" + grouped: dict[str, list[Rollout]] = {} + for r in self.rollouts: + grouped.setdefault(r.episode_id or r.id, []).append(r) + return list(grouped.values()) + + def by_agent(self) -> dict[str, AgentMetrics]: + """Per-agent metric views (``vf.Episode.by_agent`` over the subset's rollouts): each + episode contributes the agent's traces in it.""" + per_agent: dict[str, list[list[Rollout]]] = {} + for episode in self.episodes(): + traces: dict[str, list[Rollout]] = {} + for r in episode: + traces.setdefault(r.agent.name, []).append(r) + for name, agent_traces in traces.items(): + per_agent.setdefault(name, []).append(agent_traces) + return {name: AgentMetrics(episodes) for name, episodes in sorted(per_agent.items())} + + # Episode-level count metrics: one value per episode, summing its traces — the same + # aggregation as ``vf.Episode.num_turns`` / ``num_*_tokens``. @property def num_total_tokens(self) -> Stat: - return Stat([float(r.num_total_tokens) for r in self.rollouts]) + return Stat([float(sum(r.num_total_tokens for r in episode)) for episode in self.episodes()]) @property def num_input_tokens(self) -> Stat: - return Stat([float(r.num_input_tokens) for r in self.rollouts]) + return Stat([float(sum(r.num_input_tokens for r in episode)) for episode in self.episodes()]) @property def num_output_tokens(self) -> Stat: - return Stat([float(r.num_output_tokens) for r in self.rollouts]) + return Stat([float(sum(r.num_output_tokens for r in episode)) for episode in self.episodes()]) @property def num_turns(self) -> Stat: - return Stat([float(r.num_turns) for r in self.rollouts]) + return Stat([float(sum(r.num_turns for r in episode)) for episode in self.episodes()]) @property def num_branches(self) -> Stat: - return Stat([float(r.num_branches) for r in self.rollouts]) + return Stat([float(sum(r.num_branches for r in episode)) for episode in self.episodes()]) @property def timing(self) -> TimingMetrics: @@ -259,6 +317,8 @@ def to_wandb(self, *, prefix: str, subset: Subset) -> dict[str, float]: out |= self.timing.to_dict(f"{p}/timing") out |= self.metrics.to_dict(f"{p}/metrics") out |= self.rewards.to_dict(f"{p}/rewards") + for agent, agent_metrics in self.by_agent().items(): + out |= agent_metrics.to_dict(f"{p}/{agent}") out[f"{p}/is_truncated/mean"] = self.is_truncated.mean() out[f"{p}/is_completed/mean"] = self.is_completed.mean() # errors live only on the `all` subset (effective drops them), so emit the rate + the diff --git a/tests/unit/orchestrator/test_metrics.py b/tests/unit/orchestrator/test_metrics.py index 5362037833..9a47742490 100644 --- a/tests/unit/orchestrator/test_metrics.py +++ b/tests/unit/orchestrator/test_metrics.py @@ -1,4 +1,5 @@ import math +from itertools import count from types import SimpleNamespace import pytest @@ -7,10 +8,14 @@ from prime_rl.orchestrator.metrics import EvalRollouts, Stat, TrainRollouts from prime_rl.orchestrator.utils import compute_pass_metrics +_ids = count() + def mk( reward: float = 0.0, *, + episode_id: str = "", + agent_name: str = "agent", num_total_tokens: int = 10, num_input_tokens: int = 4, num_output_tokens: int = 6, @@ -36,8 +41,12 @@ def mk( finalize: float = 0.0, scoring: float = 0.0, ): - """Duck-typed stand-in for ``Rollout``, exposing only the Trace properties the metrics read.""" + """Duck-typed stand-in for ``Rollout``, exposing only the Trace properties the metrics read. + Without an ``episode_id`` each rollout is its own single-trace episode (the unique ``id`` + is the grouping fallback).""" return SimpleNamespace( + id=f"t{next(_ids)}", + episode_id=episode_id, reward=reward, rewards=rewards or {}, num_total_tokens=num_total_tokens, @@ -53,7 +62,7 @@ def mk( metrics=metrics or {}, env_name=env_name, group_id=group_id, - agent=SimpleNamespace(trainable=trainable), + agent=SimpleNamespace(trainable=trainable, name=agent_name), is_trainable=is_trainable, is_filtered=is_filtered, filter_results=filter_results or {}, @@ -107,11 +116,35 @@ def test_to_wandb_distributions(): out = m.to_wandb(prefix="train/agg", subset="all") assert out["train/agg/all/reward/mean"] == 0.5 assert out["train/agg/all/num_total_tokens/mean"] == 15.0 - assert out["train/agg/all/num_total_tokens/max"] == 20.0 # flat over rollouts, not per-group + assert out["train/agg/all/num_total_tokens/max"] == 20.0 # single-trace episodes: one value per rollout assert out["train/agg/all/num_input_tokens/mean"] == 5.0 assert out["train/agg/all/num_output_tokens/mean"] == 6.0 +def test_episode_and_agent_levels(): + # Two proposer-solver episodes: one proposer + two solvers each (the solver fan-out). + rollouts = [ + mk(reward=1.0, num_turns=1, agent_name="proposer", episode_id="e1"), + mk(reward=0.0, num_turns=2, agent_name="solver", episode_id="e1"), + mk(reward=1.0, num_turns=4, agent_name="solver", episode_id="e1"), + mk(reward=0.0, num_turns=3, agent_name="proposer", episode_id="e2"), + mk(reward=1.0, num_turns=6, agent_name="solver", episode_id="e2"), + mk(reward=0.0, num_turns=8, agent_name="solver", episode_id="e2"), + ] + m = TrainRollouts(rollouts).metrics + assert m.num_turns.mean() == 12.0 # episode-level sums: 1+2+4 and 3+6+8 + assert m.num_total_tokens.values == [30.0, 30.0] # summed across the episode's traces + out = m.to_wandb(prefix="train/agg", subset="all") + assert out["train/agg/all/num_turns/mean"] == 12.0 + assert out["train/agg/all/proposer/num_turns/mean"] == 2.0 # one trace per episode: (1 + 3) / 2 + assert out["train/agg/all/solver/num_turns/mean"] == 5.0 # mean of per-episode fan-out means (3, 7) + assert out["train/agg/all/solver/num_turns/max"] == 7.0 + assert out["train/agg/all/solver/reward/mean"] == 0.5 + assert out["train/agg/all/proposer/is_truncated/mean"] == 0.0 + assert "train/agg/all/proposer/is_truncated/p90" not in out # rates emit /mean only + assert out["train/agg/all/reward/mean"] == 0.5 # trace-level metrics stay flat over rollouts + + def test_boolean_rates_and_error_breakdown_all_only(): rc = TrainRollouts([mk(is_truncated=True), mk(has_error=True, error_type="ProviderError"), mk(is_filtered=True)]) out = rc.metrics.to_wandb(prefix="train/agg", subset="all") From de1e25e8bb83b344f870147154e2d50aa212bebf Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Thu, 30 Jul 2026 21:23:19 +0000 Subject: [PATCH 02/58] chore: bump verifiers to vf#2187 head, episode wording cleanups Co-Authored-By: Claude Fable 5 --- deps/verifiers | 2 +- skills/training/monitor-run/SKILL.md | 2 +- src/prime_rl/orchestrator/dispatcher.py | 6 +++--- src/prime_rl/orchestrator/envs.py | 6 +++--- 4 files changed, 8 insertions(+), 8 deletions(-) diff --git a/deps/verifiers b/deps/verifiers index 82919fc115..7b8c5ad68c 160000 --- a/deps/verifiers +++ b/deps/verifiers @@ -1 +1 @@ -Subproject commit 82919fc11512e447d2cb304ff0ed2ab0d888fb7c +Subproject commit 7b8c5ad68cf3af5fc741cd54f7a994585a4b6ea9 diff --git a/skills/training/monitor-run/SKILL.md b/skills/training/monitor-run/SKILL.md index ae9224c430..48eedb47ea 100644 --- a/skills/training/monitor-run/SKILL.md +++ b/skills/training/monitor-run/SKILL.md @@ -95,7 +95,7 @@ All metrics print to the console log (and W&B when configured). |--------|-------------| | `train/agg/effective/reward/mean` | mean training reward (per env: `train//effective/reward/mean`) | | `train/agg/effective/num_total_tokens/mean` | avg tokens per episode (also `num_input_tokens`, `num_output_tokens`) | -| `train/agg/effective/num_turns/mean` | avg turns per episode (multi-turn only) | +| `train/agg/effective/num_turns/mean` | avg turns per episode | | `train//effective//num_turns/mean` | per-agent avg turns (also `reward`, token counts, `is_truncated`) | | `train/agg/effective/is_truncated/mean` | fraction truncated | | `train/agg/all/has_error/mean` | fraction errored (per-type under `train/agg/all/error/`; also `dispatcher/errored/{train,eval}`) | diff --git a/src/prime_rl/orchestrator/dispatcher.py b/src/prime_rl/orchestrator/dispatcher.py index 260a6cdd38..5fee725f6a 100644 --- a/src/prime_rl/orchestrator/dispatcher.py +++ b/src/prime_rl/orchestrator/dispatcher.py @@ -5,8 +5,8 @@ N rollouts in one call reserves N permits (each bridged v0 rollout is its own single-agent episode). - Optional rate limiting via ``AsyncLimiter(tasks_per_minute, 60)``. -- Emit-everything invariant: every dispatched env-rollout eventually reaches - ``out_q`` exactly once, as one episode (a ``list[Rollout]``). Failures +- Emit-everything invariant: every dispatched episode eventually reaches + ``out_q`` exactly once, as a ``list[Rollout]``. Failures (env error, empty trajectory, task exception, off-policy cancel) carry ``trace.last_error`` set; sinks decide drop / partial-train policy. - ``DispatcherMode.PREFER_TRAIN`` / ``PREFER_EVAL`` controls which kind to @@ -495,7 +495,7 @@ def release(self, n: int) -> None: self.inflight_permits -= n async def handle_completed_rollout(self, task: asyncio.Task) -> None: - """Emit every dispatched env-rollout exactly once to ``out_q``: a ``run`` + """Emit every dispatched episode exactly once to ``out_q``: a ``run`` result as one episode, a legacy ``run_group`` result as ``rollout_count`` single-trace episodes. Task exceptions synthesize ``rollout_count`` error-marker episodes so the sink's count-to-``group_size`` finalization diff --git a/src/prime_rl/orchestrator/envs.py b/src/prime_rl/orchestrator/envs.py index 311249c3a6..25baafdb00 100644 --- a/src/prime_rl/orchestrator/envs.py +++ b/src/prime_rl/orchestrator/envs.py @@ -4,7 +4,7 @@ external one pinned by ``config.serve.address``) and an ``EnvClient`` to drive it. The orchestrator never *runs* an environment — the agents and their runtimes live only in the server — but it does own the *taskset*: a v1 env's tasks are loaded here, -once, and each dispatched env-rollout ships its task's data on the request +once, and each dispatched episode ships its task's data on the request (``task_data``); the server pydantic-validates it into the taskset's declared ``TaskData`` type and runs it. That keeps the server (and every worker in its pool) stateless about data — no per-worker dataset loads, no idx-addressed task @@ -12,7 +12,7 @@ the legacy (v0) bridge, whose dataset genuinely lives server-side, is still driven by ``task_idx`` (its count comes from ``info``). -The server answers one ``Episode`` per env-rollout, whose traces we validate into +The server answers one ``Episode`` per run request, whose traces we validate into ``Trace[WireTaskData]`` — real ``vf.Trace``\\ s (never loose dicts) whose task keeps the env's task-specific fields as extras (``WireTaskData`` allows them). """ @@ -219,7 +219,7 @@ async def run( if not episode.traces: error = episode.last_error detail = f"{error.type}: {error.message}" if error is not None else "no traces and no error recorded" - raise RuntimeError(f"env-rollout failed before any trace was minted — {detail}") + raise RuntimeError(f"episode failed before any trace was produced — {detail}") rollouts = [ROLLOUT_TYPE.model_construct(**dict(wire)) for wire in episode.traces] for rollout in rollouts: rollout.episode_id = episode.id From 840dc2d2469f6df672cb605e481b5a760930d82f Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Thu, 30 Jul 2026 21:24:23 +0000 Subject: [PATCH 03/58] refactor(orchestrator): RolloutMetrics -> EpisodeMetrics, AgentMetrics -> TraceMetrics Co-Authored-By: Claude Fable 5 --- src/prime_rl/orchestrator/metrics.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/src/prime_rl/orchestrator/metrics.py b/src/prime_rl/orchestrator/metrics.py index ce245b534d..ae12e895ba 100644 --- a/src/prime_rl/orchestrator/metrics.py +++ b/src/prime_rl/orchestrator/metrics.py @@ -158,7 +158,7 @@ def stats(self) -> dict[str, Stat]: } -class AgentMetrics(StatGroup): +class TraceMetrics(StatGroup): """Trace-level metrics for one agent, one value per episode: a fan-out (n same-agent traces in one episode, e.g. n solvers) collapses to its within-episode mean first, so ``/num_turns/mean`` is the mean over episodes of the agent's mean turns.""" @@ -187,7 +187,7 @@ def to_dict(self, prefix: str) -> dict[str, float]: return out -class RolloutMetrics: +class EpisodeMetrics: """Metrics shared by train and eval over a rollout list. Distributional metrics are ``Stat``s (mean/max/min); boolean metrics are ``Stat``s of 0/1 (use ``.mean()`` for the rate). The count metrics (tokens/turns/branches) are episode-level — one value per episode, summing its traces — @@ -205,7 +205,7 @@ def episodes(self) -> list[list[Rollout]]: grouped.setdefault(r.episode_id or r.id, []).append(r) return list(grouped.values()) - def by_agent(self) -> dict[str, AgentMetrics]: + def by_agent(self) -> dict[str, TraceMetrics]: """Per-agent metric views (``vf.Episode.by_agent`` over the subset's rollouts): each episode contributes the agent's traces in it.""" per_agent: dict[str, list[list[Rollout]]] = {} @@ -215,7 +215,7 @@ def by_agent(self) -> dict[str, AgentMetrics]: traces.setdefault(r.agent.name, []).append(r) for name, agent_traces in traces.items(): per_agent.setdefault(name, []).append(agent_traces) - return {name: AgentMetrics(episodes) for name, episodes in sorted(per_agent.items())} + return {name: TraceMetrics(episodes) for name, episodes in sorted(per_agent.items())} # Episode-level count metrics: one value per episode, summing its traces — the same # aggregation as ``vf.Episode.num_turns`` / ``num_*_tokens``. @@ -331,7 +331,7 @@ def to_wandb(self, *, prefix: str, subset: Subset) -> dict[str, float]: return out -class TrainMetrics(RolloutMetrics): +class TrainMetrics(EpisodeMetrics): """Common metrics plus the reward distribution and filter-pipeline rates.""" @property @@ -365,7 +365,7 @@ def to_wandb(self, *, prefix: str, subset: Subset) -> dict[str, float]: return out -class EvalMetrics(RolloutMetrics): +class EvalMetrics(EpisodeMetrics): """Common metrics plus the ``avg@`` score and (on the effective subset, for binary-reward tasks) pass@k / pass^k. ``group_size`` (the ``avg@k`` k) is supplied by the container so the ``all`` and ``effective`` subsets share one stable key.""" From 9520eab8bc8a1a348f61e69f69218b22a9a85cf8 Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Thu, 30 Jul 2026 21:43:58 +0000 Subject: [PATCH 04/58] feat(orchestrator)!: trace-level metrics are agent-only, timing generation span renamed agent Env level now carries only episode-level facts (count sums, pipeline rates, eval scores); reward, truncation, errors, stop conditions, timing, custom metrics and solve rates move under {scope}/{subset}//. Adapts to the vf span rename (Timing.agent, AgentSpan, split_agent_time) and re-pins deps/verifiers. Overview reward/error/truncation panels match the per-agent keys by regex. Co-Authored-By: Claude Fable 5 --- deps/verifiers | 2 +- skills/training/monitor-run/SKILL.md | 13 +- src/prime_rl/orchestrator/metrics.py | 202 ++++++++++++------------ src/prime_rl/utils/monitor/wandb.py | 29 ++-- tests/unit/orchestrator/test_metrics.py | 63 ++++---- 5 files changed, 164 insertions(+), 145 deletions(-) diff --git a/deps/verifiers b/deps/verifiers index 7b8c5ad68c..515b820490 160000 --- a/deps/verifiers +++ b/deps/verifiers @@ -1 +1 @@ -Subproject commit 7b8c5ad68cf3af5fc741cd54f7a994585a4b6ea9 +Subproject commit 515b8204903afb2fc41f6bd6aaaafa2af9c3e3d1 diff --git a/skills/training/monitor-run/SKILL.md b/skills/training/monitor-run/SKILL.md index 48eedb47ea..692ed08ab3 100644 --- a/skills/training/monitor-run/SKILL.md +++ b/skills/training/monitor-run/SKILL.md @@ -89,17 +89,18 @@ grep -E "WARNING|ERROR" {output_dir}/logs/envs/{train,eval}/*.log All metrics print to the console log (and W&B when configured). -**Progress** — orchestrator log. Rollout metrics are keyed `{scope}/{subset}//` (episode-level: the token/turn/branch counts sum an episode's traces) and `{scope}/{subset}///` (trace-level, per agent name — an in-episode fan-out like n solvers averages within the episode first): `scope` is `train/agg` (all train envs) or `train/` (`eval/` for eval); `subset` is `all` (every rollout) or `effective` (post-filter). Single-agent envs have one trace per episode, so both levels agree. +**Progress** — orchestrator log. Rollout metrics mirror the episode/trace hierarchy: `{scope}/{subset}//` carries episode-level facts only (token/turn/branch counts summed over an episode's traces, plus train pipeline rates and eval scores); every trace-level metric (reward, truncation, errors, timing, env metrics) lives under `{scope}/{subset}///`, per agent name — an in-episode fan-out like n solvers averages within the episode first. `scope` is `train/agg` (all train envs) or `train/` (`eval/` for eval); `subset` is `all` (every rollout) or `effective` (post-filter). Single-agent envs have one agent (usually `agent`) and one trace per episode. | Metric | Description | |--------|-------------| -| `train/agg/effective/reward/mean` | mean training reward (per env: `train//effective/reward/mean`) | +| `train/agg/effective//reward/mean` | mean training reward (per env: `train//effective//reward/mean`) | | `train/agg/effective/num_total_tokens/mean` | avg tokens per episode (also `num_input_tokens`, `num_output_tokens`) | | `train/agg/effective/num_turns/mean` | avg turns per episode | -| `train//effective//num_turns/mean` | per-agent avg turns (also `reward`, token counts, `is_truncated`) | -| `train/agg/effective/is_truncated/mean` | fraction truncated | -| `train/agg/all/has_error/mean` | fraction errored (per-type under `train/agg/all/error/`; also `dispatcher/errored/{train,eval}`) | -| `train//effective/metrics//mean` | env-specific metrics (e.g. pass rate) | +| `train//effective//num_turns/mean` | per-agent avg turns (also token counts, `num_branches`) | +| `train/agg/effective//is_truncated/mean` | fraction truncated | +| `train/agg/all//has_error/mean` | fraction errored (per-type under `train/agg/all//error/`; also `dispatcher/errored/{train,eval}`) | +| `train//effective//metrics//mean` | env-specific metrics (e.g. pass rate) | +| `train//effective//timing/agent/model/mean` | model vs harness share of the agent phase | | `eval//effective/{avg@k,pass@k}` | eval scores when configured | **Stability** — trainer log: diff --git a/src/prime_rl/orchestrator/metrics.py b/src/prime_rl/orchestrator/metrics.py index ae12e895ba..656a19c22f 100644 --- a/src/prime_rl/orchestrator/metrics.py +++ b/src/prime_rl/orchestrator/metrics.py @@ -6,14 +6,14 @@ ``rollouts.metrics.num_input_tokens.mean()`` works — and assembles the full ``{prefix}/{subset}//`` wandb dict via ``.to_wandb(...)``. -The wandb layout mirrors the episode/trace hierarchy: ``{prefix}/{subset}//`` reads at -the episode level (the count metrics sum an episode's traces, matching ``vf.Episode``'s aggregates) -and ``{prefix}/{subset}///`` at the trace level, grouped by agent name -(``vf.Episode.by_agent``). A single-agent env has one trace per episode, so both levels coincide. - -No I/O, no pandas — plain Python over the ``vf.Trace`` properties each rollout exposes. The count -metrics group by ``episode_id`` and the solve rates by ``group_id``; everything else is flat over -the rollout list. +The wandb layout mirrors the episode/trace hierarchy. ``{prefix}/{subset}//`` carries +only episode-level facts: the count metrics sum an episode's traces (matching ``vf.Episode``'s +aggregates), plus the train pipeline rates and the eval scores. Every trace-level metric lives under +``{prefix}/{subset}///``, grouped by agent name (``vf.Episode.by_agent``) so +agents never mix into one distribution. A single-agent env has one trace per episode, so the count +metrics coincide across levels. + +No I/O, no pandas — plain Python over the ``vf.Trace`` properties each rollout exposes. """ from __future__ import annotations @@ -96,26 +96,26 @@ class TimingMetrics(StatGroup): """Per-phase rollout durations, nested so ``metrics.timing.setup.mean()`` reads naturally. ``total`` is the per-rollout sum across all phases.""" - PHASES = ("setup", "generation", "finalize", "scoring") + PHASES = ("setup", "agent", "finalize", "scoring") @property def setup(self) -> Stat: return Stat([r.timing.setup.duration for r in self.rollouts]) @property - def generation(self) -> Stat: - return Stat([r.timing.generation.duration for r in self.rollouts]) + def agent(self) -> Stat: + return Stat([r.timing.agent.duration for r in self.rollouts]) @property - def generation_model(self) -> Stat: - """The share of the generation phase spent inside model calls (inference).""" - return Stat([r.timing.generation.model.duration for r in self.rollouts]) + def agent_model(self) -> Stat: + """The share of the agent phase spent inside model calls (inference).""" + return Stat([r.timing.agent.model.duration for r in self.rollouts]) @property - def generation_harness(self) -> Stat: - """The share of the generation phase spent outside model calls (harness, tools, + def agent_harness(self) -> Stat: + """The share of the agent phase spent outside model calls (harness, tools, user simulation).""" - return Stat([r.timing.generation.harness.duration for r in self.rollouts]) + return Stat([r.timing.agent.harness.duration for r in self.rollouts]) @property def finalize(self) -> Stat: @@ -132,8 +132,8 @@ def total(self) -> Stat: def stats(self) -> dict[str, Stat]: return { **{phase: getattr(self, phase) for phase in self.PHASES}, - "generation/model": self.generation_model, - "generation/harness": self.generation_harness, + "agent/model": self.agent_model, + "agent/harness": self.agent_harness, "total": self.total, } @@ -159,9 +159,11 @@ def stats(self) -> dict[str, Stat]: class TraceMetrics(StatGroup): - """Trace-level metrics for one agent, one value per episode: a fan-out (n same-agent traces - in one episode, e.g. n solvers) collapses to its within-episode mean first, so - ``/num_turns/mean`` is the mean over episodes of the agent's mean turns.""" + """Trace-level metrics for one agent. The reward and count distributions carry one value per + episode — a fan-out (n same-agent traces in one episode, e.g. n solvers) collapses to its + within-episode mean first, so ``/num_turns/mean`` is the mean over episodes of the + agent's mean turns. Everything else (timing, custom metrics / reward components, rates, stop + conditions, errors, solve rates) is flat over the agent's traces.""" DISTRIBUTIONS = ("reward", "num_total_tokens", "num_input_tokens", "num_output_tokens", "num_turns", "num_branches") RATES = ("is_truncated", "is_completed") @@ -176,23 +178,89 @@ def stats(self) -> dict[str, Stat]: for name in (*self.DISTRIBUTIONS, *self.RATES) } - def to_dict(self, prefix: str) -> dict[str, float]: - """Full ```` fan-out for the distributions; ``/mean`` only for the 0/1 rates.""" + @property + def timing(self) -> TimingMetrics: + return TimingMetrics(self.rollouts) + + @property + def metrics(self) -> CustomMetrics: + """Env custom ``@metric`` outputs, keyed by name.""" + return CustomMetrics(self.rollouts, "metrics") + + @property + def rewards(self) -> CustomMetrics: + """Per-component reward breakdown, keyed by name (each entry's weighted ``value``, + summed into the scalar ``reward``).""" + return CustomMetrics(self.rollouts, "rewards", value=lambda reward: reward.value) + + @property + def has_error(self) -> Stat: + return Stat([float(r.has_error) for r in self.rollouts]) + + def stop_conditions(self) -> dict[str, float]: + """``generation_truncated`` over the agent's traces, then each recorded + ``stop_condition``'s rate over the traces that recorded one.""" + out = { + "generation_truncated": sum( + 1 for r in self.rollouts if r.is_truncated and r.stop_condition != "prompt_too_long" + ) + / len(self.rollouts) + } + conditions = [r.stop_condition for r in self.rollouts if r.stop_condition is not None] + for condition in sorted(set(conditions)): + out[condition] = conditions.count(condition) / len(conditions) + return out + + def error_types(self) -> dict[str, int]: + """Count of errored traces by error type (the trace's last error — e.g. ``Cancelled``, + ``ProviderError``).""" + types = [r.last_error.type for r in self.rollouts if r.has_error and r.last_error is not None] + return {t: types.count(t) for t in sorted(set(types))} + + def solve_rates(self) -> dict[str, float]: + """Per-group solve rates over the agent's traces, assuming binary 0/1 rewards (unspecified + for other reward ranges): ``solved_none`` (the group earned no reward), ``solved_all`` + (every trace scored 1.0), and ``solved_some`` (the mixed remainder — the GRPO-signal + groups).""" + groups: dict = {} + for r in self.rollouts: + groups.setdefault(r.group_id, []).append(r) + n_groups = len(groups) + solved_none = sum(1 for g in groups.values() if sum(r.reward for r in g) == 0) + solved_all = sum(1 for g in groups.values() if all(r.reward == 1.0 for r in g)) + return { + "solved_none": solved_none / n_groups, + "solved_all": solved_all / n_groups, + "solved_some": 1 - (solved_none + solved_all) / n_groups, + } + + def to_dict(self, prefix: str, *, subset: Subset) -> dict[str, float]: + """Full ```` fan-out for the distributions; ``/mean`` only for the 0/1 rates. + Errors live only on the ``all`` subset (``effective`` drops them by construction).""" stats = self.stats() out: dict[str, float] = {} for name in self.DISTRIBUTIONS: out |= stats[name].to_dict(f"{prefix}/{name}") for name in self.RATES: out[f"{prefix}/{name}/mean"] = stats[name].mean() + out |= self.timing.to_dict(f"{prefix}/timing") + out |= self.metrics.to_dict(f"{prefix}/metrics") + out |= self.rewards.to_dict(f"{prefix}/rewards") + if subset == "all": + out[f"{prefix}/has_error/mean"] = self.has_error.mean() + out |= {f"{prefix}/error/{t}": float(count) for t, count in self.error_types().items()} + out |= {f"{prefix}/stop_condition/{k}": v for k, v in self.stop_conditions().items()} + out |= {f"{prefix}/{k}": v for k, v in self.solve_rates().items()} return out class EpisodeMetrics: - """Metrics shared by train and eval over a rollout list. Distributional metrics are ``Stat``s - (mean/max/min); boolean metrics are ``Stat``s of 0/1 (use ``.mean()`` for the rate). The count - metrics (tokens/turns/branches) are episode-level — one value per episode, summing its traces — - with the per-agent trace-level view under ``by_agent()``. ``to_wandb`` assembles the full - ``{prefix}/{subset}/...`` dict; ``TrainMetrics`` / ``EvalMetrics`` extend it.""" + """Metrics shared by train and eval over a rollout list. The count metrics (tokens/turns/ + branches) are episode-level — one value per episode, summing its traces — and are the only + per-metric keys ``to_wandb`` emits at the env level; every trace-level metric is emitted per + agent via ``by_agent()``. The boolean ``Stat`` properties (0/1 distributions, ``.mean()`` is + the rate) serve the console log lines. ``TrainMetrics`` / ``EvalMetrics`` extend ``to_wandb`` + with the train pipeline rates and the eval scores.""" def __init__(self, rollouts: list[Rollout]) -> None: self.rollouts = rollouts @@ -239,70 +307,16 @@ def num_turns(self) -> Stat: def num_branches(self) -> Stat: return Stat([float(sum(r.num_branches for r in episode)) for episode in self.episodes()]) - @property - def timing(self) -> TimingMetrics: - return TimingMetrics(self.rollouts) - - @property - def metrics(self) -> CustomMetrics: - """Env custom ``@metric`` outputs, keyed by name (``metrics.metrics["acc"].mean()``).""" - return CustomMetrics(self.rollouts, "metrics") - - @property - def rewards(self) -> CustomMetrics: - """Per-component reward breakdown, keyed by name (each entry's weighted ``value``, - summed into the scalar ``reward``).""" - return CustomMetrics(self.rollouts, "rewards", value=lambda reward: reward.value) - - # Boolean rate metrics (0/1 distributions — ``.mean()`` is the rate) + # Boolean rate metrics for the console log lines (0/1 distributions — ``.mean()`` is the + # rate); to_wandb emits their per-agent counterparts instead. @property def is_truncated(self) -> Stat: return Stat([float(r.is_truncated) for r in self.rollouts]) - @property - def is_completed(self) -> Stat: - return Stat([float(r.is_completed) for r in self.rollouts]) - @property def has_error(self) -> Stat: return Stat([float(r.has_error) for r in self.rollouts]) - def stop_conditions(self) -> dict[str, float]: - """``generation_truncated`` over all rollouts, then each recorded ``stop_condition``'s rate - over the rollouts that recorded one.""" - out = { - "generation_truncated": sum( - 1 for r in self.rollouts if r.is_truncated and r.stop_condition != "prompt_too_long" - ) - / len(self.rollouts) - } - conditions = [r.stop_condition for r in self.rollouts if r.stop_condition is not None] - for condition in sorted(set(conditions)): - out[condition] = conditions.count(condition) / len(conditions) - return out - - def error_types(self) -> dict[str, int]: - """Count of errored rollouts by error type (the rollout's last error — e.g. ``Cancelled``, - ``ProviderError``).""" - types = [r.last_error.type for r in self.rollouts if r.has_error and r.last_error is not None] - return {t: types.count(t) for t in sorted(set(types))} - - def solve_rates(self) -> dict[str, float]: - """Per-group solve rates, assuming binary 0/1 rewards (unspecified for other reward ranges): - ``solved_none`` (the group earned no reward), ``solved_all`` (every rollout scored 1.0), and - ``solved_some`` (the mixed remainder — the GRPO-signal groups).""" - groups: dict = {} - for r in self.rollouts: - groups.setdefault(r.group_id, []).append(r) - n_groups = len(groups) - solved_none = sum(1 for g in groups.values() if sum(r.reward for r in g) == 0) - solved_all = sum(1 for g in groups.values() if all(r.reward == 1.0 for r in g)) - return { - "solved_none": solved_none / n_groups, - "solved_all": solved_all / n_groups, - "solved_some": 1 - (solved_none + solved_all) / n_groups, - } - def to_wandb(self, *, prefix: str, subset: Subset) -> dict[str, float]: """The common metric dict for one ``{prefix}/{subset}`` slice. Empty input → ``{}``.""" if not self.rollouts: @@ -314,25 +328,14 @@ def to_wandb(self, *, prefix: str, subset: Subset) -> dict[str, float]: out |= self.num_output_tokens.to_dict(f"{p}/num_output_tokens") out |= self.num_turns.to_dict(f"{p}/num_turns") out |= self.num_branches.to_dict(f"{p}/num_branches") - out |= self.timing.to_dict(f"{p}/timing") - out |= self.metrics.to_dict(f"{p}/metrics") - out |= self.rewards.to_dict(f"{p}/rewards") for agent, agent_metrics in self.by_agent().items(): - out |= agent_metrics.to_dict(f"{p}/{agent}") - out[f"{p}/is_truncated/mean"] = self.is_truncated.mean() - out[f"{p}/is_completed/mean"] = self.is_completed.mean() - # errors live only on the `all` subset (effective drops them), so emit the rate + the - # per-type counts there only - if subset == "all": - out[f"{p}/has_error/mean"] = self.has_error.mean() - out |= {f"{p}/error/{t}": float(count) for t, count in self.error_types().items()} - out |= {f"{p}/stop_condition/{k}": v for k, v in self.stop_conditions().items()} - out |= {f"{p}/{k}": v for k, v in self.solve_rates().items()} + out |= agent_metrics.to_dict(f"{p}/{agent}", subset=subset) return out class TrainMetrics(EpisodeMetrics): - """Common metrics plus the reward distribution and filter-pipeline rates.""" + """Common metrics plus the filter-pipeline rates. ``reward`` (flat over all traces) serves the + console log lines and distributions; the wandb reward stats are per-agent.""" @property def reward(self) -> Stat: @@ -358,7 +361,6 @@ def to_wandb(self, *, prefix: str, subset: Subset) -> dict[str, float]: if not self.rollouts: return out p = f"{prefix}/{subset}" - out |= self.reward.to_dict(f"{p}/reward") out[f"{p}/is_trainable/mean"] = self.is_trainable.mean() out[f"{p}/is_filtered/mean"] = self.is_filtered.mean() out |= {f"{p}/filters/{k}/mean": v for k, v in self.filter_rates().items()} diff --git a/src/prime_rl/utils/monitor/wandb.py b/src/prime_rl/utils/monitor/wandb.py index a385ef7570..1669564f06 100644 --- a/src/prime_rl/utils/monitor/wandb.py +++ b/src/prime_rl/utils/monitor/wandb.py @@ -298,18 +298,22 @@ def save_final_summary(self, filename: str = "final_summary.json") -> None: OVERVIEW_NAME = "overview" -# Per-rollout metrics (as "/" under "/") shown for BOTH train and eval. -# Quality metrics read the effective subset — the all subset includes errored rollouts, whose -# zero values skew the distributions. has_error only exists on all (effective drops errors by -# construction). Only the reward metrics differ — train uses "reward/mean", eval uses "avg@k", -# each shown for the all and effective subsets — and each section builder prepends its own. +# Rollout metrics (under "/") shown for BOTH train and eval. Quality metrics read the +# effective subset — the all subset includes errored rollouts, whose zero values skew the +# distributions. has_error only exists on all (effective drops errors by construction). The count +# metrics are episode-level exact keys; the trace-level metrics (reward, truncation, errors) live +# under the per-agent subtree, whose names are data-dependent — matched by regex, one panel per +# agent. Only the reward metrics differ — train uses the per-agent "reward/mean", eval uses the +# env-level "avg@k" — and each section builder prepends its own. COMMON_METRICS = [ - "all/has_error/mean", - "effective/is_truncated/mean", "effective/num_total_tokens/mean", "effective/num_turns/mean", "effective/num_branches/mean", ] +COMMON_REGEXES = [ + "all/[^/]+/has_error/mean", + "effective/[^/]+/is_truncated/mean", +] STABILITY_METRICS = ["optim/grad_norm", "entropy/all/mean", "mismatch_kl/all/mean", "kl_ent_ratio/mean"] @@ -350,10 +354,14 @@ def section(name: str, metrics: Sequence[str] = (), regexes: Sequence[str] = ()) def train_section(name: str, scope: str) -> ws.Section: + # Env names may carry regex metacharacters (e.g. "+"), so the scope is escaped in the + # regex-matched per-agent panels. + pattern = re.escape(scope) return section( name, - metrics=[f"{scope}/all/reward/mean", f"{scope}/effective/reward/mean"] - + [f"{scope}/{m}" for m in COMMON_METRICS], + metrics=[f"{scope}/{m}" for m in COMMON_METRICS], + regexes=[f"{pattern}/all/[^/]+/reward/mean", f"{pattern}/effective/[^/]+/reward/mean"] + + [f"{pattern}/{r}" for r in COMMON_REGEXES], ) @@ -363,7 +371,8 @@ def eval_section(name: str, env_pattern: str) -> ws.Section: return section( name, regexes=[f"eval/{env_pattern}/all/avg@.*", f"eval/{env_pattern}/effective/avg@.*"] - + [f"eval/{env_pattern}/{m}" for m in COMMON_METRICS], + + [f"eval/{env_pattern}/{m}" for m in COMMON_METRICS] + + [f"eval/{env_pattern}/{r}" for r in COMMON_REGEXES], ) diff --git a/tests/unit/orchestrator/test_metrics.py b/tests/unit/orchestrator/test_metrics.py index 9a47742490..54531ea8dc 100644 --- a/tests/unit/orchestrator/test_metrics.py +++ b/tests/unit/orchestrator/test_metrics.py @@ -35,9 +35,9 @@ def mk( is_filtered: bool = False, filter_results: dict | None = None, setup: float = 0.0, - generation: float = 0.0, - generation_model: float = 0.0, - generation_harness: float = 0.0, + agent_time: float = 0.0, + agent_model: float = 0.0, + agent_harness: float = 0.0, finalize: float = 0.0, scoring: float = 0.0, ): @@ -68,10 +68,10 @@ def mk( filter_results=filter_results or {}, timing=SimpleNamespace( setup=SimpleNamespace(duration=setup), - generation=SimpleNamespace( - duration=generation, - model=SimpleNamespace(duration=generation_model), - harness=SimpleNamespace(duration=generation_harness), + agent=SimpleNamespace( + duration=agent_time, + model=SimpleNamespace(duration=agent_model), + harness=SimpleNamespace(duration=agent_harness), ), finalize=SimpleNamespace(duration=finalize), scoring=SimpleNamespace(duration=scoring), @@ -114,7 +114,8 @@ def test_to_wandb_distributions(): ).metrics assert m.num_input_tokens.mean() == 5.0 # fluent Stat access out = m.to_wandb(prefix="train/agg", subset="all") - assert out["train/agg/all/reward/mean"] == 0.5 + assert out["train/agg/all/agent/reward/mean"] == 0.5 + assert "train/agg/all/reward/mean" not in out # trace-level metrics are agent-only assert out["train/agg/all/num_total_tokens/mean"] == 15.0 assert out["train/agg/all/num_total_tokens/max"] == 20.0 # single-trace episodes: one value per rollout assert out["train/agg/all/num_input_tokens/mean"] == 5.0 @@ -142,16 +143,16 @@ def test_episode_and_agent_levels(): assert out["train/agg/all/solver/reward/mean"] == 0.5 assert out["train/agg/all/proposer/is_truncated/mean"] == 0.0 assert "train/agg/all/proposer/is_truncated/p90" not in out # rates emit /mean only - assert out["train/agg/all/reward/mean"] == 0.5 # trace-level metrics stay flat over rollouts + assert "train/agg/all/reward/mean" not in out # reward never pools across agents def test_boolean_rates_and_error_breakdown_all_only(): rc = TrainRollouts([mk(is_truncated=True), mk(has_error=True, error_type="ProviderError"), mk(is_filtered=True)]) out = rc.metrics.to_wandb(prefix="train/agg", subset="all") - assert out["train/agg/all/is_truncated/mean"] == 1 / 3 - assert out["train/agg/all/is_completed/mean"] == 1.0 - assert out["train/agg/all/has_error/mean"] == 1 / 3 - assert out["train/agg/all/error/ProviderError"] == 1 # error-type breakdown by count + assert out["train/agg/all/agent/is_truncated/mean"] == 1 / 3 + assert out["train/agg/all/agent/is_completed/mean"] == 1.0 + assert out["train/agg/all/agent/has_error/mean"] == 1 / 3 + assert out["train/agg/all/agent/error/ProviderError"] == 1 # error-type breakdown by count assert not any("no_response" in k for k in out) # removed metric # has_error + the error-type counts are structurally empty on effective, so emitted on `all` only eff = rc.effective.metrics.to_wandb(prefix="train/agg", subset="effective") @@ -161,16 +162,20 @@ def test_boolean_rates_and_error_breakdown_all_only(): def test_solve_rates(): groups = {"A": [1.0, 1.0], "B": [0.0, 0.0], "C": [1.0, 0.0], "D": [1.0, 0.0]} # all / none / some / some out = train_wandb([mk(reward=r, group_id=g) for g, rs in groups.items() for r in rs]) - rates = (out["train/agg/all/solved_all"], out["train/agg/all/solved_none"], out["train/agg/all/solved_some"]) + rates = ( + out["train/agg/all/agent/solved_all"], + out["train/agg/all/agent/solved_none"], + out["train/agg/all/agent/solved_some"], + ) assert rates == (0.25, 0.25, 0.5) def test_stop_condition_breakdown(): truncated = [mk(is_truncated=True, stop_condition=c) for c in ("length", "max_turns", "prompt_too_long")] out = train_wandb(truncated + [mk(stop_condition=None)]) - assert out["train/agg/all/stop_condition/generation_truncated"] == 0.5 # truncated & not prompt_too_long, over all - assert out["train/agg/all/stop_condition/length"] == 1 / 3 # over the 3 recorded conditions - assert out["train/agg/all/stop_condition/prompt_too_long"] == 1 / 3 + assert out["train/agg/all/agent/stop_condition/generation_truncated"] == 0.5 # truncated & not prompt_too_long + assert out["train/agg/all/agent/stop_condition/length"] == 1 / 3 # over the 3 recorded conditions + assert out["train/agg/all/agent/stop_condition/prompt_too_long"] == 1 / 3 def test_nested_metrics_and_rewards(): @@ -179,24 +184,26 @@ def test_nested_metrics_and_rewards(): mk(metrics={"acc": 3.0, "fmt": 5.0}, rewards={"correct": vf.Reward(score=0.0), "format": vf.Reward(score=1.0)}), ] m = TrainRollouts(rollouts).metrics - assert m.metrics["acc"].mean() == 2.0 and m.rewards["correct"].mean() == 0.5 # nested group access + agent = m.by_agent()["agent"] + assert agent.metrics["acc"].mean() == 2.0 and agent.rewards["correct"].mean() == 0.5 # nested group access out = m.to_wandb(prefix="train/agg", subset="all") - assert out["train/agg/all/metrics/acc/mean"] == 2.0 # averaged over reporters - assert out["train/agg/all/metrics/fmt/mean"] == 5.0 # single reporter - assert out["train/agg/all/rewards/format/mean"] == 0.5 + assert out["train/agg/all/agent/metrics/acc/mean"] == 2.0 # averaged over reporters + assert out["train/agg/all/agent/metrics/fmt/mean"] == 5.0 # single reporter + assert out["train/agg/all/agent/rewards/format/mean"] == 0.5 def test_nested_timing(): m = TrainRollouts( - [mk(setup=1.0, generation=2.0, generation_model=1.5, generation_harness=0.5, finalize=0.5, scoring=0.5)] + [mk(setup=1.0, agent_time=2.0, agent_model=1.5, agent_harness=0.5, finalize=0.5, scoring=0.5)] ).metrics - assert m.timing.setup.mean() == 1.0 and m.timing.total.mean() == 4.0 # total sums all four phases - assert m.timing.generation_model.mean() == 1.5 and m.timing.generation_harness.mean() == 0.5 + timing = m.by_agent()["agent"].timing + assert timing.setup.mean() == 1.0 and timing.total.mean() == 4.0 # total sums all four phases + assert timing.agent_model.mean() == 1.5 and timing.agent_harness.mean() == 0.5 out = m.to_wandb(prefix="train/agg", subset="all") - assert out["train/agg/all/timing/setup/mean"] == 1.0 - assert out["train/agg/all/timing/total/mean"] == 4.0 - assert out["train/agg/all/timing/generation/model/mean"] == 1.5 - assert out["train/agg/all/timing/generation/harness/mean"] == 0.5 + assert out["train/agg/all/agent/timing/setup/mean"] == 1.0 + assert out["train/agg/all/agent/timing/total/mean"] == 4.0 + assert out["train/agg/all/agent/timing/agent/model/mean"] == 1.5 + assert out["train/agg/all/agent/timing/agent/harness/mean"] == 0.5 def test_train_only_metrics_absent_from_eval(): From a073f60e4eb70eca3114b5aea2db2118199549f2 Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Thu, 30 Jul 2026 22:08:08 +0000 Subject: [PATCH 05/58] feat(configs): agentic-judge reverse-text debug config Co-Authored-By: Claude Fable 5 --- configs/debug/envs/README.md | 7 ++++ configs/debug/envs/agentic_judge.toml | 46 +++++++++++++++++++++++++++ 2 files changed, 53 insertions(+) create mode 100644 configs/debug/envs/README.md create mode 100644 configs/debug/envs/agentic_judge.toml diff --git a/configs/debug/envs/README.md b/configs/debug/envs/README.md new file mode 100644 index 0000000000..920c5481b4 --- /dev/null +++ b/configs/debug/envs/README.md @@ -0,0 +1,7 @@ +# Envs — Debug Configs + +Minimal end-to-end configs for bundled multi-agent envs, using `PrimeIntellect/Qwen3-0.6B-Reverse-Text-SFT` as the policy. + +| Config | Env | Notes | +|---|---|---| +| `agentic_judge.toml` | `agentic-judge` over `reverse-text-v1` | solver in a docker box, frozen `deepseek/deepseek-v4-flash` judge grades in the same box (needs docker + a Prime Inference key) | diff --git a/configs/debug/envs/agentic_judge.toml b/configs/debug/envs/agentic_judge.toml new file mode 100644 index 0000000000..d05d69d0a0 --- /dev/null +++ b/configs/debug/envs/agentic_judge.toml @@ -0,0 +1,46 @@ +max_steps = 20 +seq_len = 2048 + +[model] +name = "PrimeIntellect/Qwen3-0.6B-Reverse-Text-SFT" + +[wandb] +project = "reverse-text" +name = "reverse-text-agentic-judge" + +[orchestrator] +batch_size = 128 +group_size = 16 + +[orchestrator.train.sampling] +max_completion_tokens = 128 + +# The agentic judge over reverse-text: the solver plays the task in a docker box, +# a frozen judge then grades the work in the same box and its verdict composes +# with the taskset's own reward on the solver's trace. +[[orchestrator.train.source]] +name = "reverse-text-judge" +env.id = "agentic-judge" +env.taskset = { id = "reverse-text-v1" } +env.solver.harness = { id = "null" } +env.solver.runtime = { type = "docker" } +env.judge.harness = { id = "bash" } +env.judge.model = "deepseek/deepseek-v4-flash" +env.judge.client = { type = "eval" } +# The run's sampling caps completions at the solver's 128 tokens — the judge +# reasons before writing its verdict, so it gets its own budget. +env.judge.sampling = { max_tokens = 4096 } +# The judge explores the trace record step by step before writing its verdict — +# 10 turns is routinely not enough. +env.judge.max_turns = 25 +env.score = { task_weight = 1.0 } + +[orchestrator.renderer] +name = "prime-qwen3" + +[trainer.optim] +lr = 3e-6 + +[ckpt] + +[inference] From 7fd65a9d72338d42a94032379ab667cea878af99 Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Thu, 30 Jul 2026 22:33:45 +0000 Subject: [PATCH 06/58] chore(configs): deterministic two-call judge policy for the agentic-judge debug config Co-Authored-By: Claude Fable 5 --- configs/debug/envs/README.md | 2 +- configs/debug/envs/agentic_judge.toml | 11 ++++++++--- configs/debug/envs/judge_policy.md | 21 +++++++++++++++++++++ 3 files changed, 30 insertions(+), 4 deletions(-) create mode 100644 configs/debug/envs/judge_policy.md diff --git a/configs/debug/envs/README.md b/configs/debug/envs/README.md index 920c5481b4..dadd90f5b7 100644 --- a/configs/debug/envs/README.md +++ b/configs/debug/envs/README.md @@ -4,4 +4,4 @@ Minimal end-to-end configs for bundled multi-agent envs, using `PrimeIntellect/Q | Config | Env | Notes | |---|---|---| -| `agentic_judge.toml` | `agentic-judge` over `reverse-text-v1` | solver in a docker box, frozen `deepseek/deepseek-v4-flash` judge grades in the same box (needs docker + a Prime Inference key) | +| `agentic_judge.toml` | `agentic-judge` over `reverse-text-v1` | solver in a docker box, frozen `deepseek/deepseek-v4-flash` judge grades in the same box via `judge_policy.md` — a two-tool-call mirror of the deterministic check (needs docker + a Prime Inference key; run from the repo root) | diff --git a/configs/debug/envs/agentic_judge.toml b/configs/debug/envs/agentic_judge.toml index d05d69d0a0..c22435e203 100644 --- a/configs/debug/envs/agentic_judge.toml +++ b/configs/debug/envs/agentic_judge.toml @@ -11,6 +11,9 @@ name = "reverse-text-agentic-judge" [orchestrator] batch_size = 128 group_size = 16 +# Each in-flight episode holds a docker box (solver + judge share it); unbounded +# in-flight stacks hundreds of writable layers and can fill the disk. +max_inflight_episodes = 64 [orchestrator.train.sampling] max_completion_tokens = 128 @@ -30,9 +33,11 @@ env.judge.client = { type = "eval" } # The run's sampling caps completions at the solver's 128 tokens — the judge # reasons before writing its verdict, so it gets its own budget. env.judge.sampling = { max_tokens = 4096 } -# The judge explores the trace record step by step before writing its verdict — -# 10 turns is routinely not enough. -env.judge.max_turns = 25 +env.judge.max_turns = 8 +# Debug grading policy: mirror reverse-text's deterministic check in two tool +# calls, instead of the default open-ended trace investigation. The path is +# relative to the repo root — run from there. +env.task.prompt = "configs/debug/envs/judge_policy.md" env.score = { task_weight = 1.0 } [orchestrator.renderer] diff --git a/configs/debug/envs/judge_policy.md b/configs/debug/envs/judge_policy.md new file mode 100644 index 0000000000..c9ecc370db --- /dev/null +++ b/configs/debug/envs/judge_policy.md @@ -0,0 +1,21 @@ +You are verifying a text-reversal attempt. This is a DEBUG judge that mirrors a +deterministic checker — do not investigate the trace beyond the recipe below, +and use AT MOST TWO tool calls total. + +## The task the agent was given + +{prompt} + +## Recipe + +First tool call — one python3 command that: + +1. loads `/tmp/trace.json`, +2. takes the task prompt (`data["task"]["data"]["prompt"]`) and the last + assistant message (`[n["message"] for n in data["nodes"] if n["message"]["role"] == "assistant"][-1]["content"]`), +3. prints `PASS` if the answer equals the prompt reversed character-by-character + (compare with surrounding whitespace stripped), else prints `FAIL` and both + strings. + +Second tool call — write your verdict file: `solved` is `yes` on `PASS`, `no` +otherwise. Then state your verdict and stop. From e0baafd5fdeb24bc70fcb1929aa398891e16aeaf Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Thu, 30 Jul 2026 22:55:28 +0000 Subject: [PATCH 07/58] chore: bump verifiers to vf#2187 head Co-Authored-By: Claude Fable 5 --- deps/verifiers | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deps/verifiers b/deps/verifiers index 515b820490..3b61e6406c 160000 --- a/deps/verifiers +++ b/deps/verifiers @@ -1 +1 @@ -Subproject commit 515b8204903afb2fc41f6bd6aaaafa2af9c3e3d1 +Subproject commit 3b61e6406cf9e156f3f6420e6cf4fde37a9454ad From 9220746482d3432568370bcd0c04555eb3914629 Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Fri, 31 Jul 2026 00:23:16 +0000 Subject: [PATCH 08/58] chore: bump verifiers to vf#2187 head (main merged) Co-Authored-By: Claude Fable 5 --- deps/verifiers | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deps/verifiers b/deps/verifiers index 3b61e6406c..31a447a4a9 160000 --- a/deps/verifiers +++ b/deps/verifiers @@ -1 +1 @@ -Subproject commit 3b61e6406cf9e156f3f6420e6cf4fde37a9454ad +Subproject commit 31a447a4a96d9782e8e46bbdd0c1ea1cc11bf723 From b76e153030068b422a70ee1646d178a9ba7729cd Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Fri, 31 Jul 2026 00:32:28 +0000 Subject: [PATCH 09/58] chore: bump verifiers pin Co-Authored-By: Claude Fable 5 --- deps/verifiers | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deps/verifiers b/deps/verifiers index 31a447a4a9..8a941dbcb3 160000 --- a/deps/verifiers +++ b/deps/verifiers @@ -1 +1 @@ -Subproject commit 31a447a4a96d9782e8e46bbdd0c1ea1cc11bf723 +Subproject commit 8a941dbcb38987c09958d752c2a10e58190e267c From fb49e3490a33e4888454b537a71f581008d06c29 Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Fri, 31 Jul 2026 00:39:19 +0000 Subject: [PATCH 10/58] chore: bump verifiers pin Co-Authored-By: Claude Fable 5 --- deps/verifiers | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deps/verifiers b/deps/verifiers index 8a941dbcb3..c9ee978483 160000 --- a/deps/verifiers +++ b/deps/verifiers @@ -1 +1 @@ -Subproject commit 8a941dbcb38987c09958d752c2a10e58190e267c +Subproject commit c9ee9784832395ce27631b3f66578085b925337b From 5f1404f2f396904307db4342793bf59823af3493 Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Fri, 31 Jul 2026 17:13:34 +0000 Subject: [PATCH 11/58] chore: re-pin verifiers to the vf#2187 merge commit Co-Authored-By: Claude Fable 5 --- deps/verifiers | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deps/verifiers b/deps/verifiers index c9ee978483..2719276245 160000 --- a/deps/verifiers +++ b/deps/verifiers @@ -1 +1 @@ -Subproject commit c9ee9784832395ce27631b3f66578085b925337b +Subproject commit 27192762454de76add51434d0383e918315da8be From 60eb2e5c0225324869dcf8e072018ade39498e16 Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Sat, 1 Aug 2026 03:06:54 +0000 Subject: [PATCH 12/58] chore: bump verifiers and research-environments to latest main verifiers f646beb37 -> f6e420b99 (0.2.2.dev66, 17 commits) and research-environments 6a2dee6fc -> ccb375ce5 (main + the scicode-v1 pin relax from research-environments#740). Adaptations to the verifiers changes: - Episode.error is now Episode.errors + a last_error property (vf#2187) - verifiers.v1.loaders / verifiers.v1.push moved into verifiers.v1.utils.{loaders,platform} (vf#2204) - verifiers floors bumped to 0.2.2.dev66 (root and prime-rl-configs) Co-Authored-By: Claude Fable 5 --- deps/research-environments | 2 +- deps/verifiers | 2 +- packages/prime-rl-configs/pyproject.toml | 2 +- pyproject.toml | 2 +- src/prime_rl/orchestrator/envs.py | 4 +- src/prime_rl/trainer/runs.py | 2 +- src/prime_rl/utils/monitor/prime.py | 4 +- uv.lock | 170 +++++++++++++++++------ 8 files changed, 134 insertions(+), 54 deletions(-) diff --git a/deps/research-environments b/deps/research-environments index 6a2dee6fcd..ccb375ce55 160000 --- a/deps/research-environments +++ b/deps/research-environments @@ -1 +1 @@ -Subproject commit 6a2dee6fcd805f41128dba0024f7073afd496169 +Subproject commit ccb375ce550e48af9e898351145975521a2b6897 diff --git a/deps/verifiers b/deps/verifiers index f646beb37e..f6e420b990 160000 --- a/deps/verifiers +++ b/deps/verifiers @@ -1 +1 @@ -Subproject commit f646beb37eef51869f886f244456e3d07818e4d6 +Subproject commit f6e420b9908ae14d625f079881f13c15011ee1c9 diff --git a/packages/prime-rl-configs/pyproject.toml b/packages/prime-rl-configs/pyproject.toml index 5785a3aa8d..aabadc1268 100644 --- a/packages/prime-rl-configs/pyproject.toml +++ b/packages/prime-rl-configs/pyproject.toml @@ -10,7 +10,7 @@ dependencies = [ "renderers>=0.1.9.dev10", "tomli>=2.2.1", "tomli-w>=1.2.0", - "verifiers>=0.2.2.dev43", + "verifiers>=0.2.2.dev66", ] [build-system] diff --git a/pyproject.toml b/pyproject.toml index 91214d92b3..9cbeeec4a9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -33,7 +33,7 @@ dependencies = [ "setproctitle>=1.3.0", "uvloop>=0.21.0", "torchtitan", - "verifiers[harbor]>=0.2.2.dev43", + "verifiers[harbor]>=0.2.2.dev66", "renderers", "dion", "tilelang>=0.1.8", diff --git a/src/prime_rl/orchestrator/envs.py b/src/prime_rl/orchestrator/envs.py index 3bffb3504d..311249c3a6 100644 --- a/src/prime_rl/orchestrator/envs.py +++ b/src/prime_rl/orchestrator/envs.py @@ -217,14 +217,14 @@ async def run( sampling=self._sampling(cache_salt), ) if not episode.traces: - error = episode.error + error = episode.last_error detail = f"{error.type}: {error.message}" if error is not None else "no traces and no error recorded" raise RuntimeError(f"env-rollout failed before any trace was minted — {detail}") rollouts = [ROLLOUT_TYPE.model_construct(**dict(wire)) for wire in episode.traces] for rollout in rollouts: rollout.episode_id = episode.id if not episode.ok and rollout.ok: - error = episode.error or vf.Error( + error = episode.last_error or vf.Error( type="EpisodeFailed", message="A sibling trace in this episode failed" ) rollout.errors = [*rollout.errors, error] diff --git a/src/prime_rl/trainer/runs.py b/src/prime_rl/trainer/runs.py index 24dd48dc2b..6dc384c9d0 100644 --- a/src/prime_rl/trainer/runs.py +++ b/src/prime_rl/trainer/runs.py @@ -236,7 +236,7 @@ def get_orchestrator_config(self, run_id: str) -> Optional["OrchestratorConfig"] with open(config_path, "rb") as f: config_dict = tomli.load(f) - from verifiers.v1.loaders import skip_plugin_install + from verifiers.v1.utils.loaders import skip_plugin_install from prime_rl.configs.orchestrator import OrchestratorConfig diff --git a/src/prime_rl/utils/monitor/prime.py b/src/prime_rl/utils/monitor/prime.py index e0d0afec5b..eee454a9aa 100644 --- a/src/prime_rl/utils/monitor/prime.py +++ b/src/prime_rl/utils/monitor/prime.py @@ -15,7 +15,7 @@ import pyarrow.parquet as pq from prime_cli.core.config import Config as PrimeConfig from transformers.tokenization_utils import PreTrainedTokenizer -from verifiers.v1.push import trace_to_sample +from verifiers.v1.utils.platform import trace_to_sample from prime_rl.configs.orchestrator import OrchestratorConfig from prime_rl.configs.shared import PrimeMonitorConfig @@ -295,7 +295,7 @@ def _rollouts_to_parquet_bytes(self, rollouts: list[Rollout], step: int) -> byte """Convert rollouts to Parquet bytes for upload. One row per rollout. The conversation is the unit (no prompt/completion split — meaningless mid-branch): `completion` is the last branch's messages and `trajectory` is one message list per branch. Shares - `verifiers.v1.push.trace_to_sample` with verifiers' eval `--push`, so a training-run + `verifiers.v1.utils.platform.trace_to_sample` with verifiers' eval `--push`, so a training-run sample and an eval sample land on the platform identically; the RFT-only columns (run/step/advantage/problem_id/env_name) are layered on here.""" now = datetime.now(timezone.utc) diff --git a/uv.lock b/uv.lock index 263191aea2..7a1957a9b1 100644 --- a/uv.lock +++ b/uv.lock @@ -11,7 +11,7 @@ supported-markers = [ ] [options] -exclude-newer = "2026-07-23T19:34:50.942572Z" +exclude-newer = "0001-01-01T00:00:00Z" # This has no effect and is included for backwards compatibility when using relative exclude-newer values. exclude-newer-span = "P7D" [options.exclude-newer-package] @@ -51,12 +51,14 @@ members = [ "color-codeword-v1", "compact", "deepdive-v1", + "deepseek-prover-v1", "deepwiki-v1", "enterprise-ops-gym-v1", "forth-lang-v1", "frontierscience-v1", "general-agent-v1", "glossary-v1", + "goedel-pset-v1", "gpqa-v1", "graphwalks-v1", "gsm8k-v1", @@ -68,8 +70,8 @@ members = [ "i3-science-v1", "ifbench-v1", "ifeval-v1", + "kimina-v1", "kuhn-poker-v1", - "lean-v1", "livecodebench-v1", "longbenchpro-v1", "longcot-mini-v1", @@ -78,12 +80,14 @@ members = [ "math-env-v1", "math500-v1", "mcp-atlas-v1", + "minif2f-v1", "mmlu-pro-v1", "mmmu-pro-v1", "mmmu-v1", "mrcr-v2-v1", "multiswe-v1", "nl2repobench-v1", + "numina-v1", "oolong-pairs-v1", "oolong-real-v1", "oolong-synth-v1", @@ -98,6 +102,7 @@ members = [ "programbench-v1", "prolog-v1", "proposer-solver-v1", + "proverbench-v1", "r2e-gym-v1", "redsearcher-v1", "reverse-text-v1", @@ -484,7 +489,7 @@ requires-dist = [ { name = "datasets" }, { name = "mcp", specifier = "<2" }, { name = "openai" }, - { name = "verifiers", specifier = ">=0.2.1" }, + { name = "verifiers", specifier = ">=0.2.2.dev65" }, ] [[package]] @@ -608,7 +613,7 @@ dependencies = [ requires-dist = [ { name = "bfcl-eval", git = "https://github.com/mikasenghaas/gorilla.git?subdirectory=berkeley-function-call-leaderboard&rev=898763a" }, { name = "soundfile", specifier = ">=0.13.0" }, - { name = "verifiers", specifier = ">=0.2.2.dev21" }, + { name = "verifiers", specifier = ">=0.2.2.dev65" }, ] [[package]] @@ -730,7 +735,7 @@ requires-dist = [ { name = "openai" }, { name = "pystemmer" }, { name = "tokenizers" }, - { name = "verifiers", specifier = ">=0.2.0" }, + { name = "verifiers", specifier = ">=0.2.2.dev65" }, ] [[package]] @@ -1413,6 +1418,21 @@ dependencies = [ requires-dist = [ { name = "datasets" }, { name = "openai" }, + { name = "verifiers", specifier = ">=0.2.2.dev65" }, +] + +[[package]] +name = "deepseek-prover-v1" +version = "0.1.0" +source = { editable = "deps/research-environments/environments/lean/deepseek_prover_v1" } +dependencies = [ + { name = "datasets", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "verifiers", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] + +[package.metadata] +requires-dist = [ + { name = "datasets", specifier = ">=2.0.0" }, { name = "verifiers", specifier = ">=0.2.0" }, ] @@ -1604,7 +1624,7 @@ dependencies = [ [package.metadata] requires-dist = [ { name = "datasets", specifier = ">=4.0.0,<5" }, - { name = "verifiers", specifier = ">=0.2.0" }, + { name = "verifiers", specifier = ">=0.2.2.dev65" }, ] [[package]] @@ -1983,7 +2003,7 @@ dependencies = [ [package.metadata] requires-dist = [ { name = "datasets" }, - { name = "verifiers", specifier = ">=0.2.0" }, + { name = "verifiers", specifier = ">=0.2.2.dev65" }, ] [[package]] @@ -2056,7 +2076,7 @@ dependencies = [ [package.metadata] requires-dist = [ { name = "pydantic" }, - { name = "verifiers", extras = ["harbor"], specifier = ">=0.2.0" }, + { name = "verifiers", extras = ["harbor"], specifier = ">=0.2.2.dev65" }, ] [[package]] @@ -2130,6 +2150,21 @@ dependencies = [ [package.metadata] requires-dist = [{ name = "verifiers" }] +[[package]] +name = "goedel-pset-v1" +version = "0.1.0" +source = { editable = "deps/research-environments/environments/lean/goedel_pset_v1" } +dependencies = [ + { name = "datasets", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "verifiers", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] + +[package.metadata] +requires-dist = [ + { name = "datasets", specifier = ">=2.0.0" }, + { name = "verifiers", specifier = ">=0.2.0" }, +] + [[package]] name = "google-auth" version = "2.55.0" @@ -3177,6 +3212,21 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/81/db/e655086b7f3a705df045bf0933bdd9c2f79bb3c97bfef1384598bb79a217/keyring-25.7.0-py3-none-any.whl", hash = "sha256:be4a0b195f149690c166e850609a477c532ddbfbaed96a404d4e43f8d5e2689f", size = 39160, upload-time = "2025-11-16T16:26:08.402Z" }, ] +[[package]] +name = "kimina-v1" +version = "0.1.0" +source = { editable = "deps/research-environments/environments/lean/kimina_v1" } +dependencies = [ + { name = "datasets", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "verifiers", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] + +[package.metadata] +requires-dist = [ + { name = "datasets", specifier = ">=2.0.0" }, + { name = "verifiers", specifier = ">=0.2.0" }, +] + [[package]] name = "kiwisolver" version = "1.5.0" @@ -3272,21 +3322,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/e9/61/f75cd1fa54d8434276126034aed54dd120747de9a8fa013cdd79545ccbeb/latex2sympy2_extended-1.11.0-py3-none-any.whl", hash = "sha256:aebb77d52ce269e25028e4bea89ddb14d242ba36bcf7b636496fb5fd9728d234", size = 209050, upload-time = "2026-01-10T01:43:19.458Z" }, ] -[[package]] -name = "lean-v1" -version = "0.1.0" -source = { editable = "deps/research-environments/environments/math/lean_v1" } -dependencies = [ - { name = "datasets", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "verifiers", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, -] - -[package.metadata] -requires-dist = [ - { name = "datasets", specifier = ">=2.0.0" }, - { name = "verifiers", specifier = ">=0.2.0" }, -] - [[package]] name = "linkify-it-py" version = "2.1.0" @@ -3646,7 +3681,7 @@ requires-dist = [ { name = "datasets", specifier = ">=4.0.0,<5" }, { name = "httpx" }, { name = "prime-sandboxes", specifier = ">=0.2.19" }, - { name = "verifiers", specifier = ">=0.2.0" }, + { name = "verifiers", specifier = ">=0.2.2.dev65" }, ] [[package]] @@ -3670,6 +3705,21 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/b3/38/89ba8ad64ae25be8de66a6d463314cf1eb366222074cfda9ee839c56a4b4/mdurl-0.1.2-py3-none-any.whl", hash = "sha256:84008a41e51615a49fc9966191ff91509e3c40b939176e643fd50a5c2196b8f8", size = 9979, upload-time = "2022-08-14T12:40:09.779Z" }, ] +[[package]] +name = "minif2f-v1" +version = "0.1.0" +source = { editable = "deps/research-environments/environments/lean/minif2f_v1" } +dependencies = [ + { name = "datasets", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "verifiers", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] + +[package.metadata] +requires-dist = [ + { name = "datasets", specifier = ">=2.0.0" }, + { name = "verifiers", specifier = ">=0.2.0" }, +] + [[package]] name = "mistral-common" version = "1.11.7" @@ -4155,6 +4205,21 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/0c/61/ce753a1d7646dd477e16d15e89473703faebb8995d2f71d7ad69a540b565/numba-0.65.0-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:da8e371e328c06d0010c3d8b44b21858652831b85bcfba78cb22c042e22dbd8e", size = 3501622, upload-time = "2026-04-01T03:51:36.348Z" }, ] +[[package]] +name = "numina-v1" +version = "0.1.0" +source = { editable = "deps/research-environments/environments/lean/numina_v1" } +dependencies = [ + { name = "datasets", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "verifiers", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] + +[package.metadata] +requires-dist = [ + { name = "datasets", specifier = ">=2.0.0" }, + { name = "verifiers", specifier = ">=0.2.0" }, +] + [[package]] name = "numpy" version = "2.3.5" @@ -4593,7 +4658,7 @@ dependencies = [ [package.metadata] requires-dist = [ { name = "datasets" }, - { name = "verifiers", specifier = ">=0.2.0" }, + { name = "verifiers", specifier = ">=0.2.2.dev65" }, ] [[package]] @@ -4743,11 +4808,11 @@ name = "openthoughts-tblite-v1" version = "0.1.1" source = { editable = "deps/research-environments/environments/terminal/openthoughts_tblite_v1" } dependencies = [ - { name = "verifiers", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "verifiers", extra = ["harbor"], marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] [package.metadata] -requires-dist = [{ name = "verifiers", specifier = ">=0.2.1" }] +requires-dist = [{ name = "verifiers", extras = ["harbor"], specifier = ">=0.2.1" }] [[package]] name = "orderly-set" @@ -4828,7 +4893,7 @@ dependencies = [ requires-dist = [ { name = "datasets" }, { name = "openai" }, - { name = "verifiers", specifier = ">=0.2.0" }, + { name = "verifiers", specifier = ">=0.2.2.dev65" }, ] [[package]] @@ -4924,7 +4989,7 @@ requires-dist = [ { name = "openai" }, { name = "prime-sandboxes", specifier = ">=0.2.19" }, { name = "pyyaml", specifier = ">=6.0.1" }, - { name = "verifiers", specifier = ">=0.2.0" }, + { name = "verifiers", specifier = ">=0.2.2.dev65" }, ] [[package]] @@ -5286,7 +5351,7 @@ requires-dist = [ { name = "renderers", specifier = ">=0.1.9.dev10" }, { name = "tomli", specifier = ">=2.2.1" }, { name = "tomli-w", specifier = ">=1.2.0" }, - { name = "verifiers", specifier = ">=0.2.2.dev43" }, + { name = "verifiers", specifier = ">=0.2.2.dev66" }, ] [[package]] @@ -5440,6 +5505,21 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/c4/72/02445137af02769918a93807b2b7890047c32bfb9f90371cbc12688819eb/protobuf-6.33.6-py3-none-any.whl", hash = "sha256:77179e006c476e69bf8e8ce866640091ec42e1beb80b213c3900006ecfba6901", size = 170656, upload-time = "2026-03-18T19:04:59.826Z" }, ] +[[package]] +name = "proverbench-v1" +version = "0.1.0" +source = { editable = "deps/research-environments/environments/lean/proverbench_v1" } +dependencies = [ + { name = "datasets", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "verifiers", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] + +[package.metadata] +requires-dist = [ + { name = "datasets", specifier = ">=2.0.0" }, + { name = "verifiers", specifier = ">=0.2.0" }, +] + [[package]] name = "psutil" version = "7.2.2" @@ -6058,7 +6138,7 @@ dependencies = [ [package.metadata] requires-dist = [ { name = "datasets" }, - { name = "verifiers", specifier = ">=0.2.0" }, + { name = "verifiers", specifier = ">=0.2.2.dev65" }, ] [[package]] @@ -6259,7 +6339,7 @@ dependencies = [ [package.metadata] requires-dist = [ { name = "huggingface-hub" }, - { name = "verifiers", specifier = ">=0.2.0" }, + { name = "verifiers", specifier = ">=0.2.2.dev65" }, ] [[package]] @@ -6328,7 +6408,7 @@ wheels = [ [[package]] name = "scicode-v1" -version = "0.2.0" +version = "0.3.0" source = { editable = "deps/research-environments/environments/code/scicode_v1" } dependencies = [ { name = "datasets", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, @@ -6338,7 +6418,7 @@ dependencies = [ [package.metadata] requires-dist = [ { name = "datasets" }, - { name = "verifiers", specifier = ">=0.2.0" }, + { name = "verifiers", specifier = ">=0.2.2.dev65" }, ] [[package]] @@ -6429,7 +6509,7 @@ dependencies = [ ] [package.metadata] -requires-dist = [{ name = "verifiers", extras = ["harbor"], specifier = ">=0.2.2.dev27" }] +requires-dist = [{ name = "verifiers", extras = ["harbor"], specifier = ">=0.2.2.dev65" }] [[package]] name = "sentence-transformers" @@ -6847,29 +6927,29 @@ dependencies = [ ] [package.metadata] -requires-dist = [{ name = "verifiers", extras = ["harbor"], specifier = ">=0.2.2.dev27" }] +requires-dist = [{ name = "verifiers", extras = ["harbor"], specifier = ">=0.2.2.dev65" }] [[package]] name = "swebench-pro-v1" version = "0.1.0" source = { editable = "deps/research-environments/environments/swe/swebench_pro_v1" } dependencies = [ - { name = "verifiers", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "verifiers", extra = ["harbor"], marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] [package.metadata] -requires-dist = [{ name = "verifiers", specifier = ">=0.2.2.dev27" }] +requires-dist = [{ name = "verifiers", extras = ["harbor"], specifier = ">=0.2.2.dev27" }] [[package]] name = "swebench-verified-v1" version = "0.1.0" source = { editable = "deps/research-environments/environments/swe/swebench_verified_v1" } dependencies = [ - { name = "verifiers", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "verifiers", extra = ["harbor"], marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] [package.metadata] -requires-dist = [{ name = "verifiers", specifier = ">=0.2.2.dev27" }] +requires-dist = [{ name = "verifiers", extras = ["harbor"], specifier = ">=0.2.2.dev65" }] [[package]] name = "swelego-v1" @@ -7059,11 +7139,11 @@ name = "terminal-bench-2-v1" version = "0.1.0" source = { editable = "deps/research-environments/environments/terminal/terminal_bench_2_v1" } dependencies = [ - { name = "verifiers", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "verifiers", extra = ["harbor"], marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] [package.metadata] -requires-dist = [{ name = "verifiers", specifier = ">=0.2.1" }] +requires-dist = [{ name = "verifiers", extras = ["harbor"], specifier = ">=0.2.1" }] [[package]] name = "terminal-lego-v1" @@ -7071,13 +7151,13 @@ version = "0.1.1" source = { editable = "deps/research-environments/environments/terminal/terminal_lego_v1" } dependencies = [ { name = "huggingface-hub", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "verifiers", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "verifiers", extra = ["harbor"], marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] [package.metadata] requires-dist = [ { name = "huggingface-hub" }, - { name = "verifiers", specifier = ">=0.2.0" }, + { name = "verifiers", extras = ["harbor"], specifier = ">=0.2.0" }, ] [[package]] @@ -8548,7 +8628,7 @@ dependencies = [ ] [package.metadata] -requires-dist = [{ name = "verifiers", specifier = ">=0.2.0" }] +requires-dist = [{ name = "verifiers", specifier = ">=0.2.2.dev65" }] [[package]] name = "wordle-v1" From 72b31bc17ead862acd2bcf60434f43c8d37c5c09 Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Mon, 3 Aug 2026 19:37:14 +0000 Subject: [PATCH 13/58] chore: bump research-environments to latest main ccb375ce5 -> 9a57c713c: research-environments#740 refreshed on current main, picking up the ViRL39K v1 taskset (RE#689) and the taskset prompt config fixes (RE#742). verifiers main is unchanged since f6e420b99. Co-Authored-By: Claude Fable 5 --- deps/research-environments | 2 +- uv.lock | 28 ++++++++++++++++++++++++---- 2 files changed, 25 insertions(+), 5 deletions(-) diff --git a/deps/research-environments b/deps/research-environments index ccb375ce55..9a57c713c3 160000 --- a/deps/research-environments +++ b/deps/research-environments @@ -1 +1 @@ -Subproject commit ccb375ce550e48af9e898351145975521a2b6897 +Subproject commit 9a57c713c3770b81f669cd7e36c970c0df4cae57 diff --git a/uv.lock b/uv.lock index 7a1957a9b1..f61db79dd6 100644 --- a/uv.lock +++ b/uv.lock @@ -127,6 +127,7 @@ members = [ "unscramble-v1", "uuid-ctf-v1", "verbatim-copy-v1", + "virl39k-v1", "wideseek-v1", "wiki-search-v1", "wikispeedia-v1", @@ -1993,7 +1994,7 @@ wheels = [ [[package]] name = "forth-lang-v1" -version = "0.1.0" +version = "0.2.0" source = { editable = "deps/research-environments/environments/code/forth_lang_v1" } dependencies = [ { name = "datasets", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, @@ -2464,7 +2465,7 @@ wheels = [ [[package]] name = "hle-v1" -version = "0.2.0" +version = "0.3.0" source = { editable = "deps/research-environments/environments/knowledge/hle_v1" } dependencies = [ { name = "datasets", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, @@ -3667,7 +3668,7 @@ wheels = [ [[package]] name = "mcp-atlas-v1" -version = "0.2.0" +version = "0.3.0" source = { editable = "deps/research-environments/environments/tool_use/mcp_atlas_v1" } dependencies = [ { name = "datasets", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, @@ -5427,7 +5428,7 @@ requires-dist = [ [[package]] name = "prolog-v1" -version = "0.4.0" +version = "0.5.0" source = { editable = "deps/research-environments/environments/reasoning/prolog_v1" } dependencies = [ { name = "verifiers", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, @@ -7976,6 +7977,25 @@ examples = [ { name = "wordle-v1", editable = "deps/verifiers/environments/wordle_v1" }, ] +[[package]] +name = "virl39k-v1" +version = "0.1.0" +source = { editable = "deps/research-environments/environments/multimodal/virl39k_v1" } +dependencies = [ + { name = "datasets", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "huggingface-hub", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "math-verify", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "verifiers", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] + +[package.metadata] +requires-dist = [ + { name = "datasets" }, + { name = "huggingface-hub" }, + { name = "math-verify" }, + { name = "verifiers", specifier = ">=0.2.2.dev65" }, +] + [[package]] name = "virtualenv" version = "21.3.3" From b67c4908ad633eda98ac5940b2e43a3be3cdc52f Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Mon, 3 Aug 2026 19:50:13 +0000 Subject: [PATCH 14/58] chore: bump verifiers to 0.2.2.dev73 (per-rollout clients) f6e420b99 -> 576506d66 (7 commits). vf#2218 builds one client per rollout with a process-wide elastic renderer pool: TrainClientConfig drops pool_size for a multiplex knob, and the config module moved from clients.config to configs.client. The orchestrator config follows suit (pool_size -> multiplex, None = client default). Floors to dev73. Co-Authored-By: Claude Fable 5 --- deps/verifiers | 2 +- packages/prime-rl-configs/pyproject.toml | 2 +- .../src/prime_rl/configs/orchestrator.py | 19 ++++++++++--------- pyproject.toml | 2 +- src/prime_rl/orchestrator/utils.py | 2 +- src/prime_rl/utils/client.py | 18 ++++++++++-------- src/prime_rl/utils/elastic.py | 10 +++++----- .../orchestrator/test_orchestrator_setup.py | 8 ++++---- tests/unit/utils/test_client.py | 2 +- tests/unit/utils/test_elastic.py | 2 +- uv.lock | 2 +- 11 files changed, 36 insertions(+), 33 deletions(-) diff --git a/deps/verifiers b/deps/verifiers index f6e420b990..576506d66f 160000 --- a/deps/verifiers +++ b/deps/verifiers @@ -1 +1 @@ -Subproject commit f6e420b9908ae14d625f079881f13c15011ee1c9 +Subproject commit 576506d66f37a2d6d996a85b8eb2ef53bbfddd9f diff --git a/packages/prime-rl-configs/pyproject.toml b/packages/prime-rl-configs/pyproject.toml index aabadc1268..e7eb616b08 100644 --- a/packages/prime-rl-configs/pyproject.toml +++ b/packages/prime-rl-configs/pyproject.toml @@ -10,7 +10,7 @@ dependencies = [ "renderers>=0.1.9.dev10", "tomli>=2.2.1", "tomli-w>=1.2.0", - "verifiers>=0.2.2.dev66", + "verifiers>=0.2.2.dev73", ] [build-system] diff --git a/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py b/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py index 434a361f28..da829fcfbe 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py @@ -467,10 +467,11 @@ class OrchestratorConfig(BaseConfig): ``tokenizer.name_or_path`` via ``MODEL_RENDERER_MAP``. RL/OPD roll out through the renderer client; SFT uses it to backfill tokens for its chat-completions teacher.""" - pool_size: int | None = Field(None, ge=1) - """Number of renderer slots shared across concurrent rollouts. Bump - for long multi-turn prompts where client-side jinja tokenization - serializes.""" + multiplex: int | None = Field(None, ge=1) + """Concurrent rollouts that share one renderer. The pool warms one and grows on + demand, so this bounds tokenizer count at ~concurrency/multiplex rather than + fixing it. Lower it for long multi-turn prompts where client-side jinja + tokenization serializes; None keeps the client default.""" optim: OptimizerConfig = OptimizerConfig() """Per-run optimizer configuration for multi-run training.""" @@ -620,18 +621,18 @@ def any_policy_sourced(self) -> bool: return any(env.algo is not None and env.algo.sampling.source == "policy" for env in self.train.source) @model_validator(mode="after") - def validate_pool_size(self): - """``pool_size`` sizes the renderer-client pool for policy-sourced + def validate_multiplex(self): + """``multiplex`` sizes the renderer-client pool for policy-sourced sampling. Reject it when that path never runs — no train env samples from the policy — so callers don't silently pass it and wonder why it's ignored.""" - if self.pool_size is None: + if self.multiplex is None: return self if not self.any_policy_sourced: raise ValueError( - f"orchestrator.pool_size={self.pool_size!r} is set but no train env samples " + f"orchestrator.multiplex={self.multiplex!r} is set but no train env samples " "from the policy — the renderer-client sampling pool never runs (the renderer " - "is still used for client-side tokenization). Remove pool_size." + "is still used for client-side tokenization). Remove multiplex." ) return self diff --git a/pyproject.toml b/pyproject.toml index 9cbeeec4a9..4705ee165e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -33,7 +33,7 @@ dependencies = [ "setproctitle>=1.3.0", "uvloop>=0.21.0", "torchtitan", - "verifiers[harbor]>=0.2.2.dev66", + "verifiers[harbor]>=0.2.2.dev73", "renderers", "dion", "tilelang>=0.1.8", diff --git a/src/prime_rl/orchestrator/utils.py b/src/prime_rl/orchestrator/utils.py index 2b23369045..29f0559661 100644 --- a/src/prime_rl/orchestrator/utils.py +++ b/src/prime_rl/orchestrator/utils.py @@ -45,7 +45,7 @@ async def setup_policy_inference_pool(*, config: OrchestratorConfig, tokenizer): train_client_type="renderer", eval_client_type="openai_chat_completions", renderer_config=config.renderer, - pool_size=config.pool_size, + multiplex=config.multiplex, ) return renderer, inference_pool diff --git a/src/prime_rl/utils/client.py b/src/prime_rl/utils/client.py index c054e66cd9..e7675384b4 100644 --- a/src/prime_rl/utils/client.py +++ b/src/prime_rl/utils/client.py @@ -13,7 +13,7 @@ from openai import AsyncOpenAI from renderers import RendererConfig from tenacity import AsyncRetrying, retry, retry_if_exception, stop_after_attempt, stop_after_delay, wait_exponential -from verifiers.v1.clients.config import EvalClientConfig, TrainClientConfig +from verifiers.v1.configs.client import EvalClientConfig, TrainClientConfig from prime_rl.configs.shared import ClientConfig from prime_rl.utils.logger import get_logger @@ -122,7 +122,7 @@ def __init__( train_client_type: str = "openai_chat_completions", eval_client_type: str = "openai_chat_completions", renderer_config: RendererConfig | None = None, - pool_size: int | None = None, + multiplex: int | None = None, ): renderer_model_name = model_name if train_client_type == "renderer" else None self._train_clients = setup_clients( @@ -130,7 +130,7 @@ def __init__( client_type=train_client_type, renderer_config=renderer_config, renderer_model_name=renderer_model_name, - pool_size=pool_size, + multiplex=multiplex, ) self._eval_clients = setup_clients(client_config, client_type=eval_client_type) self._admin_clients = setup_admin_clients(client_config) @@ -195,7 +195,7 @@ async def setup_inference_pool( train_client_type: str = "openai_chat_completions", eval_client_type: str = "openai_chat_completions", renderer_config: RendererConfig | None = None, - pool_size: int | None = None, + multiplex: int | None = None, ) -> InferencePool: """Create an inference pool from config (static or elastic).""" if client_config.is_elastic: @@ -207,7 +207,7 @@ async def setup_inference_pool( train_client_type=train_client_type, eval_client_type=eval_client_type, renderer_config=renderer_config, - pool_size=pool_size, + multiplex=multiplex, ) return StaticInferencePool( @@ -216,7 +216,7 @@ async def setup_inference_pool( train_client_type=train_client_type, eval_client_type=eval_client_type, renderer_config=renderer_config, - pool_size=pool_size, + multiplex=multiplex, ) @@ -225,7 +225,7 @@ def setup_clients( client_type: str = "openai_chat_completions", renderer_config: RendererConfig | None = None, renderer_model_name: str | None = None, - pool_size: int | None = None, + multiplex: int | None = None, ) -> list[vf.ClientConfig]: """Build one v1 client config per base URL. ``client_type`` ``renderer`` → token-in/out (``TrainClientConfig``, with the renderer the env @@ -237,9 +237,11 @@ def setup_clients( if is_renderer: renderer_extra = { "renderer": renderer_config, - "pool_size": pool_size or 1, "renderer_model_name": renderer_model_name, } + # Unset leaves the config's own default; pinning a number here would override it. + if multiplex is not None: + renderer_extra["multiplex"] = multiplex env_headers = { k: v for k, v in ((k, os.getenv(v)) for k, v in client_config.headers_from_env.items()) if v is not None } diff --git a/src/prime_rl/utils/elastic.py b/src/prime_rl/utils/elastic.py index dbcf3c0825..ebbe4ca11d 100644 --- a/src/prime_rl/utils/elastic.py +++ b/src/prime_rl/utils/elastic.py @@ -116,7 +116,7 @@ def __init__( train_client_type: str = "openai_chat_completions", eval_client_type: str = "openai_chat_completions", renderer_config: RendererConfig | None = None, - pool_size: int | None = None, + multiplex: int | None = None, ): self.logger = get_logger() self.client_config = client_config @@ -129,7 +129,7 @@ def __init__( self.train_client_type = train_client_type self.eval_client_type = eval_client_type self.renderer_config = renderer_config - self.pool_size = pool_size + self.multiplex = multiplex self.router_url = client_config.router_url self._servers: dict[str, ServerState] = {} @@ -155,7 +155,7 @@ async def from_config( train_client_type: str = "openai_chat_completions", eval_client_type: str = "openai_chat_completions", renderer_config: RendererConfig | None = None, - pool_size: int | None = None, + multiplex: int | None = None, ) -> ElasticInferencePool: if client_config.elastic is None: raise ValueError("Elastic inference pool requires elastic config") @@ -165,7 +165,7 @@ async def from_config( train_client_type=train_client_type, eval_client_type=eval_client_type, renderer_config=renderer_config, - pool_size=pool_size, + multiplex=multiplex, ) await pool.start() return pool @@ -211,7 +211,7 @@ def _rebuild_clients(self) -> None: client_type=self.train_client_type, renderer_config=self.renderer_config, renderer_model_name=self.renderer_model_name, - pool_size=self.pool_size, + multiplex=self.multiplex, ) if urls else [] diff --git a/tests/unit/orchestrator/test_orchestrator_setup.py b/tests/unit/orchestrator/test_orchestrator_setup.py index 73768a7698..449c9f266e 100644 --- a/tests/unit/orchestrator/test_orchestrator_setup.py +++ b/tests/unit/orchestrator/test_orchestrator_setup.py @@ -17,7 +17,7 @@ async def run() -> None: name="policy-model", ), renderer=renderer_settings, - pool_size=None, + multiplex=None, any_policy_sourced=True, ) renderer = object() @@ -44,7 +44,7 @@ async def run() -> None: train_client_type="renderer", eval_client_type="openai_chat_completions", renderer_config=renderer_settings, - pool_size=None, + multiplex=None, ) asyncio.run(run()) @@ -65,7 +65,7 @@ async def run() -> None: name="policy-model", ), renderer=renderer_settings, - pool_size=None, + multiplex=None, any_policy_sourced=False, ) renderer = object() @@ -92,7 +92,7 @@ async def run() -> None: train_client_type="renderer", eval_client_type="openai_chat_completions", renderer_config=renderer_settings, - pool_size=None, + multiplex=None, ) asyncio.run(run()) diff --git a/tests/unit/utils/test_client.py b/tests/unit/utils/test_client.py index 0bcd606cc4..797169f905 100644 --- a/tests/unit/utils/test_client.py +++ b/tests/unit/utils/test_client.py @@ -3,7 +3,7 @@ from unittest.mock import AsyncMock, MagicMock, patch import httpx -from verifiers.v1.clients.config import EvalClientConfig +from verifiers.v1.configs.client import EvalClientConfig from prime_rl.configs.shared import ClientConfig from prime_rl.utils.client import _is_retryable_lora_error, check_health, load_lora_adapter, setup_clients diff --git a/tests/unit/utils/test_elastic.py b/tests/unit/utils/test_elastic.py index 0a1679b1e5..9530b0ac62 100644 --- a/tests/unit/utils/test_elastic.py +++ b/tests/unit/utils/test_elastic.py @@ -4,7 +4,7 @@ from unittest.mock import AsyncMock, MagicMock, patch import httpx -from verifiers.v1.clients.config import TrainClientConfig +from verifiers.v1.configs.client import TrainClientConfig from prime_rl.utils.elastic import ( AdapterState, diff --git a/uv.lock b/uv.lock index f61db79dd6..da67f439cd 100644 --- a/uv.lock +++ b/uv.lock @@ -5352,7 +5352,7 @@ requires-dist = [ { name = "renderers", specifier = ">=0.1.9.dev10" }, { name = "tomli", specifier = ">=2.2.1" }, { name = "tomli-w", specifier = ">=1.2.0" }, - { name = "verifiers", specifier = ">=0.2.2.dev66" }, + { name = "verifiers", specifier = ">=0.2.2.dev73" }, ] [[package]] From c85bd88a28a286221ac04d1266772e865276dfba Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Mon, 3 Aug 2026 20:03:12 +0000 Subject: [PATCH 15/58] chore!: pin research-environments to the scicode merge commit, drop orchestrator.multiplex deps/research-environments -> 21a284618 (research-environments#740 on main). The orchestrator-level renderer knob (pool_size, briefly multiplex) is removed instead of renamed: clients resolve inside env- server workers, whose concurrency is already bounded per env by serve.pool.multiplex (default 128, below the client-side renderer default of 256), so each worker warms one renderer and the elastic pool grows on demand. Co-Authored-By: Claude Fable 5 --- deps/research-environments | 2 +- .../src/prime_rl/configs/orchestrator.py | 22 ------------------- src/prime_rl/orchestrator/utils.py | 1 - src/prime_rl/utils/client.py | 9 -------- src/prime_rl/utils/elastic.py | 5 ----- .../orchestrator/test_orchestrator_setup.py | 4 ---- 6 files changed, 1 insertion(+), 42 deletions(-) diff --git a/deps/research-environments b/deps/research-environments index 9a57c713c3..21a284618f 160000 --- a/deps/research-environments +++ b/deps/research-environments @@ -1 +1 @@ -Subproject commit 9a57c713c3770b81f669cd7e36c970c0df4cae57 +Subproject commit 21a284618ff72091788237c517369675560cc447 diff --git a/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py b/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py index da829fcfbe..e9fdbf2a6e 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py @@ -467,12 +467,6 @@ class OrchestratorConfig(BaseConfig): ``tokenizer.name_or_path`` via ``MODEL_RENDERER_MAP``. RL/OPD roll out through the renderer client; SFT uses it to backfill tokens for its chat-completions teacher.""" - multiplex: int | None = Field(None, ge=1) - """Concurrent rollouts that share one renderer. The pool warms one and grows on - demand, so this bounds tokenizer count at ~concurrency/multiplex rather than - fixing it. Lower it for long multi-turn prompts where client-side jinja - tokenization serializes; None keeps the client default.""" - optim: OptimizerConfig = OptimizerConfig() """Per-run optimizer configuration for multi-run training.""" @@ -620,22 +614,6 @@ def any_policy_sourced(self) -> bool: """True when at least one train env samples rollouts from the live policy.""" return any(env.algo is not None and env.algo.sampling.source == "policy" for env in self.train.source) - @model_validator(mode="after") - def validate_multiplex(self): - """``multiplex`` sizes the renderer-client pool for policy-sourced - sampling. Reject it when that path never runs — no train env samples - from the policy — so callers don't silently pass it and wonder why - it's ignored.""" - if self.multiplex is None: - return self - if not self.any_policy_sourced: - raise ValueError( - f"orchestrator.multiplex={self.multiplex!r} is set but no train env samples " - "from the policy — the renderer-client sampling pool never runs (the renderer " - "is still used for client-side tokenization). Remove multiplex." - ) - return self - @model_validator(mode="after") def validate_renderer_auto_resolves(self): """Reject the silent DefaultRenderer fallback at config time. diff --git a/src/prime_rl/orchestrator/utils.py b/src/prime_rl/orchestrator/utils.py index 29f0559661..374a1070f0 100644 --- a/src/prime_rl/orchestrator/utils.py +++ b/src/prime_rl/orchestrator/utils.py @@ -45,7 +45,6 @@ async def setup_policy_inference_pool(*, config: OrchestratorConfig, tokenizer): train_client_type="renderer", eval_client_type="openai_chat_completions", renderer_config=config.renderer, - multiplex=config.multiplex, ) return renderer, inference_pool diff --git a/src/prime_rl/utils/client.py b/src/prime_rl/utils/client.py index e7675384b4..2d35f122ef 100644 --- a/src/prime_rl/utils/client.py +++ b/src/prime_rl/utils/client.py @@ -122,7 +122,6 @@ def __init__( train_client_type: str = "openai_chat_completions", eval_client_type: str = "openai_chat_completions", renderer_config: RendererConfig | None = None, - multiplex: int | None = None, ): renderer_model_name = model_name if train_client_type == "renderer" else None self._train_clients = setup_clients( @@ -130,7 +129,6 @@ def __init__( client_type=train_client_type, renderer_config=renderer_config, renderer_model_name=renderer_model_name, - multiplex=multiplex, ) self._eval_clients = setup_clients(client_config, client_type=eval_client_type) self._admin_clients = setup_admin_clients(client_config) @@ -195,7 +193,6 @@ async def setup_inference_pool( train_client_type: str = "openai_chat_completions", eval_client_type: str = "openai_chat_completions", renderer_config: RendererConfig | None = None, - multiplex: int | None = None, ) -> InferencePool: """Create an inference pool from config (static or elastic).""" if client_config.is_elastic: @@ -207,7 +204,6 @@ async def setup_inference_pool( train_client_type=train_client_type, eval_client_type=eval_client_type, renderer_config=renderer_config, - multiplex=multiplex, ) return StaticInferencePool( @@ -216,7 +212,6 @@ async def setup_inference_pool( train_client_type=train_client_type, eval_client_type=eval_client_type, renderer_config=renderer_config, - multiplex=multiplex, ) @@ -225,7 +220,6 @@ def setup_clients( client_type: str = "openai_chat_completions", renderer_config: RendererConfig | None = None, renderer_model_name: str | None = None, - multiplex: int | None = None, ) -> list[vf.ClientConfig]: """Build one v1 client config per base URL. ``client_type`` ``renderer`` → token-in/out (``TrainClientConfig``, with the renderer the env @@ -239,9 +233,6 @@ def setup_clients( "renderer": renderer_config, "renderer_model_name": renderer_model_name, } - # Unset leaves the config's own default; pinning a number here would override it. - if multiplex is not None: - renderer_extra["multiplex"] = multiplex env_headers = { k: v for k, v in ((k, os.getenv(v)) for k, v in client_config.headers_from_env.items()) if v is not None } diff --git a/src/prime_rl/utils/elastic.py b/src/prime_rl/utils/elastic.py index ebbe4ca11d..5d64c48794 100644 --- a/src/prime_rl/utils/elastic.py +++ b/src/prime_rl/utils/elastic.py @@ -116,7 +116,6 @@ def __init__( train_client_type: str = "openai_chat_completions", eval_client_type: str = "openai_chat_completions", renderer_config: RendererConfig | None = None, - multiplex: int | None = None, ): self.logger = get_logger() self.client_config = client_config @@ -129,7 +128,6 @@ def __init__( self.train_client_type = train_client_type self.eval_client_type = eval_client_type self.renderer_config = renderer_config - self.multiplex = multiplex self.router_url = client_config.router_url self._servers: dict[str, ServerState] = {} @@ -155,7 +153,6 @@ async def from_config( train_client_type: str = "openai_chat_completions", eval_client_type: str = "openai_chat_completions", renderer_config: RendererConfig | None = None, - multiplex: int | None = None, ) -> ElasticInferencePool: if client_config.elastic is None: raise ValueError("Elastic inference pool requires elastic config") @@ -165,7 +162,6 @@ async def from_config( train_client_type=train_client_type, eval_client_type=eval_client_type, renderer_config=renderer_config, - multiplex=multiplex, ) await pool.start() return pool @@ -211,7 +207,6 @@ def _rebuild_clients(self) -> None: client_type=self.train_client_type, renderer_config=self.renderer_config, renderer_model_name=self.renderer_model_name, - multiplex=self.multiplex, ) if urls else [] diff --git a/tests/unit/orchestrator/test_orchestrator_setup.py b/tests/unit/orchestrator/test_orchestrator_setup.py index 449c9f266e..60f13e01aa 100644 --- a/tests/unit/orchestrator/test_orchestrator_setup.py +++ b/tests/unit/orchestrator/test_orchestrator_setup.py @@ -17,7 +17,6 @@ async def run() -> None: name="policy-model", ), renderer=renderer_settings, - multiplex=None, any_policy_sourced=True, ) renderer = object() @@ -44,7 +43,6 @@ async def run() -> None: train_client_type="renderer", eval_client_type="openai_chat_completions", renderer_config=renderer_settings, - multiplex=None, ) asyncio.run(run()) @@ -65,7 +63,6 @@ async def run() -> None: name="policy-model", ), renderer=renderer_settings, - multiplex=None, any_policy_sourced=False, ) renderer = object() @@ -92,7 +89,6 @@ async def run() -> None: train_client_type="renderer", eval_client_type="openai_chat_completions", renderer_config=renderer_settings, - multiplex=None, ) asyncio.run(run()) From 40174c81030e3a12c6141c1868a200dc18b52048 Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Mon, 3 Aug 2026 21:01:57 +0000 Subject: [PATCH 16/58] feat(orchestrator)!: pipeline verdicts and eval scores are per-agent too is_trainable / is_filtered / filters and avg@k / pass@k all score a single trace, so they read per agent like every other trace-level metric instead of pooling an episode's seats. Co-Authored-By: Claude Fable 5 --- src/prime_rl/orchestrator/metrics.py | 91 ++++++++++++------------- src/prime_rl/utils/monitor/wandb.py | 6 +- tests/unit/orchestrator/test_metrics.py | 15 ++-- 3 files changed, 53 insertions(+), 59 deletions(-) diff --git a/src/prime_rl/orchestrator/metrics.py b/src/prime_rl/orchestrator/metrics.py index 656a19c22f..1bf656c02f 100644 --- a/src/prime_rl/orchestrator/metrics.py +++ b/src/prime_rl/orchestrator/metrics.py @@ -7,11 +7,11 @@ ``{prefix}/{subset}//`` wandb dict via ``.to_wandb(...)``. The wandb layout mirrors the episode/trace hierarchy. ``{prefix}/{subset}//`` carries -only episode-level facts: the count metrics sum an episode's traces (matching ``vf.Episode``'s -aggregates), plus the train pipeline rates and the eval scores. Every trace-level metric lives under -``{prefix}/{subset}///``, grouped by agent name (``vf.Episode.by_agent``) so -agents never mix into one distribution. A single-agent env has one trace per episode, so the count -metrics coincide across levels. +only episode-level facts: the count metrics, summing an episode's traces (matching ``vf.Episode``'s +aggregates). Every trace-level metric — reward, rates, timing, custom metrics, the pipeline verdicts, +and the eval scores — lives under ``{prefix}/{subset}///``, grouped by agent +name (``vf.Episode.by_agent``) so agents never mix into one distribution. A single-agent env has one +trace per episode, so the count metrics coincide across levels. No I/O, no pandas — plain Python over the ``vf.Trace`` properties each rollout exposes. """ @@ -334,43 +334,49 @@ def to_wandb(self, *, prefix: str, subset: Subset) -> dict[str, float]: class TrainMetrics(EpisodeMetrics): - """Common metrics plus the filter-pipeline rates. ``reward`` (flat over all traces) serves the - console log lines and distributions; the wandb reward stats are per-agent.""" + """Common metrics plus the per-agent filter-pipeline rates. ``reward`` (flat over all traces) + serves the console log lines; the wandb reward stats are per-agent.""" @property def reward(self) -> Stat: return Stat([float(r.reward) for r in self.rollouts]) - @property - def is_trainable(self) -> Stat: - return Stat([float(r.is_trainable) for r in self.rollouts]) - - @property - def is_filtered(self) -> Stat: - return Stat([float(r.is_filtered) for r in self.rollouts]) - - def filter_rates(self) -> dict[str, float]: - """Per-filter detection rate over all rollouts.""" - names = sorted({name for r in self.rollouts for name in r.filter_results}) - return { - name: sum(1 for r in self.rollouts if r.filter_results.get(name)) / len(self.rollouts) for name in names - } - def to_wandb(self, *, prefix: str, subset: Subset) -> dict[str, float]: out = super().to_wandb(prefix=prefix, subset=subset) - if not self.rollouts: - return out - p = f"{prefix}/{subset}" - out[f"{p}/is_trainable/mean"] = self.is_trainable.mean() - out[f"{p}/is_filtered/mean"] = self.is_filtered.mean() - out |= {f"{p}/filters/{k}/mean": v for k, v in self.filter_rates().items()} + # The pipeline verdicts are per-trace (an untrainable seat is 0.0 throughout, and filters + # only ever run on trainable survivors), so they read per agent like the rest. + for agent, traces in self.by_agent().items(): + p = f"{prefix}/{subset}/{agent}" + rollouts = traces.rollouts + out[f"{p}/is_trainable/mean"] = sum(float(r.is_trainable) for r in rollouts) / len(rollouts) + out[f"{p}/is_filtered/mean"] = sum(float(r.is_filtered) for r in rollouts) / len(rollouts) + names = sorted({name for r in rollouts for name in r.filter_results}) + out |= { + f"{p}/filters/{name}/mean": sum(1 for r in rollouts if r.filter_results.get(name)) / len(rollouts) + for name in names + } return out +def pass_at_k(rollouts: list[Rollout]) -> dict[str, float]: + """pass@k / pass^k averaged over examples; ``{}`` for non-binary rewards.""" + rewards = [r.reward for r in rollouts] + if not set(rewards).issubset({0.0, 1.0}): + return {} + by_example: dict = {} + for r in rollouts: + by_example.setdefault(r.group_id, []).append(r.reward) + per_example = [compute_pass_metrics(rs) for rs in by_example.values()] + keys = sorted({k for d in per_example for k in d}) + return {k: sum(d[k] for d in per_example if k in d) / sum(1 for d in per_example if k in d) for k in keys} + + class EvalMetrics(EpisodeMetrics): - """Common metrics plus the ``avg@`` score and (on the effective subset, for - binary-reward tasks) pass@k / pass^k. ``group_size`` (the ``avg@k`` k) is supplied by the - container so the ``all`` and ``effective`` subsets share one stable key.""" + """Common metrics plus the per-agent ``avg@`` score and (on the effective subset, + for binary-reward tasks) pass@k / pass^k. Both score an agent's own traces, so they live in its + subtree like every other trace-level metric. ``group_size`` (the ``avg@k`` k, the run's rollouts + per example) is supplied by the container so the ``all`` and ``effective`` subsets — and every + agent — share one stable key.""" def __init__(self, rollouts: list[Rollout], group_size: int) -> None: super().__init__(rollouts) @@ -380,26 +386,13 @@ def __init__(self, rollouts: list[Rollout], group_size: int) -> None: def reward(self) -> Stat: return Stat([float(r.reward) for r in self.rollouts]) - def pass_at_k(self) -> dict[str, float]: - """pass@k / pass^k averaged over examples; ``{}`` for non-binary rewards.""" - rewards = [r.reward for r in self.rollouts] - if not set(rewards).issubset({0.0, 1.0}): - return {} - by_example: dict = {} - for r in self.rollouts: - by_example.setdefault(r.group_id, []).append(r.reward) - per_example = [compute_pass_metrics(rs) for rs in by_example.values()] - keys = sorted({k for d in per_example for k in d}) - return {k: sum(d[k] for d in per_example if k in d) / sum(1 for d in per_example if k in d) for k in keys} - def to_wandb(self, *, prefix: str, subset: Subset) -> dict[str, float]: out = super().to_wandb(prefix=prefix, subset=subset) - if not self.rollouts: - return out - p = f"{prefix}/{subset}" - out[f"{p}/avg@{self.group_size}"] = self.reward.mean() - if subset == "effective": - out |= {f"{p}/{k}": v for k, v in self.pass_at_k().items()} + for agent, traces in self.by_agent().items(): + p = f"{prefix}/{subset}/{agent}" + out[f"{p}/avg@{self.group_size}"] = traces.stats()["reward"].mean() + if subset == "effective": + out |= {f"{p}/{k}": v for k, v in pass_at_k(traces.rollouts).items()} return out diff --git a/src/prime_rl/utils/monitor/wandb.py b/src/prime_rl/utils/monitor/wandb.py index 1669564f06..68b8eb5a32 100644 --- a/src/prime_rl/utils/monitor/wandb.py +++ b/src/prime_rl/utils/monitor/wandb.py @@ -366,11 +366,11 @@ def train_section(name: str, scope: str) -> ws.Section: def eval_section(name: str, env_pattern: str) -> ws.Section: - # Same metrics as train, but eval's reward is "avg@k" (dynamic k → regex). Everything is a regex so - # one section can also serve any env (env_pattern=".*"). + # Same metrics as train, but eval's reward is the per-agent "avg@k" (dynamic k → regex). + # Everything is a regex so one section can also serve any env (env_pattern=".*"). return section( name, - regexes=[f"eval/{env_pattern}/all/avg@.*", f"eval/{env_pattern}/effective/avg@.*"] + regexes=[f"eval/{env_pattern}/all/[^/]+/avg@.*", f"eval/{env_pattern}/effective/[^/]+/avg@.*"] + [f"eval/{env_pattern}/{m}" for m in COMMON_METRICS] + [f"eval/{env_pattern}/{r}" for r in COMMON_REGEXES], ) diff --git a/tests/unit/orchestrator/test_metrics.py b/tests/unit/orchestrator/test_metrics.py index f9f3f5beb5..e187f153b3 100644 --- a/tests/unit/orchestrator/test_metrics.py +++ b/tests/unit/orchestrator/test_metrics.py @@ -210,9 +210,10 @@ def test_train_only_metrics_absent_from_eval(): mk(is_trainable=False, filter_results={"gibberish": False}), ] out = train_wandb(rollouts) - assert out["train/agg/all/is_trainable/mean"] == 0.5 - assert out["train/agg/all/is_filtered/mean"] == 0.5 - assert out["train/agg/all/filters/gibberish/mean"] == 0.5 + assert out["train/agg/all/agent/is_trainable/mean"] == 0.5 + assert out["train/agg/all/agent/is_filtered/mean"] == 0.5 + assert out["train/agg/all/agent/filters/gibberish/mean"] == 0.5 + assert "train/agg/all/is_trainable/mean" not in out # pipeline verdicts are per-trace eval_out = EvalRollouts(rollouts).metrics.to_wandb(prefix="eval/x", subset="all") assert not any("is_trainable" in k or "is_filtered" in k or "/filters/" in k for k in eval_out) @@ -220,11 +221,11 @@ def test_train_only_metrics_absent_from_eval(): def test_eval_avg_at_k_and_pass_k(): binary = EvalRollouts([mk(reward=1.0, group_id="g0"), mk(reward=0.0, group_id="g0")]) eff = binary.effective.metrics.to_wandb(prefix="eval/x", subset="effective") - assert eff["eval/x/effective/avg@2"] == 0.5 # mean reward under avg@ (k derived from the groups) - assert not any(k.startswith("eval/x/effective/reward") for k in eff) - assert eff["eval/x/effective/pass@1"] == 0.5 and eff["eval/x/effective/pass^2"] == 0.0 + assert eff["eval/x/effective/agent/avg@2"] == 0.5 # mean reward under avg@ (k from the groups) + assert "eval/x/effective/avg@2" not in eff # scores are per-agent, never pooled + assert eff["eval/x/effective/agent/pass@1"] == 0.5 and eff["eval/x/effective/agent/pass^2"] == 0.0 all_out = binary.metrics.to_wandb(prefix="eval/x", subset="all") - assert all_out["eval/x/all/avg@2"] == 0.5 + assert all_out["eval/x/all/agent/avg@2"] == 0.5 assert not any("pass@" in k or "pass^" in k for k in all_out) # pass@k effective-only non_binary = EvalRollouts([mk(reward=0.5, group_id="g0"), mk(reward=1.0, group_id="g0")]) assert not any("pass@" in k for k in non_binary.effective.metrics.to_wandb(prefix="eval/x", subset="effective")) From 7dffc3690841a42c4dfb9861e3abd37324ef81ab Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Mon, 3 Aug 2026 21:10:37 +0000 Subject: [PATCH 17/58] chore: drop the agentic-judge debug config from this PR Co-Authored-By: Claude Fable 5 --- configs/debug/envs/README.md | 7 ---- configs/debug/envs/agentic_judge.toml | 51 --------------------------- configs/debug/envs/judge_policy.md | 23 ------------ 3 files changed, 81 deletions(-) delete mode 100644 configs/debug/envs/README.md delete mode 100644 configs/debug/envs/agentic_judge.toml delete mode 100644 configs/debug/envs/judge_policy.md diff --git a/configs/debug/envs/README.md b/configs/debug/envs/README.md deleted file mode 100644 index dadd90f5b7..0000000000 --- a/configs/debug/envs/README.md +++ /dev/null @@ -1,7 +0,0 @@ -# Envs — Debug Configs - -Minimal end-to-end configs for bundled multi-agent envs, using `PrimeIntellect/Qwen3-0.6B-Reverse-Text-SFT` as the policy. - -| Config | Env | Notes | -|---|---|---| -| `agentic_judge.toml` | `agentic-judge` over `reverse-text-v1` | solver in a docker box, frozen `deepseek/deepseek-v4-flash` judge grades in the same box via `judge_policy.md` — a two-tool-call mirror of the deterministic check (needs docker + a Prime Inference key; run from the repo root) | diff --git a/configs/debug/envs/agentic_judge.toml b/configs/debug/envs/agentic_judge.toml deleted file mode 100644 index c22435e203..0000000000 --- a/configs/debug/envs/agentic_judge.toml +++ /dev/null @@ -1,51 +0,0 @@ -max_steps = 20 -seq_len = 2048 - -[model] -name = "PrimeIntellect/Qwen3-0.6B-Reverse-Text-SFT" - -[wandb] -project = "reverse-text" -name = "reverse-text-agentic-judge" - -[orchestrator] -batch_size = 128 -group_size = 16 -# Each in-flight episode holds a docker box (solver + judge share it); unbounded -# in-flight stacks hundreds of writable layers and can fill the disk. -max_inflight_episodes = 64 - -[orchestrator.train.sampling] -max_completion_tokens = 128 - -# The agentic judge over reverse-text: the solver plays the task in a docker box, -# a frozen judge then grades the work in the same box and its verdict composes -# with the taskset's own reward on the solver's trace. -[[orchestrator.train.source]] -name = "reverse-text-judge" -env.id = "agentic-judge" -env.taskset = { id = "reverse-text-v1" } -env.solver.harness = { id = "null" } -env.solver.runtime = { type = "docker" } -env.judge.harness = { id = "bash" } -env.judge.model = "deepseek/deepseek-v4-flash" -env.judge.client = { type = "eval" } -# The run's sampling caps completions at the solver's 128 tokens — the judge -# reasons before writing its verdict, so it gets its own budget. -env.judge.sampling = { max_tokens = 4096 } -env.judge.max_turns = 8 -# Debug grading policy: mirror reverse-text's deterministic check in two tool -# calls, instead of the default open-ended trace investigation. The path is -# relative to the repo root — run from there. -env.task.prompt = "configs/debug/envs/judge_policy.md" -env.score = { task_weight = 1.0 } - -[orchestrator.renderer] -name = "prime-qwen3" - -[trainer.optim] -lr = 3e-6 - -[ckpt] - -[inference] diff --git a/configs/debug/envs/judge_policy.md b/configs/debug/envs/judge_policy.md deleted file mode 100644 index 4003fe3003..0000000000 --- a/configs/debug/envs/judge_policy.md +++ /dev/null @@ -1,23 +0,0 @@ -You are verifying a text-reversal attempt. This is a DEBUG judge that mirrors a -deterministic checker — do not investigate the trace beyond the recipe below, -and use AT MOST TWO tool calls total. - -## The task the agent was given - -{prompt} - -## Recipe - -First tool call — one python3 command that: - -1. loads `/tmp/trace.json`, -2. takes the task prompt (`data["task"]["data"]["prompt"]`) and the last - assistant message (`[n["message"] for n in data["nodes"] if n["message"]["role"] == "assistant"][-1]["content"]`), -3. extracts the answer: the text inside `...` if - those tags are present, the whole message otherwise, -4. prints `PASS` if the answer equals the prompt reversed character-by-character - (compare with surrounding whitespace stripped), else prints `FAIL` and both - strings. - -Second tool call — write your verdict file: `solved` is `yes` on `PASS`, `no` -otherwise. Then state your verdict and stop. From fc2b0c9b732a1e4030aadb8094679a16ce76389c Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Mon, 3 Aug 2026 21:11:23 +0000 Subject: [PATCH 18/58] docs: spell out the per-agent metric level in the monitor-run skill Co-Authored-By: Claude Fable 5 --- skills/training/monitor-run/SKILL.md | 26 ++++++++++++++++---------- 1 file changed, 16 insertions(+), 10 deletions(-) diff --git a/skills/training/monitor-run/SKILL.md b/skills/training/monitor-run/SKILL.md index 692ed08ab3..a377400179 100644 --- a/skills/training/monitor-run/SKILL.md +++ b/skills/training/monitor-run/SKILL.md @@ -89,19 +89,25 @@ grep -E "WARNING|ERROR" {output_dir}/logs/envs/{train,eval}/*.log All metrics print to the console log (and W&B when configured). -**Progress** — orchestrator log. Rollout metrics mirror the episode/trace hierarchy: `{scope}/{subset}//` carries episode-level facts only (token/turn/branch counts summed over an episode's traces, plus train pipeline rates and eval scores); every trace-level metric (reward, truncation, errors, timing, env metrics) lives under `{scope}/{subset}///`, per agent name — an in-episode fan-out like n solvers averages within the episode first. `scope` is `train/agg` (all train envs) or `train/` (`eval/` for eval); `subset` is `all` (every rollout) or `effective` (post-filter). Single-agent envs have one agent (usually `agent`) and one trace per episode. +**Progress** — orchestrator log. Rollout metrics mirror the episode/trace hierarchy, at two levels: + +- `{scope}/{subset}//` — episode-level facts only: the token/turn/branch counts, summed over an episode's traces. +- `{scope}/{subset}///` — every trace-level metric (reward, truncation, errors, timing, env metrics, filter verdicts, eval scores), keyed by agent name so seats never mix. An in-episode fan-out like n solvers averages within the episode first. + +`scope` is `train/agg` (all train envs) or `train/` (`eval/` for eval); `subset` is `all` (every rollout) or `effective` (post-filter). Single-agent envs have one agent — usually `agent` — and one trace per episode, so both levels agree; multi-agent envs name each seat (`proposer`, `solver`, `judge`, …). | Metric | Description | |--------|-------------| -| `train/agg/effective//reward/mean` | mean training reward (per env: `train//effective//reward/mean`) | -| `train/agg/effective/num_total_tokens/mean` | avg tokens per episode (also `num_input_tokens`, `num_output_tokens`) | -| `train/agg/effective/num_turns/mean` | avg turns per episode | -| `train//effective//num_turns/mean` | per-agent avg turns (also token counts, `num_branches`) | -| `train/agg/effective//is_truncated/mean` | fraction truncated | -| `train/agg/all//has_error/mean` | fraction errored (per-type under `train/agg/all//error/`; also `dispatcher/errored/{train,eval}`) | -| `train//effective//metrics//mean` | env-specific metrics (e.g. pass rate) | -| `train//effective//timing/agent/model/mean` | model vs harness share of the agent phase | -| `eval//effective/{avg@k,pass@k}` | eval scores when configured | +| `train/agg/effective//reward/mean` | mean training reward for that agent (per env: `train//effective//reward/mean`) | +| `train/agg/effective/num_total_tokens/mean` | avg tokens per episode, summed over its agents (also `num_input_tokens`, `num_output_tokens`) | +| `train/agg/effective/num_turns/mean` | avg turns per episode, summed over its agents | +| `train//effective//num_turns/mean` | avg turns for that agent alone (also token counts, `num_branches`) | +| `train/agg/effective//is_truncated/mean` | fraction of that agent's rollouts truncated | +| `train/agg/all//has_error/mean` | fraction of that agent's rollouts errored (per-type under `train/agg/all//error/`; also `dispatcher/errored/{train,eval}`) | +| `train/agg/all//is_trainable/mean` | fraction carrying a training signal — 0.0 for a frozen seat like a judge (also `is_filtered`, `filters/`) | +| `train//effective//metrics//mean` | env-specific metrics for that agent (e.g. pass rate) | +| `train//effective//timing/agent/model/mean` | model vs harness share of that agent's phase | +| `eval//effective//{avg@k,pass@k}` | eval scores for that agent, when configured | **Stability** — trainer log: From 4ab035dfa0bcda29b2f5e13cc6d55e8276bc21b5 Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Mon, 3 Aug 2026 21:31:49 +0000 Subject: [PATCH 19/58] docs: eval's overview score is the per-agent avg@k Co-Authored-By: Claude Fable 5 --- src/prime_rl/utils/monitor/wandb.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/prime_rl/utils/monitor/wandb.py b/src/prime_rl/utils/monitor/wandb.py index 68b8eb5a32..7340b3c6de 100644 --- a/src/prime_rl/utils/monitor/wandb.py +++ b/src/prime_rl/utils/monitor/wandb.py @@ -303,8 +303,8 @@ def save_final_summary(self, filename: str = "final_summary.json") -> None: # distributions. has_error only exists on all (effective drops errors by construction). The count # metrics are episode-level exact keys; the trace-level metrics (reward, truncation, errors) live # under the per-agent subtree, whose names are data-dependent — matched by regex, one panel per -# agent. Only the reward metrics differ — train uses the per-agent "reward/mean", eval uses the -# env-level "avg@k" — and each section builder prepends its own. +# agent. Only the score metric differs — train scores with "reward/mean", eval with "avg@k" (its k +# dynamic, so also a regex) — and each section builder prepends its own. COMMON_METRICS = [ "effective/num_total_tokens/mean", "effective/num_turns/mean", From 7c78ea4986419d5dbfe83ff58d4067dd38523a2c Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Mon, 3 Aug 2026 21:41:22 +0000 Subject: [PATCH 20/58] fix(orchestrator): agent-level rates are flat over traces MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit is_truncated / is_completed rode the per-episode mean path with the distributions, so an uneven fan-out — what the effective subset leaves whenever a sibling errors — reweighted them away from the plain fraction their siblings and the docs promise. Co-Authored-By: Claude Fable 5 --- src/prime_rl/orchestrator/metrics.py | 7 +++++-- tests/unit/orchestrator/test_metrics.py | 15 +++++++++++++++ 2 files changed, 20 insertions(+), 2 deletions(-) diff --git a/src/prime_rl/orchestrator/metrics.py b/src/prime_rl/orchestrator/metrics.py index 1bf656c02f..d8a01fa2cd 100644 --- a/src/prime_rl/orchestrator/metrics.py +++ b/src/prime_rl/orchestrator/metrics.py @@ -173,10 +173,13 @@ def __init__(self, episodes: list[list[Rollout]]) -> None: self.episodes = episodes def stats(self) -> dict[str, Stat]: + """One value per episode for the distributions (the fan-out collapses to its mean first); + one value per trace for the rates, so a rate stays the plain fraction of the agent's + rollouts — an episode that kept more traces than its siblings must not count for less.""" return { name: Stat([sum(float(getattr(r, name)) for r in episode) / len(episode) for episode in self.episodes]) - for name in (*self.DISTRIBUTIONS, *self.RATES) - } + for name in self.DISTRIBUTIONS + } | {name: Stat([float(getattr(r, name)) for r in self.rollouts]) for name in self.RATES} @property def timing(self) -> TimingMetrics: diff --git a/tests/unit/orchestrator/test_metrics.py b/tests/unit/orchestrator/test_metrics.py index e187f153b3..056af991db 100644 --- a/tests/unit/orchestrator/test_metrics.py +++ b/tests/unit/orchestrator/test_metrics.py @@ -146,6 +146,21 @@ def test_episode_and_agent_levels(): assert "train/agg/all/reward/mean" not in out # reward never pools across agents +def test_agent_rates_are_flat_over_traces(): + """An uneven fan-out (one surviving solver trace here, three there — what the effective subset + leaves behind whenever a sibling errors) must not reweight a rate: it stays the plain fraction + of the agent's rollouts, unlike the distributions, which collapse per episode first.""" + rollouts = [ + mk(agent_name="solver", episode_id="e1", is_truncated=True, num_turns=1), + *[mk(agent_name="solver", episode_id="e2", is_truncated=False, num_turns=1) for _ in range(3)], + ] + out = TrainRollouts(rollouts).metrics.to_wandb(prefix="train/agg", subset="all") + assert out["train/agg/all/solver/is_truncated/mean"] == 0.25 # 1 of 4 traces, not (1.0 + 0.0) / 2 + assert out["train/agg/all/solver/is_completed/mean"] == 1.0 + assert out["train/agg/all/solver/is_trainable/mean"] == 1.0 # sibling rates agree + assert out["train/agg/all/solver/num_turns/mean"] == 1.0 # distributions still collapse per episode + + def test_boolean_rates_and_error_breakdown_all_only(): rc = TrainRollouts([mk(is_truncated=True), mk(has_error=True, error_type="ProviderError"), mk(is_filtered=True)]) out = rc.metrics.to_wandb(prefix="train/agg", subset="all") From f81a178aeb7ac3ca45d53b58da46708f8787e462 Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Mon, 3 Aug 2026 22:21:54 +0000 Subject: [PATCH 21/58] refactor(orchestrator)!: agent-level metrics are flat over traces MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit One aggregation per level: the episode level sums an episode's traces, the agent level takes each trace as a sample. Reward, avg@k and the per-agent counts had been collapsing each episode's fan-out to a mean first, which left them the odd ones out among the seat's own metrics and made avg@k's k — a trace count — disagree with what it averaged. Co-Authored-By: Claude Fable 5 --- skills/training/monitor-run/SKILL.md | 2 +- src/prime_rl/orchestrator/metrics.py | 51 ++++++++++--------------- tests/unit/orchestrator/test_metrics.py | 20 +++++----- 3 files changed, 31 insertions(+), 42 deletions(-) diff --git a/skills/training/monitor-run/SKILL.md b/skills/training/monitor-run/SKILL.md index a377400179..70e95165b6 100644 --- a/skills/training/monitor-run/SKILL.md +++ b/skills/training/monitor-run/SKILL.md @@ -92,7 +92,7 @@ All metrics print to the console log (and W&B when configured). **Progress** — orchestrator log. Rollout metrics mirror the episode/trace hierarchy, at two levels: - `{scope}/{subset}//` — episode-level facts only: the token/turn/branch counts, summed over an episode's traces. -- `{scope}/{subset}///` — every trace-level metric (reward, truncation, errors, timing, env metrics, filter verdicts, eval scores), keyed by agent name so seats never mix. An in-episode fan-out like n solvers averages within the episode first. +- `{scope}/{subset}///` — every trace-level metric (reward, truncation, errors, timing, env metrics, filter verdicts, eval scores), keyed by agent name so seats never mix. Flat over that agent's traces: one sample is one trace, so an in-episode fan-out like n solvers contributes n samples. `scope` is `train/agg` (all train envs) or `train/` (`eval/` for eval); `subset` is `all` (every rollout) or `effective` (post-filter). Single-agent envs have one agent — usually `agent` — and one trace per episode, so both levels agree; multi-agent envs name each seat (`proposer`, `solver`, `judge`, …). diff --git a/src/prime_rl/orchestrator/metrics.py b/src/prime_rl/orchestrator/metrics.py index d8a01fa2cd..6cd11a4546 100644 --- a/src/prime_rl/orchestrator/metrics.py +++ b/src/prime_rl/orchestrator/metrics.py @@ -6,12 +6,15 @@ ``rollouts.metrics.num_input_tokens.mean()`` works — and assembles the full ``{prefix}/{subset}//`` wandb dict via ``.to_wandb(...)``. -The wandb layout mirrors the episode/trace hierarchy. ``{prefix}/{subset}//`` carries -only episode-level facts: the count metrics, summing an episode's traces (matching ``vf.Episode``'s -aggregates). Every trace-level metric — reward, rates, timing, custom metrics, the pipeline verdicts, -and the eval scores — lives under ``{prefix}/{subset}///``, grouped by agent -name (``vf.Episode.by_agent``) so agents never mix into one distribution. A single-agent env has one -trace per episode, so the count metrics coincide across levels. +The wandb layout mirrors the episode/trace hierarchy, one aggregation per level: + +- ``{prefix}/{subset}//`` — episode level, one value per episode summing its traces + (matching ``vf.Episode``'s aggregates). Only the count metrics live here. +- ``{prefix}/{subset}///`` — agent level, one value per trace, grouped by agent + name (``vf.Episode.by_agent``) so agents never mix into one distribution. Everything trace-scoped + lives here: reward, rates, timing, custom metrics, the pipeline verdicts, the eval scores. + +A single-agent env has one trace per episode, so the count metrics coincide across levels. No I/O, no pandas — plain Python over the ``vf.Trace`` properties each rollout exposes. """ @@ -159,27 +162,18 @@ def stats(self) -> dict[str, Stat]: class TraceMetrics(StatGroup): - """Trace-level metrics for one agent. The reward and count distributions carry one value per - episode — a fan-out (n same-agent traces in one episode, e.g. n solvers) collapses to its - within-episode mean first, so ``/num_turns/mean`` is the mean over episodes of the - agent's mean turns. Everything else (timing, custom metrics / reward components, rates, stop - conditions, errors, solve rates) is flat over the agent's traces.""" + """Trace-level metrics for one agent, every one of them flat over that agent's traces: one + sample is one trace, so a fan-out (n same-agent traces in one episode, e.g. n solvers) simply + contributes n samples. Weighing whole episodes against each other is the episode level's job; + inside a seat the trace is the unit, which is also what the advantage computation samples.""" DISTRIBUTIONS = ("reward", "num_total_tokens", "num_input_tokens", "num_output_tokens", "num_turns", "num_branches") RATES = ("is_truncated", "is_completed") - def __init__(self, episodes: list[list[Rollout]]) -> None: - super().__init__([r for episode in episodes for r in episode]) - self.episodes = episodes - def stats(self) -> dict[str, Stat]: - """One value per episode for the distributions (the fan-out collapses to its mean first); - one value per trace for the rates, so a rate stays the plain fraction of the agent's - rollouts — an episode that kept more traces than its siblings must not count for less.""" return { - name: Stat([sum(float(getattr(r, name)) for r in episode) / len(episode) for episode in self.episodes]) - for name in self.DISTRIBUTIONS - } | {name: Stat([float(getattr(r, name)) for r in self.rollouts]) for name in self.RATES} + name: Stat([float(getattr(r, name)) for r in self.rollouts]) for name in (*self.DISTRIBUTIONS, *self.RATES) + } @property def timing(self) -> TimingMetrics: @@ -277,16 +271,11 @@ def episodes(self) -> list[list[Rollout]]: return list(grouped.values()) def by_agent(self) -> dict[str, TraceMetrics]: - """Per-agent metric views (``vf.Episode.by_agent`` over the subset's rollouts): each - episode contributes the agent's traces in it.""" - per_agent: dict[str, list[list[Rollout]]] = {} - for episode in self.episodes(): - traces: dict[str, list[Rollout]] = {} - for r in episode: - traces.setdefault(r.agent.name, []).append(r) - for name, agent_traces in traces.items(): - per_agent.setdefault(name, []).append(agent_traces) - return {name: TraceMetrics(episodes) for name, episodes in sorted(per_agent.items())} + """Per-agent metric views (``vf.Episode.by_agent`` over the subset's rollouts).""" + per_agent: dict[str, list[Rollout]] = {} + for r in self.rollouts: + per_agent.setdefault(r.agent.name, []).append(r) + return {name: TraceMetrics(rollouts) for name, rollouts in sorted(per_agent.items())} # Episode-level count metrics: one value per episode, summing its traces — the same # aggregation as ``vf.Episode.num_turns`` / ``num_*_tokens``. diff --git a/tests/unit/orchestrator/test_metrics.py b/tests/unit/orchestrator/test_metrics.py index 056af991db..475d92be93 100644 --- a/tests/unit/orchestrator/test_metrics.py +++ b/tests/unit/orchestrator/test_metrics.py @@ -137,28 +137,28 @@ def test_episode_and_agent_levels(): assert m.num_total_tokens.values == [30.0, 30.0] # summed across the episode's traces out = m.to_wandb(prefix="train/agg", subset="all") assert out["train/agg/all/num_turns/mean"] == 12.0 - assert out["train/agg/all/proposer/num_turns/mean"] == 2.0 # one trace per episode: (1 + 3) / 2 - assert out["train/agg/all/solver/num_turns/mean"] == 5.0 # mean of per-episode fan-out means (3, 7) - assert out["train/agg/all/solver/num_turns/max"] == 7.0 + assert out["train/agg/all/proposer/num_turns/mean"] == 2.0 # (1 + 3) / 2 + assert out["train/agg/all/solver/num_turns/mean"] == 5.0 # flat over the 4 solver traces + assert out["train/agg/all/solver/num_turns/max"] == 8.0 # a real trace, not an episode mean assert out["train/agg/all/solver/reward/mean"] == 0.5 assert out["train/agg/all/proposer/is_truncated/mean"] == 0.0 assert "train/agg/all/proposer/is_truncated/p90" not in out # rates emit /mean only assert "train/agg/all/reward/mean" not in out # reward never pools across agents -def test_agent_rates_are_flat_over_traces(): - """An uneven fan-out (one surviving solver trace here, three there — what the effective subset - leaves behind whenever a sibling errors) must not reweight a rate: it stays the plain fraction - of the agent's rollouts, unlike the distributions, which collapse per episode first.""" +def test_agent_metrics_are_flat_over_traces(): + """Inside a seat the trace is the unit of aggregation, so an uneven fan-out (one solver trace + from this episode, three from that) never reweights anything: every agent-level metric is the + plain figure over that agent's rollouts.""" rollouts = [ - mk(agent_name="solver", episode_id="e1", is_truncated=True, num_turns=1), - *[mk(agent_name="solver", episode_id="e2", is_truncated=False, num_turns=1) for _ in range(3)], + mk(agent_name="solver", episode_id="e1", is_truncated=True, reward=1.0), + *[mk(agent_name="solver", episode_id="e2", is_truncated=False, reward=0.0) for _ in range(3)], ] out = TrainRollouts(rollouts).metrics.to_wandb(prefix="train/agg", subset="all") assert out["train/agg/all/solver/is_truncated/mean"] == 0.25 # 1 of 4 traces, not (1.0 + 0.0) / 2 assert out["train/agg/all/solver/is_completed/mean"] == 1.0 assert out["train/agg/all/solver/is_trainable/mean"] == 1.0 # sibling rates agree - assert out["train/agg/all/solver/num_turns/mean"] == 1.0 # distributions still collapse per episode + assert out["train/agg/all/solver/reward/mean"] == 0.25 # 1 of 4 traces scored, not (1.0 + 0.0) / 2 def test_boolean_rates_and_error_breakdown_all_only(): From a7eebcc994f756c3056fcbac5b6d4cdfe305e22b Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Mon, 3 Aug 2026 22:51:33 +0000 Subject: [PATCH 22/58] feat(orchestrator)!: cancellations and task failures are first-class A dispatched episode that never produced a trace used to be reported as a stand-in Rollout carrying a fake error. That trace had no real agent, so it defaulted to the 'agent' seat and its failure landed under a phantom subtree in a multi-agent env, inflating that seat's error rate while the real seats looked clean. EpisodeResult now carries an episode's scheduling facts plus either its traces or an EpisodeFailure, so the outcome is representable without inventing a rollout, and failures are counted at the episode level where they belong. Co-Authored-By: Claude Fable 5 --- src/prime_rl/orchestrator/dispatcher.py | 97 +++++++++++++---------- src/prime_rl/orchestrator/eval_sink.py | 26 +++--- src/prime_rl/orchestrator/metrics.py | 17 +++- src/prime_rl/orchestrator/orchestrator.py | 35 ++++---- src/prime_rl/orchestrator/train_sink.py | 27 ++++--- src/prime_rl/orchestrator/types.py | 42 +++++++++- tests/unit/orchestrator/test_metrics.py | 45 ++++++++++- 7 files changed, 210 insertions(+), 79 deletions(-) diff --git a/src/prime_rl/orchestrator/dispatcher.py b/src/prime_rl/orchestrator/dispatcher.py index 5fee725f6a..b8641da997 100644 --- a/src/prime_rl/orchestrator/dispatcher.py +++ b/src/prime_rl/orchestrator/dispatcher.py @@ -5,10 +5,11 @@ N rollouts in one call reserves N permits (each bridged v0 rollout is its own single-agent episode). - Optional rate limiting via ``AsyncLimiter(tasks_per_minute, 60)``. -- Emit-everything invariant: every dispatched episode eventually reaches - ``out_q`` exactly once, as a ``list[Rollout]``. Failures - (env error, empty trajectory, task exception, off-policy cancel) carry - ``trace.last_error`` set; sinks decide drop / partial-train policy. +- Emit-everything invariant: every dispatched episode eventually reaches ``out_q`` exactly once, + as an ``EpisodeResult``. An env-side failure (env error, empty trajectory) rides the trace that + hit it (``trace.last_error``); a prime-rl-side one (task exception, off-policy cancel) has no + trace to ride, so it arrives as ``EpisodeResult.failure`` rather than as a stand-in rollout. + Sinks decide drop / partial-train policy. - ``DispatcherMode.PREFER_TRAIN`` / ``PREFER_EVAL`` controls which kind to schedule next. Transitions are level-triggered (driven by the eval source's emptiness), so in-flight rollouts of the opposite kind drain @@ -40,6 +41,8 @@ from prime_rl.orchestrator.eval_source import EvalSource from prime_rl.orchestrator.train_source import TrainSource from prime_rl.orchestrator.types import ( + EpisodeFailure, + EpisodeResult, GroupState, InflightRollout, Policy, @@ -156,7 +159,7 @@ def __init__( # Bounded so the dispatcher backpressures on a slow sink. One entry per # episode — the sinks count episodes, never loose traces. - self.out_q: asyncio.Queue[list[Rollout]] = asyncio.Queue(maxsize=max(8, self.max_inflight)) + self.out_q: asyncio.Queue[EpisodeResult] = asyncio.Queue(maxsize=max(8, self.max_inflight)) self.mode: DispatcherMode = DispatcherMode.PREFER_TRAIN # Set by the orchestrator after the final train step; pipeline then @@ -509,27 +512,21 @@ async def handle_completed_rollout(self, task: asyncio.Task) -> None: self.release(meta.rollout_count) group = self.groups.get(meta.group_id) - is_synth_exception = False try: result = task.result() rollouts: list[Rollout] = result if isinstance(result, list) else [result] - if not rollouts: - raise RuntimeError("env run returned an empty episode (no traces)") except asyncio.CancelledError: return except Exception as exc: get_logger().warning(f"Rollout task failed in group {meta.group_id} ({meta.env_name}): {exc!r}") - task_idx = group.task_idx if group is not None else -1 - rollouts = [ - Rollout( - task=vf.TraceTask(type="Task", data=vf.TaskData(idx=task_idx, prompt=None)), - agent=vf.AgentInfo(config=vf.AgentConfig()), - ) - for _ in range(meta.rollout_count) - ] - for r in rollouts: - r.record_error(exc) - is_synth_exception = True + await self.emit_failed_episodes(meta, group, EpisodeFailure(type="TaskFailed", message=repr(exc))) + return + if not rollouts: + get_logger().warning(f"Env returned no traces in group {meta.group_id} ({meta.env_name})") + await self.emit_failed_episodes( + meta, group, EpisodeFailure(type="EmptyEpisode", message="env run returned an episode with no traces") + ) + return for r in rollouts: if not r.has_error and r.num_turns == 0: @@ -540,7 +537,7 @@ async def handle_completed_rollout(self, task: asyncio.Task) -> None: get_logger().warning(f"Empty trajectory in group {meta.group_id} ({meta.env_name})") if r.has_error: self.metrics.record_error(kind=meta.kind, env_name=meta.env_name) - if not is_synth_exception and r.last_error is not None: + if r.last_error is not None: get_logger().warning( f"Rollout failed in group {meta.group_id} ({meta.env_name}) — {r.last_error.type}: {r.last_error.message}" ) @@ -552,10 +549,17 @@ async def handle_completed_rollout(self, task: asyncio.Task) -> None: for r in rollouts: await self.emit_episode(meta, group, [r]) - async def emit_episode(self, meta: InflightRollout, group: GroupState | None, rollouts: list[Rollout]) -> None: - """Stamp prime-rl metadata onto one completed episode and put it on - ``out_q``. Pops the group from ``self.groups`` once every owed episode - has been emitted.""" + async def emit_episode( + self, + meta: InflightRollout, + group: GroupState | None, + rollouts: list[Rollout], + failure: EpisodeFailure | None = None, + ) -> None: + """Put one completed episode on ``out_q`` — its traces, or the ``failure`` that stands in + for them. Pops the group from ``self.groups`` once every owed episode has been emitted. + The scheduling facts also ride each trace, where the sinks and the saved records read + them.""" eval_step = meta.eval_step policy_version = meta.policy_version if group is not None: @@ -564,6 +568,8 @@ async def emit_episode(self, meta: InflightRollout, group: GroupState | None, ro group.emitted += 1 if group.emitted >= group.target_rollouts: self.groups.pop(meta.group_id, None) + if meta.kind == "eval": + assert eval_step is not None, "eval episode missing eval_step" for rollout in rollouts: rollout.kind = meta.kind @@ -572,9 +578,28 @@ async def emit_episode(self, meta: InflightRollout, group: GroupState | None, ro rollout.policy_version = policy_version rollout.off_policy_steps = meta.off_policy_steps if meta.kind == "eval": - assert eval_step is not None, "eval rollout missing eval_step" rollout.eval_step = eval_step - await self.out_q.put(rollouts) + await self.out_q.put( + EpisodeResult( + kind=meta.kind, + env_name=meta.env_name, + group_id=meta.group_id, + policy_version=policy_version, + off_policy_steps=meta.off_policy_steps, + eval_step=eval_step if meta.kind == "eval" else None, + rollouts=rollouts, + failure=failure, + ) + ) + + async def emit_failed_episodes( + self, meta: InflightRollout, group: GroupState | None, failure: EpisodeFailure + ) -> None: + """Emit one failed episode per rollout the task owed, so the sink still counts its way to + ``group_size`` — prime-rl's own failure, carried as itself rather than as a stand-in trace.""" + for _ in range(meta.rollout_count): + self.metrics.record_error(kind=meta.kind, env_name=meta.env_name) + await self.emit_episode(meta, group, [], failure=failure) async def drop_group(self, group_id: uuid.UUID) -> int: """Cancel remaining in-flight tasks for this group and emit a @@ -582,7 +607,6 @@ async def drop_group(self, group_id: uuid.UUID) -> int: (both in-flight and not-yet-scheduled). Returns the count for off-policy metrics.""" group = self.groups.pop(group_id, None) - task_idx = group.task_idx if group is not None else -1 # Sync claim phase: pop matching tasks from ``self.inflight`` and # release their permits in one non-yielding sweep. After this loop @@ -600,16 +624,10 @@ async def drop_group(self, group_id: uuid.UUID) -> int: tasks_to_cancel = [task for task, _ in claimed] inflight_cancelled = sum(meta.rollout_count for _, meta in claimed) last_meta: InflightRollout | None = claimed[-1][1] if claimed else None + cancel = EpisodeFailure(type="Cancelled", message="Off-policy cancel") for _, meta in claimed: for _ in range(meta.rollout_count): - trace = Rollout( - task=vf.TraceTask(type="Task", data=vf.TaskData(idx=task_idx, prompt=None)), - agent=vf.AgentInfo(config=vf.AgentConfig()), - ok=False, - errors=[vf.Error(type="Cancelled", message="Off-policy cancel")], - stop_condition="error", - ) - await self.emit_episode(meta, group, [trace]) + await self.emit_episode(meta, group, [], failure=cancel) # For non-group-scoring envs, the group may have rollouts that # were never dispatched (``rollouts_to_schedule > 0``). Emit @@ -630,14 +648,7 @@ async def drop_group(self, group_id: uuid.UUID) -> int: ) unscheduled_cancelled = group.rollouts_to_schedule for _ in range(unscheduled_cancelled): - trace = Rollout( - task=vf.TraceTask(type="Task", data=vf.TaskData(idx=task_idx, prompt=None)), - agent=vf.AgentInfo(config=vf.AgentConfig()), - ok=False, - errors=[vf.Error(type="Cancelled", message="Off-policy cancel")], - stop_condition="error", - ) - await self.emit_episode(fallback_meta, group, [trace]) + await self.emit_episode(fallback_meta, group, [], failure=cancel) cancelled = inflight_cancelled + unscheduled_cancelled if cancelled > 0: diff --git a/src/prime_rl/orchestrator/eval_sink.py b/src/prime_rl/orchestrator/eval_sink.py index c02ebab7b9..7e93a45ceb 100644 --- a/src/prime_rl/orchestrator/eval_sink.py +++ b/src/prime_rl/orchestrator/eval_sink.py @@ -8,7 +8,7 @@ 3. ``process_batch`` — at ``num_examples × group_size`` episodes, return an ``EvalBatch`` with the full returned cohort (metrics are computed downstream). -``add()`` takes one episode (``list[Rollout]``) and returns ``EvalBatch | None``; +``add()`` takes one ``EpisodeResult`` and returns ``EvalBatch | None``; all accounting counts episodes, never loose traces. """ @@ -19,7 +19,7 @@ from prime_rl.orchestrator.envs import EvalEnvs from prime_rl.orchestrator.metrics import EvalRollouts -from prime_rl.orchestrator.types import EvalBatch, Rollout +from prime_rl.orchestrator.types import EpisodeFailure, EpisodeResult, EvalBatch, Rollout from prime_rl.utils.logger import get_logger @@ -33,16 +33,21 @@ def __init__(self, *, eval_envs: EvalEnvs) -> None: self.pending_group_episodes: dict[uuid.UUID, int] = defaultdict(int) self.pending_batches: dict[tuple[str, int], list[Rollout]] = defaultdict(list) self.pending_batch_episodes: dict[tuple[str, int], int] = defaultdict(int) + # Episodes of the epoch that produced no traces at all (cancelled, or the task failed). + self.pending_batch_failures: dict[tuple[str, int], list[EpisodeFailure]] = defaultdict(list) - def add(self, episode: list[Rollout]) -> EvalBatch | None: + def add(self, episode: EpisodeResult) -> EvalBatch | None: """Process one episode arrival; finalize the group on the ``group_size``-th - episode and the per-env epoch on the ``num_examples × group_size``-th.""" - env_name = episode[0].env_name - group_id = episode[0].group_id - for rollout in episode: + episode and the per-env epoch on the ``num_examples × group_size``-th. A failed + episode brings no rollouts but still counts toward both.""" + env_name = episode.env_name + group_id = episode.group_id + for rollout in episode.rollouts: self.process_rollout(rollout) - bkey = (env_name, episode[0].eval_step) - self.pending_groups[group_id].extend(episode) + bkey = (env_name, episode.eval_step) + if episode.failure is not None: + self.pending_batch_failures[bkey].append(episode.failure) + self.pending_groups[group_id].extend(episode.rollouts) self.pending_group_episodes[group_id] += 1 if self.pending_group_episodes[group_id] >= self.group_size_for(env_name): self.process_group(group_id) @@ -125,4 +130,5 @@ def process_batch(self, key: tuple[str, int]) -> EvalBatch: env_name, step = key rollouts = self.pending_batches.pop(key, []) self.pending_batch_episodes.pop(key, None) - return EvalBatch(env_name=env_name, step=step, rollouts=EvalRollouts(rollouts)) + failures = self.pending_batch_failures.pop(key, []) + return EvalBatch(env_name=env_name, step=step, rollouts=EvalRollouts(rollouts), failures=failures) diff --git a/src/prime_rl/orchestrator/metrics.py b/src/prime_rl/orchestrator/metrics.py index 4e67190dc3..cb5c8cc184 100644 --- a/src/prime_rl/orchestrator/metrics.py +++ b/src/prime_rl/orchestrator/metrics.py @@ -26,11 +26,26 @@ from prime_rl.orchestrator.utils import compute_pass_metrics if TYPE_CHECKING: - from prime_rl.orchestrator.types import Rollout + from prime_rl.orchestrator.types import EpisodeFailure, Rollout Subset = Literal["all", "effective"] +def episode_failure_metrics(failures: list[EpisodeFailure], *, prefix: str, total: int) -> dict[str, float]: + """Episodes that produced no traces, by reason (``{prefix}/episode_failure/``) plus their + share of the window. These are prime-rl's own outcomes — a cancellation or a task that never + reached the env — so they belong to no agent and are counted whole rather than folded into a + seat's error rate. ``total`` is the window's rollout count, for the rate's denominator.""" + counts: dict[str, int] = {} + for failure in failures: + counts[failure.type] = counts.get(failure.type, 0) + 1 + out = {f"{prefix}/episode_failure/{name}": float(n) for name, n in sorted(counts.items())} + denominator = total + len(failures) + if denominator: + out[f"{prefix}/episode_failure/rate"] = len(failures) / denominator + return out + + class Stat: """A distribution of per-rollout values with mean/max/min and p10/p90 accessors.""" diff --git a/src/prime_rl/orchestrator/orchestrator.py b/src/prime_rl/orchestrator/orchestrator.py index d6b8d07e19..19314f2838 100644 --- a/src/prime_rl/orchestrator/orchestrator.py +++ b/src/prime_rl/orchestrator/orchestrator.py @@ -3,8 +3,8 @@ ``Orchestrator`` owns the shared state (policy, progress, ckpt, monitor) and drives the pipeline. Components are single-purpose: -- ``RolloutDispatcher`` schedules rollouts; emits ``Rollout`` (train/eval - discriminated by ``kind``) on its queue. +- ``RolloutDispatcher`` schedules rollouts; emits ``EpisodeResult`` (train/eval + discriminated by ``kind``) on its queue — its traces, or the failure that replaced them. - ``TrainSink`` ingests train rollouts (tokenize → advantages → filters) and returns a ``TrainBatch`` when the threshold is met. - ``EvalSink`` ingests eval rollouts and returns an ``EvalBatch`` (the full @@ -49,6 +49,7 @@ from prime_rl.orchestrator.eval_source import EvalSource from prime_rl.orchestrator.filters import setup_filters from prime_rl.orchestrator.inference_metrics import InferenceMetricsCollector +from prime_rl.orchestrator.metrics import episode_failure_metrics from prime_rl.orchestrator.patches import ( monkey_patch_chat_completion_logprobs, monkey_patch_oai_iterable_types, @@ -57,10 +58,10 @@ from prime_rl.orchestrator.train_sink import TrainSink from prime_rl.orchestrator.train_source import TrainSource from prime_rl.orchestrator.types import ( + EpisodeResult, EvalBatch, Policy, Progress, - Rollout, TrainBatch, ) from prime_rl.orchestrator.utils import ( @@ -510,7 +511,7 @@ async def start(self) -> None: trim_process_memory() async def main_loop(self) -> None: - """Consume episodes (``list[Rollout]``) from the dispatcher and route them + """Consume ``EpisodeResult``\\ s from the dispatcher and route them to the train / eval sink. Both sinks return a finalized batch (or ``None``) from ``add()``; we just dispatch on the result.""" while not self.stopped.is_set(): @@ -520,23 +521,24 @@ async def main_loop(self) -> None: break try: - episode: list[Rollout] = await asyncio.wait_for(self.dispatcher.out_q.get(), timeout=0.5) + episode: EpisodeResult = await asyncio.wait_for(self.dispatcher.out_q.get(), timeout=0.5) except asyncio.TimeoutError: continue # Every completed rollout — errored, filtered, or never batched — lands in the # ``all`` trace file the moment it arrives, so it survives crashes and drains. # Train rollouts belong to the batch window currently collecting (``progress.step``), - # eval rollouts to the step whose eval triggered them. - kind = episode[0].kind - step = episode[0].eval_step if kind == "eval" else self.progress.step + # eval rollouts to the step whose eval triggered them. A failed episode has no trace + # to write; it is counted in the batch's failure tally instead. + kind = episode.kind + step = episode.eval_step if kind == "eval" else self.progress.step assert step is not None run: vf.RunInfo = ( vf.EvalRunInfo(id=self.run_id, step=step) if kind == "eval" else vf.TrainRunInfo(id=self.run_id, step=step) ) - for rollout in episode: + for rollout in episode.rollouts: rollout.record_run( run, env_name=rollout.env_name, @@ -544,11 +546,12 @@ async def main_loop(self) -> None: episode_id=rollout.episode_id, policy_version=rollout.policy_version, ) - await asyncio.to_thread( - save_rollouts, - [rollout.to_record() for rollout in episode], - get_trace_path(self.config.output_dir, step, kind, "all"), - ) + if episode.rollouts: + await asyncio.to_thread( + save_rollouts, + [rollout.to_record() for rollout in episode.rollouts], + get_trace_path(self.config.output_dir, step, kind, "all"), + ) if kind == "eval": assert self.eval_sink is not None # eval rollouts only emitted when eval is configured @@ -660,6 +663,9 @@ async def finalize_train_batch(self, batch: TrainBatch) -> None: metrics |= pool.metrics.to_wandb(prefix="train/agg", subset=subset) for env_name, env_pool in pool.by_env().items(): metrics |= env_pool.metrics.to_wandb(prefix=f"train/{env_name}", subset=subset) + # Episodes that produced no trace at all are an episode-level fact — they belong to no + # agent, so they are counted here rather than folded into any seat's error rate. + metrics |= episode_failure_metrics(batch.failures, prefix="train/agg", total=len(batch.rollouts)) # Progress / timing / env-share / pre-filter accounting (assembled here, not in the metrics # objects). ``num_tokens`` is over the full arrival window; the input/output breakdown is over @@ -876,6 +882,7 @@ async def finalize_eval_batch(self, batch: EvalBatch) -> None: metrics: dict[str, float] = {} for subset, pool in (("all", rollouts), ("effective", effective)): metrics |= pool.metrics.to_wandb(prefix=f"eval/{batch.env_name}", subset=subset) + metrics |= episode_failure_metrics(batch.failures, prefix=f"eval/{batch.env_name}", total=len(rollouts)) metrics[f"eval/{batch.env_name}/policy_version"] = float(policy_version) metrics["step"] = float(batch.step) self.monitor.log(metrics, step=batch.step) diff --git a/src/prime_rl/orchestrator/train_sink.py b/src/prime_rl/orchestrator/train_sink.py index fb0dc6abc1..fc60b8b733 100644 --- a/src/prime_rl/orchestrator/train_sink.py +++ b/src/prime_rl/orchestrator/train_sink.py @@ -10,7 +10,7 @@ 3. ``process_batch`` — applies post-batch filter annotations and assembles the trainer-bound ``TrainingSample`` list. Returns a ``TrainBatch``. -``add()`` takes one episode (``list[Rollout]``) and returns +``add()`` takes one ``EpisodeResult`` and returns ``TrainBatch | None``; group accounting counts episodes, never loose traces. I/O concerns (ship to trainer, save_rollouts, monitor.log) live on the orchestrator. @@ -27,7 +27,7 @@ from prime_rl.orchestrator.filters import RolloutFilter, apply_filters from prime_rl.orchestrator.metrics import TrainRollouts from prime_rl.orchestrator.trajectories import trace_to_samples -from prime_rl.orchestrator.types import Rollout, TrainBatch +from prime_rl.orchestrator.types import EpisodeFailure, EpisodeResult, Rollout, TrainBatch from prime_rl.transport import TrainingSample from prime_rl.utils.logger import get_logger @@ -76,6 +76,10 @@ def __init__( # finalized since the last ship (errored + filtered + survivors). # In-progress groups stay out until they finalize. self.pending_rollouts: TrainRollouts = TrainRollouts() + # Episodes in that window that produced no traces at all (cancelled, or the task + # itself failed). They have no rollout to be counted through, so they are tallied + # here and reported alongside the window's metrics. + self.pending_failures: list[EpisodeFailure] = [] # Keyed by the dispatcher's group UUID. ``(env_name, task_idx)`` # isn't unique — the same task can be re-sampled while an # earlier group is still in flight @@ -124,16 +128,19 @@ def pending_batch_by_env(self) -> dict[str, int]: counts[r.env_name] += 1 return dict(counts) - async def add(self, episode: list[Rollout]) -> TrainBatch | None: + async def add(self, episode: EpisodeResult) -> TrainBatch | None: """Process one episode arrival; finalize the group on the ``group_size``-th episode; return a ``TrainBatch`` if the finalization pushed (or left) the batch over its threshold. Arrivals into - still-incomplete groups never ship a batch.""" - group_id = episode[0].group_id - env_name = episode[0].env_name - for rollout in episode: + still-incomplete groups never ship a batch. A failed episode brings no + rollouts, but still counts toward the group so finalization triggers.""" + group_id = episode.group_id + env_name = episode.env_name + if episode.failure is not None: + self.pending_failures.append(episode.failure) + for rollout in episode.rollouts: await self.process_rollout(rollout) - self.pending_groups[group_id].extend(episode) + self.pending_groups[group_id].extend(episode.rollouts) self.pending_group_episodes[group_id] += 1 if self.pending_group_episodes[group_id] < self.group_size_for(env_name): return None @@ -286,9 +293,11 @@ def process_batch(self) -> TrainBatch: # samples) — an empty batch is dropped unlogged by the orchestrator, so keep accumulating its # finalized groups (and any overflow) into the next shipped batch's window. rollouts = self.pending_rollouts + failures = self.pending_failures if samples: self.pending_rollouts = TrainRollouts() - return TrainBatch(rollouts=rollouts, samples=samples) + self.pending_failures = [] + return TrainBatch(rollouts=rollouts, samples=samples, failures=failures) def reset_pre_filter_stats(self) -> None: self.pre_filter_seen = 0 diff --git a/src/prime_rl/orchestrator/types.py b/src/prime_rl/orchestrator/types.py index 65c56d3181..016f5f7979 100644 --- a/src/prime_rl/orchestrator/types.py +++ b/src/prime_rl/orchestrator/types.py @@ -3,7 +3,7 @@ from __future__ import annotations import uuid -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import TYPE_CHECKING, Generic, Literal, Protocol import verifiers.v1 as vf @@ -140,6 +140,41 @@ def is_trainable(self) -> bool: return bool(self.advantages) and any(a != 0.0 for a in self.advantages) +@dataclass +class EpisodeFailure: + """Why a dispatched episode carries no traces of its own. This is prime-rl's own verdict — + the episode never reached the env, or was pulled back before it finished — as opposed to an + env-side error, which rides on the trace that hit it (``Trace.last_error``).""" + + type: Literal["Cancelled", "TaskFailed", "EmptyEpisode"] + message: str + + +@dataclass +class EpisodeResult: + """One dispatched episode as it arrives from the dispatcher: the scheduling facts that + describe the whole episode, plus either the env's traces or the reason there are none. + + The scheduling facts live here rather than on each trace because that is what they describe — + a group is dispatched, cancelled, and counted whole. ``rollouts`` is empty exactly when + ``failure`` is set: prime-rl never mints a stand-in trace to carry an outcome the env did not + produce, so a cancellation can't be mistaken for an agent's rollout.""" + + kind: RolloutKind + env_name: str + group_id: uuid.UUID + policy_version: int + off_policy_steps: int = 0 + eval_step: int | None = None + rollouts: list[Rollout] = field(default_factory=list) + failure: EpisodeFailure | None = None + + def __post_init__(self) -> None: + assert bool(self.rollouts) != (self.failure is not None), ( + "an episode carries either traces or a failure, never both or neither" + ) + + @dataclass class TrainBatch: """``rollouts`` is the observation window since the last ship — every rollout of every group @@ -151,6 +186,9 @@ class TrainBatch: rollouts: TrainRollouts samples: list[TrainingSample] + failures: list[EpisodeFailure] = field(default_factory=list) + """Episodes in the window that produced no traces at all — they appear nowhere in + ``rollouts``, so their count would otherwise be invisible.""" @dataclass @@ -161,6 +199,8 @@ class EvalBatch: env_name: str step: int rollouts: EvalRollouts + failures: list[EpisodeFailure] = field(default_factory=list) + """Episodes in the epoch that produced no traces at all.""" class VersionObserver(Protocol): diff --git a/tests/unit/orchestrator/test_metrics.py b/tests/unit/orchestrator/test_metrics.py index 546e594d9a..f5ed8a8d43 100644 --- a/tests/unit/orchestrator/test_metrics.py +++ b/tests/unit/orchestrator/test_metrics.py @@ -1,11 +1,13 @@ import math from itertools import count from types import SimpleNamespace +from uuid import uuid4 import pytest import verifiers.v1 as vf -from prime_rl.orchestrator.metrics import EvalRollouts, Stat, TrainRollouts +from prime_rl.orchestrator.metrics import EvalRollouts, Stat, TrainRollouts, episode_failure_metrics +from prime_rl.orchestrator.types import EpisodeFailure, EpisodeResult from prime_rl.orchestrator.utils import compute_pass_metrics _ids = count() @@ -265,3 +267,44 @@ def test_compute_pass_metrics_matches_closed_form(): assert out["pass@2"] == 1.0 - math.comb(2, 2) / math.comb(4, 2) assert out["pass^2"] == math.comb(2, 2) / math.comb(4, 2) assert set(out) == {"pass@1", "pass@2", "pass@4", "pass^1", "pass^2", "pass^4"} + + +def test_episode_failure_metrics(): + """Episodes that never produced a trace are counted whole, by reason. They belong to no agent, + so they must never appear under an agent subtree — the phantom-seat bug this replaced.""" + failures = [ + EpisodeFailure(type="Cancelled", message="Off-policy cancel"), + EpisodeFailure(type="Cancelled", message="Off-policy cancel"), + EpisodeFailure(type="TaskFailed", message="boom"), + ] + out = episode_failure_metrics(failures, prefix="train/agg", total=9) + assert out["train/agg/episode_failure/Cancelled"] == 2.0 + assert out["train/agg/episode_failure/TaskFailed"] == 1.0 + assert out["train/agg/episode_failure/rate"] == 0.25 # 3 of the window's 12 episodes + assert not any("/agent/" in k for k in out) # never attributed to a seat + assert episode_failure_metrics([], prefix="train/agg", total=4) == {"train/agg/episode_failure/rate": 0.0} + assert episode_failure_metrics([], prefix="train/agg", total=0) == {} # nothing arrived at all + + +def test_episode_result_carries_traces_xor_failure(): + """The envelope makes the phantom trace unrepresentable: an outcome is traces or a failure.""" + rollout = mk() + assert EpisodeResult(kind="train", env_name="env", group_id=uuid4(), policy_version=0, rollouts=[rollout]) + assert EpisodeResult( + kind="train", + env_name="env", + group_id=uuid4(), + policy_version=0, + failure=EpisodeFailure(type="Cancelled", message="Off-policy cancel"), + ) + with pytest.raises(AssertionError): # neither + EpisodeResult(kind="train", env_name="env", group_id=uuid4(), policy_version=0) + with pytest.raises(AssertionError): # both + EpisodeResult( + kind="train", + env_name="env", + group_id=uuid4(), + policy_version=0, + rollouts=[rollout], + failure=EpisodeFailure(type="Cancelled", message="Off-policy cancel"), + ) From d1141e90148d73186a8dfd56c5c5ee2b892be633 Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Mon, 3 Aug 2026 23:01:22 +0000 Subject: [PATCH 23/58] refactor(orchestrator): keep the env's vf.Episode instead of rebuilding it envs.py used to read .traces/.id/.ok/.last_error off the episode and throw the envelope away, so metrics.py regrouped the flattened traces on episode_id and hand-rolled by_agent plus the token/turn sums that vf.Episode already exposes. The episode now rides through: Env.run returns it with its traces re-typed, EpisodeResult carries it, the sinks bucket episodes rather than loose rollouts, and the metric containers narrow an episode to its surviving traces for a subset view. The count metrics and by_agent are now reads off vf.Episode. Co-Authored-By: Claude Fable 5 --- src/prime_rl/orchestrator/dispatcher.py | 29 +++--- src/prime_rl/orchestrator/envs.py | 25 ++--- src/prime_rl/orchestrator/eval_sink.py | 28 +++--- src/prime_rl/orchestrator/metrics.py | 116 +++++++++++++----------- src/prime_rl/orchestrator/train_sink.py | 20 ++-- src/prime_rl/orchestrator/types.py | 18 ++-- tests/unit/orchestrator/test_metrics.py | 84 ++++++++++------- 7 files changed, 183 insertions(+), 137 deletions(-) diff --git a/src/prime_rl/orchestrator/dispatcher.py b/src/prime_rl/orchestrator/dispatcher.py index b8641da997..f38a805554 100644 --- a/src/prime_rl/orchestrator/dispatcher.py +++ b/src/prime_rl/orchestrator/dispatcher.py @@ -46,7 +46,6 @@ GroupState, InflightRollout, Policy, - Rollout, RolloutKind, ) from prime_rl.utils.async_utils import safe_cancel, safe_cancel_all @@ -514,21 +513,21 @@ async def handle_completed_rollout(self, task: asyncio.Task) -> None: try: result = task.result() - rollouts: list[Rollout] = result if isinstance(result, list) else [result] + episodes: list[vf.WireEpisode] = result if isinstance(result, list) else [result] except asyncio.CancelledError: return except Exception as exc: get_logger().warning(f"Rollout task failed in group {meta.group_id} ({meta.env_name}): {exc!r}") await self.emit_failed_episodes(meta, group, EpisodeFailure(type="TaskFailed", message=repr(exc))) return - if not rollouts: + if not any(episode.traces for episode in episodes): get_logger().warning(f"Env returned no traces in group {meta.group_id} ({meta.env_name})") await self.emit_failed_episodes( meta, group, EpisodeFailure(type="EmptyEpisode", message="env run returned an episode with no traces") ) return - for r in rollouts: + for r in (r for episode in episodes for r in episode.traces): if not r.has_error and r.num_turns == 0: # Empty trajectory: promote to an explicit error so the sink # treats it like any other failure (``has_error`` reads ``ok``) @@ -541,19 +540,15 @@ async def handle_completed_rollout(self, task: asyncio.Task) -> None: get_logger().warning( f"Rollout failed in group {meta.group_id} ({meta.env_name}) — {r.last_error.type}: {r.last_error.message}" ) - if meta.rollout_count == 1: - # A ``run`` task: the whole result is one episode. - await self.emit_episode(meta, group, rollouts) - else: - # A legacy ``run_group`` task: one single-trace episode per trace. - for r in rollouts: - await self.emit_episode(meta, group, [r]) + # A ``run`` task answers one episode; a legacy ``run_group`` task one per rollout. + for episode in episodes: + await self.emit_episode(meta, group, episode) async def emit_episode( self, meta: InflightRollout, group: GroupState | None, - rollouts: list[Rollout], + episode: vf.WireEpisode | None, failure: EpisodeFailure | None = None, ) -> None: """Put one completed episode on ``out_q`` — its traces, or the ``failure`` that stands in @@ -571,7 +566,7 @@ async def emit_episode( if meta.kind == "eval": assert eval_step is not None, "eval episode missing eval_step" - for rollout in rollouts: + for rollout in episode.traces if episode is not None else []: rollout.kind = meta.kind rollout.env_name = meta.env_name rollout.group_id = meta.group_id @@ -587,7 +582,7 @@ async def emit_episode( policy_version=policy_version, off_policy_steps=meta.off_policy_steps, eval_step=eval_step if meta.kind == "eval" else None, - rollouts=rollouts, + episode=episode, failure=failure, ) ) @@ -599,7 +594,7 @@ async def emit_failed_episodes( ``group_size`` — prime-rl's own failure, carried as itself rather than as a stand-in trace.""" for _ in range(meta.rollout_count): self.metrics.record_error(kind=meta.kind, env_name=meta.env_name) - await self.emit_episode(meta, group, [], failure=failure) + await self.emit_episode(meta, group, None, failure=failure) async def drop_group(self, group_id: uuid.UUID) -> int: """Cancel remaining in-flight tasks for this group and emit a @@ -627,7 +622,7 @@ async def drop_group(self, group_id: uuid.UUID) -> int: cancel = EpisodeFailure(type="Cancelled", message="Off-policy cancel") for _, meta in claimed: for _ in range(meta.rollout_count): - await self.emit_episode(meta, group, [], failure=cancel) + await self.emit_episode(meta, group, None, failure=cancel) # For non-group-scoring envs, the group may have rollouts that # were never dispatched (``rollouts_to_schedule > 0``). Emit @@ -648,7 +643,7 @@ async def drop_group(self, group_id: uuid.UUID) -> int: ) unscheduled_cancelled = group.rollouts_to_schedule for _ in range(unscheduled_cancelled): - await self.emit_episode(fallback_meta, group, [], failure=cancel) + await self.emit_episode(fallback_meta, group, None, failure=cancel) cancelled = inflight_cancelled + unscheduled_cancelled if cancelled > 0: diff --git a/src/prime_rl/orchestrator/envs.py b/src/prime_rl/orchestrator/envs.py index 25baafdb00..c27fa1d42c 100644 --- a/src/prime_rl/orchestrator/envs.py +++ b/src/prime_rl/orchestrator/envs.py @@ -44,6 +44,8 @@ # fields without importing the env package — the orchestrator never reads them typed (only # task.idx + task.model_dump). ROLLOUT_TYPE = Rollout[vf.WireTaskData] +# The env server answers a ``WireEpisode``; we keep that envelope and only re-type its traces. +EPISODE_TYPE = vf.WireEpisode # Max wait for a spawned env server to bind and report its address. A legacy # child loads its dataset before reporting, so this is generous. @@ -203,12 +205,15 @@ async def run( cache_salt: str | None, task_data: dict | None = None, task_idx: int | None = None, - ) -> list[Rollout]: - """Run one episode; return its typed Traces. A v1 env takes the task itself - (``task_data``); the legacy bridge is addressed by dataset row (``task_idx``). - A zero-trace episode raises (the dispatcher synthesizes the error marker); a - not-``ok`` episode marks its clean traces failed so partial episodes never - train.""" + ) -> EPISODE_TYPE: + """Run one episode; return it with its traces re-typed as ``Rollout``. A v1 env takes the + task itself (``task_data``); the legacy bridge is addressed by dataset row (``task_idx``). + A zero-trace episode comes back as-is for the dispatcher to report as a failure; a + not-``ok`` episode marks its clean traces failed so partial episodes never train. + + The episode itself is kept rather than flattened away: it is what groups an env's seats, + and downstream reads its aggregates (``by_agent``, ``num_turns``) instead of rebuilding + them from loose traces.""" episode = await self.env_client.run( task_data=task_data, task_idx=task_idx, @@ -216,10 +221,6 @@ async def run( model=model_name, sampling=self._sampling(cache_salt), ) - if not episode.traces: - error = episode.last_error - detail = f"{error.type}: {error.message}" if error is not None else "no traces and no error recorded" - raise RuntimeError(f"episode failed before any trace was produced — {detail}") rollouts = [ROLLOUT_TYPE.model_construct(**dict(wire)) for wire in episode.traces] for rollout in rollouts: rollout.episode_id = episode.id @@ -229,7 +230,7 @@ async def run( ) rollout.errors = [*rollout.errors, error] rollout.ok = False - return rollouts + return episode.model_copy(update={"traces": rollouts}) async def run_group( self, client: vf.ClientConfig, task_idx: int, model_name: str, group_size: int, cache_salt: str | None @@ -242,7 +243,7 @@ async def run_group( model=model_name, sampling=self._sampling(cache_salt), ) - return [ROLLOUT_TYPE.model_construct(**dict(wire)) for wire in wires] + return [EPISODE_TYPE.model_construct(traces=[ROLLOUT_TYPE.model_construct(**dict(wire))]) for wire in wires] def shutdown(self) -> None: if self._env_server_process is None: diff --git a/src/prime_rl/orchestrator/eval_sink.py b/src/prime_rl/orchestrator/eval_sink.py index 7e93a45ceb..8a6eaae10c 100644 --- a/src/prime_rl/orchestrator/eval_sink.py +++ b/src/prime_rl/orchestrator/eval_sink.py @@ -17,6 +17,8 @@ import uuid from collections import defaultdict +import verifiers.v1 as vf + from prime_rl.orchestrator.envs import EvalEnvs from prime_rl.orchestrator.metrics import EvalRollouts from prime_rl.orchestrator.types import EpisodeFailure, EpisodeResult, EvalBatch, Rollout @@ -28,10 +30,10 @@ class EvalSink: def __init__(self, *, eval_envs: EvalEnvs) -> None: self.eval_envs = eval_envs - self.pending_groups: dict[uuid.UUID, list[Rollout]] = defaultdict(list) + self.pending_groups: dict[uuid.UUID, list[vf.WireEpisode]] = defaultdict(list) # Episodes arrived per group / per batch bucket — the finalization counts. self.pending_group_episodes: dict[uuid.UUID, int] = defaultdict(int) - self.pending_batches: dict[tuple[str, int], list[Rollout]] = defaultdict(list) + self.pending_batches: dict[tuple[str, int], list[vf.WireEpisode]] = defaultdict(list) self.pending_batch_episodes: dict[tuple[str, int], int] = defaultdict(int) # Episodes of the epoch that produced no traces at all (cancelled, or the task failed). self.pending_batch_failures: dict[tuple[str, int], list[EpisodeFailure]] = defaultdict(list) @@ -47,7 +49,8 @@ def add(self, episode: EpisodeResult) -> EvalBatch | None: bkey = (env_name, episode.eval_step) if episode.failure is not None: self.pending_batch_failures[bkey].append(episode.failure) - self.pending_groups[group_id].extend(episode.rollouts) + if episode.episode is not None: + self.pending_groups[group_id].append(episode.episode) self.pending_group_episodes[group_id] += 1 if self.pending_group_episodes[group_id] >= self.group_size_for(env_name): self.process_group(group_id) @@ -71,13 +74,13 @@ def batch_progress(self) -> list[tuple[str, int, int, int, int]]: ``buffered`` is partial-group episode arrivals from non-group-scoring envs.""" batch_counts: dict[tuple[str, int], int] = dict(self.pending_batch_episodes) buffered: dict[tuple[str, int], int] = {} - for group_id, rollouts in self.pending_groups.items(): - if not rollouts: + for group_id, episodes in self.pending_groups.items(): + if not episodes: continue - env_name = rollouts[0].env_name + env_name = episodes[0].traces[0].env_name if self.eval_envs.get(env_name).requires_group_scoring: continue - bkey = (env_name, rollouts[0].eval_step) + bkey = (env_name, episodes[0].traces[0].eval_step) buffered[bkey] = buffered.get(bkey, 0) + self.pending_group_episodes.get(group_id, 0) return [ ( @@ -102,15 +105,16 @@ def process_rollout(self, rollout: Rollout) -> None: # ── level 2: per-group (move into batch bucket) ─────────────────────── def process_group(self, group_id: uuid.UUID) -> None: - group = self.pending_groups.pop(group_id, []) + finished = self.pending_groups.pop(group_id, []) episodes = self.pending_group_episodes.pop(group_id, 0) - if not group: + if not finished: return + group = [trace for episode in finished for trace in episode.traces] env_name = group[0].env_name task_idx = group[0].task.data.idx eval_step = group[0].eval_step bucket = self.pending_batches[(env_name, eval_step)] - bucket.extend(group) + bucket.extend(finished) self.pending_batch_episodes[(env_name, eval_step)] += episodes survivors = [r for r in group if not r.has_error] @@ -128,7 +132,7 @@ def process_batch(self, key: tuple[str, int]) -> EvalBatch: downstream via ``EvalBatch.rollouts.metrics`` over the all/effective subsets, so the sink does no aggregation.""" env_name, step = key - rollouts = self.pending_batches.pop(key, []) + episodes = self.pending_batches.pop(key, []) self.pending_batch_episodes.pop(key, None) failures = self.pending_batch_failures.pop(key, []) - return EvalBatch(env_name=env_name, step=step, rollouts=EvalRollouts(rollouts), failures=failures) + return EvalBatch(env_name=env_name, step=step, rollouts=EvalRollouts(episodes), failures=failures) diff --git a/src/prime_rl/orchestrator/metrics.py b/src/prime_rl/orchestrator/metrics.py index cb5c8cc184..a92b199fa7 100644 --- a/src/prime_rl/orchestrator/metrics.py +++ b/src/prime_rl/orchestrator/metrics.py @@ -23,6 +23,8 @@ from typing import TYPE_CHECKING, Any, Callable, Iterator, Literal +import verifiers.v1 as vf + from prime_rl.orchestrator.utils import compute_pass_metrics if TYPE_CHECKING: @@ -274,52 +276,45 @@ def to_dict(self, prefix: str, *, subset: Subset) -> dict[str, float]: class EpisodeMetrics: - """Metrics shared by train and eval over a rollout list. The count metrics (tokens/turns/ - branches) are episode-level — one value per episode, summing its traces — and are the only - per-metric keys ``to_wandb`` emits at the env level; every trace-level metric is emitted per - agent via ``by_agent()``. The boolean ``Stat`` properties (0/1 distributions, ``.mean()`` is - the rate) serve the console log lines. ``TrainMetrics`` / ``EvalMetrics`` extend ``to_wandb`` - with the train pipeline rates and the eval scores.""" - - def __init__(self, rollouts: list[Rollout]) -> None: - self.rollouts = rollouts + """Metrics shared by train and eval over a list of episodes. The count metrics (tokens/turns/ + branches) are read straight off ``vf.Episode``, which already sums an episode's traces, and are + the only per-metric keys ``to_wandb`` emits at the env level; every trace-level metric is + emitted per agent via ``by_agent()``. The boolean ``Stat`` properties (0/1 distributions, + ``.mean()`` is the rate) serve the console log lines. ``TrainMetrics`` / ``EvalMetrics`` extend + ``to_wandb`` with the train pipeline rates and the eval scores.""" - def episodes(self) -> list[list[Rollout]]: - """The subset's rollouts grouped into their episodes. A rollout without an - ``episode_id`` (legacy envs, synthesized error markers) is its own episode.""" - grouped: dict[str, list[Rollout]] = {} - for r in self.rollouts: - grouped.setdefault(r.episode_id or r.id, []).append(r) - return list(grouped.values()) + def __init__(self, episodes: list[vf.WireEpisode]) -> None: + self.episodes = episodes + self.rollouts: list[Rollout] = [trace for episode in episodes for trace in episode.traces] def by_agent(self) -> dict[str, TraceMetrics]: - """Per-agent metric views (``vf.Episode.by_agent`` over the subset's rollouts).""" + """Per-agent metric views, merging each episode's own ``vf.Episode.by_agent`` grouping.""" per_agent: dict[str, list[Rollout]] = {} - for r in self.rollouts: - per_agent.setdefault(r.agent.name, []).append(r) + for episode in self.episodes: + for name, traces in episode.by_agent.items(): + per_agent.setdefault(name, []).extend(traces) return {name: TraceMetrics(rollouts) for name, rollouts in sorted(per_agent.items())} - # Episode-level count metrics: one value per episode, summing its traces — the same - # aggregation as ``vf.Episode.num_turns`` / ``num_*_tokens``. + # Episode-level count metrics, one value per episode — ``vf.Episode``'s own aggregates. @property def num_total_tokens(self) -> Stat: - return Stat([float(sum(r.num_total_tokens for r in episode)) for episode in self.episodes()]) + return Stat([float(episode.num_total_tokens) for episode in self.episodes]) @property def num_input_tokens(self) -> Stat: - return Stat([float(sum(r.num_input_tokens for r in episode)) for episode in self.episodes()]) + return Stat([float(episode.num_input_tokens) for episode in self.episodes]) @property def num_output_tokens(self) -> Stat: - return Stat([float(sum(r.num_output_tokens for r in episode)) for episode in self.episodes()]) + return Stat([float(episode.num_output_tokens) for episode in self.episodes]) @property def num_turns(self) -> Stat: - return Stat([float(sum(r.num_turns for r in episode)) for episode in self.episodes()]) + return Stat([float(episode.num_turns) for episode in self.episodes]) @property def num_branches(self) -> Stat: - return Stat([float(sum(r.num_branches for r in episode)) for episode in self.episodes()]) + return Stat([float(sum(trace.num_branches for trace in episode.traces)) for episode in self.episodes]) # Boolean rate metrics for the console log lines (0/1 distributions — ``.mean()`` is the # rate); to_wandb emits their per-agent counterparts instead. @@ -392,8 +387,8 @@ class EvalMetrics(EpisodeMetrics): per example) is supplied by the container so the ``all`` and ``effective`` subsets — and every agent — share one stable key.""" - def __init__(self, rollouts: list[Rollout], group_size: int) -> None: - super().__init__(rollouts) + def __init__(self, episodes: list[vf.WireEpisode], group_size: int) -> None: + super().__init__(episodes) self.group_size = group_size @property @@ -410,50 +405,70 @@ def to_wandb(self, *, prefix: str, subset: Subset) -> dict[str, float]: return out +def keep_traces(episode: vf.WireEpisode, keep: Callable[[Rollout], bool]) -> vf.WireEpisode | None: + """The episode narrowed to the traces that pass ``keep``, or ``None`` if none do. A subset view + stays a list of episodes so the episode-level aggregates keep describing what survived.""" + traces = [trace for trace in episode.traces if keep(trace)] + return episode.model_copy(update={"traces": traces}) if traces else None + + class TrainRollouts: - """A list of train rollouts (everything that came back, errored + filtered + untrainable - included). ``effective`` is the clean trainable subset (a view of the same traces); - ``metrics`` builds ``TrainMetrics`` over them.""" + """The train episodes of one window (everything that came back, errored + filtered + + untrainable included). ``effective`` is the clean trainable subset — the same episodes, each + narrowed to its surviving traces; ``metrics`` builds ``TrainMetrics`` over them. Sized and + iterated by *rollout*, since that is the unit almost every consumer wants.""" + + def __init__(self, episodes: list[vf.WireEpisode] | None = None) -> None: + self.episodes = episodes if episodes is not None else [] - def __init__(self, rollouts: list[Rollout] | None = None) -> None: - self.rollouts = rollouts if rollouts is not None else [] + def append(self, episode: vf.WireEpisode) -> None: + self.episodes.append(episode) - def append(self, rollout: Rollout) -> None: - self.rollouts.append(rollout) + @property + def rollouts(self) -> list[Rollout]: + return [trace for episode in self.episodes for trace in episode.traces] def __len__(self) -> int: - return len(self.rollouts) + return sum(len(episode.traces) for episode in self.episodes) def __iter__(self) -> Iterator[Rollout]: return iter(self.rollouts) @property def effective(self) -> TrainRollouts: - return TrainRollouts([r for r in self.rollouts if not r.has_error and not r.is_filtered and r.agent.trainable]) + kept = ( + keep_traces(episode, lambda r: not r.has_error and not r.is_filtered and r.agent.trainable) + for episode in self.episodes + ) + return TrainRollouts([episode for episode in kept if episode is not None]) def by_env(self) -> dict[str, TrainRollouts]: - grouped: dict[str, list[Rollout]] = {} - for r in self.rollouts: - grouped.setdefault(r.env_name, []).append(r) - return {env: TrainRollouts(rs) for env, rs in grouped.items()} + grouped: dict[str, list[vf.WireEpisode]] = {} + for episode in self.episodes: + grouped.setdefault(episode.traces[0].env_name, []).append(episode) + return {env: TrainRollouts(episodes) for env, episodes in grouped.items()} @property def metrics(self) -> TrainMetrics: - return TrainMetrics(self.rollouts) + return TrainMetrics(self.episodes) class EvalRollouts: - """A list of eval rollouts (errored + untrainable included). ``effective`` is the - non-errored trainable subset (a view). + """The eval episodes of one epoch (errored + untrainable included). ``effective`` is the + non-errored trainable subset — the same episodes, each narrowed to its surviving traces. ``group_size`` (rollouts per example, the ``avg@k`` k) is derived from the full epoch and carried onto ``effective`` so both subsets share one stable key; ``metrics`` builds ``EvalMetrics``.""" - def __init__(self, rollouts: list[Rollout] | None = None, group_size: int | None = None) -> None: - self.rollouts = rollouts if rollouts is not None else [] + def __init__(self, episodes: list[vf.WireEpisode] | None = None, group_size: int | None = None) -> None: + self.episodes = episodes if episodes is not None else [] self._group_size = group_size + @property + def rollouts(self) -> list[Rollout]: + return [trace for episode in self.episodes for trace in episode.traces] + def __len__(self) -> int: - return len(self.rollouts) + return sum(len(episode.traces) for episode in self.episodes) def __iter__(self) -> Iterator[Rollout]: return iter(self.rollouts) @@ -474,10 +489,9 @@ def group_size(self) -> int: @property def effective(self) -> EvalRollouts: - return EvalRollouts( - [r for r in self.rollouts if not r.has_error and r.agent.trainable], group_size=self.group_size - ) + kept = (keep_traces(episode, lambda r: not r.has_error and r.agent.trainable) for episode in self.episodes) + return EvalRollouts([episode for episode in kept if episode is not None], group_size=self.group_size) @property def metrics(self) -> EvalMetrics: - return EvalMetrics(self.rollouts, self.group_size) + return EvalMetrics(self.episodes, self.group_size) diff --git a/src/prime_rl/orchestrator/train_sink.py b/src/prime_rl/orchestrator/train_sink.py index fc60b8b733..35d0c56f77 100644 --- a/src/prime_rl/orchestrator/train_sink.py +++ b/src/prime_rl/orchestrator/train_sink.py @@ -22,6 +22,8 @@ import uuid from collections import defaultdict +import verifiers.v1 as vf + from prime_rl.configs.orchestrator import OrchestratorConfig from prime_rl.orchestrator.envs import TrainEnvs from prime_rl.orchestrator.filters import RolloutFilter, apply_filters @@ -83,7 +85,7 @@ def __init__( # Keyed by the dispatcher's group UUID. ``(env_name, task_idx)`` # isn't unique — the same task can be re-sampled while an # earlier group is still in flight - self.pending_groups: dict[uuid.UUID, list[Rollout]] = defaultdict(list) + self.pending_groups: dict[uuid.UUID, list[vf.WireEpisode]] = defaultdict(list) # Episodes arrived per group — the finalization count (an episode may # add several traces to ``pending_groups`` but counts once here). self.pending_group_episodes: dict[uuid.UUID, int] = defaultdict(int) @@ -116,8 +118,8 @@ def buffered_count(self) -> int: (non-group-scoring envs) — buffered in the sink ahead of the batch.""" return sum( self.pending_group_episodes.get(group_id, 0) - for group_id, rollouts in self.pending_groups.items() - if rollouts and not self.train_envs.get(rollouts[0].env_name).requires_group_scoring + for group_id, episodes in self.pending_groups.items() + if episodes and not self.train_envs.get(episodes[0].traces[0].env_name).requires_group_scoring ) def pending_batch_by_env(self) -> dict[str, int]: @@ -140,7 +142,8 @@ async def add(self, episode: EpisodeResult) -> TrainBatch | None: self.pending_failures.append(episode.failure) for rollout in episode.rollouts: await self.process_rollout(rollout) - self.pending_groups[group_id].extend(episode.rollouts) + if episode.episode is not None: + self.pending_groups[group_id].append(episode.episode) self.pending_group_episodes[group_id] += 1 if self.pending_group_episodes[group_id] < self.group_size_for(env_name): return None @@ -180,16 +183,17 @@ async def process_group(self, group_id: uuid.UUID) -> None: """Finalize one GRPO group: drop errored rollouts (the whole group when ``requires_group_scoring`` and any failed), assign advantages, run pre-batch filters, append survivors to ``pending_batch``.""" - group = self.pending_groups.pop(group_id, []) + episodes = self.pending_groups.pop(group_id, []) self.pending_group_episodes.pop(group_id, None) - if not group: + if not episodes: return + group = [trace for episode in episodes for trace in episode.traces] # Window membership follows group finalization, not arrival: a rollout # only becomes observable (metrics / persistence) once its whole group # is finalized, so a batch's window never claims rollouts of a group # that ships later. Dropped groups still land here — they were observed. - for r in group: - self.pending_rollouts.append(r) + for episode in episodes: + self.pending_rollouts.append(episode) env_name = group[0].env_name task_idx = group[0].task.data.idx survivors = [r for r in group if not r.has_error] diff --git a/src/prime_rl/orchestrator/types.py b/src/prime_rl/orchestrator/types.py index 016f5f7979..54ad33c0af 100644 --- a/src/prime_rl/orchestrator/types.py +++ b/src/prime_rl/orchestrator/types.py @@ -156,9 +156,10 @@ class EpisodeResult: describe the whole episode, plus either the env's traces or the reason there are none. The scheduling facts live here rather than on each trace because that is what they describe — - a group is dispatched, cancelled, and counted whole. ``rollouts`` is empty exactly when - ``failure`` is set: prime-rl never mints a stand-in trace to carry an outcome the env did not - produce, so a cancellation can't be mistaken for an agent's rollout.""" + a group is dispatched, cancelled, and counted whole. ``episode`` is the env's own + ``vf.Episode``, kept whole so its aggregates (``by_agent``, ``num_turns``) are read rather than + rebuilt; it is set exactly when ``failure`` is not, since prime-rl never mints a stand-in trace + to carry an outcome the env did not produce.""" kind: RolloutKind env_name: str @@ -166,14 +167,19 @@ class EpisodeResult: policy_version: int off_policy_steps: int = 0 eval_step: int | None = None - rollouts: list[Rollout] = field(default_factory=list) + episode: vf.WireEpisode | None = None failure: EpisodeFailure | None = None def __post_init__(self) -> None: - assert bool(self.rollouts) != (self.failure is not None), ( - "an episode carries either traces or a failure, never both or neither" + assert (self.episode is not None) != (self.failure is not None), ( + "an episode carries either its traces or a failure, never both or neither" ) + @property + def rollouts(self) -> list[Rollout]: + """The episode's traces; empty when it failed before producing any.""" + return list(self.episode.traces) if self.episode is not None else [] + @dataclass class TrainBatch: diff --git a/tests/unit/orchestrator/test_metrics.py b/tests/unit/orchestrator/test_metrics.py index f5ed8a8d43..ffe3903d35 100644 --- a/tests/unit/orchestrator/test_metrics.py +++ b/tests/unit/orchestrator/test_metrics.py @@ -81,8 +81,19 @@ def mk( ) +def ep(*rollouts): + """One episode over these traces. ``model_construct`` skips validation so the duck-typed + stand-ins above can stand in for real ones.""" + return vf.WireEpisode.model_construct(id=f"e{next(_ids)}", traces=list(rollouts)) + + +def solo(rollouts): + """Each rollout as its own single-trace episode — the single-agent shape.""" + return [ep(r) for r in rollouts] + + def train_wandb(rollouts, subset: str = "all") -> dict: - return TrainRollouts(rollouts).metrics.to_wandb(prefix="train/agg", subset=subset) + return TrainRollouts(solo(rollouts)).metrics.to_wandb(prefix="train/agg", subset=subset) def test_stat(): @@ -95,7 +106,7 @@ def test_stat(): def test_container_effective_by_env_and_listlike(): rc = TrainRollouts( - [mk(env_name="a"), mk(env_name="a", has_error=True), mk(env_name="b", is_filtered=True), mk(env_name="b")] + solo([mk(env_name="a"), mk(env_name="a", has_error=True), mk(env_name="b", is_filtered=True), mk(env_name="b")]) ) assert len(rc) == 4 and [r.env_name for r in rc] == ["a", "a", "b", "b"] # sized + iterable eff = rc.effective @@ -103,16 +114,18 @@ def test_container_effective_by_env_and_listlike(): assert all(not r.has_error and not r.is_filtered and r in rc.rollouts for r in eff) # view of references by_env = rc.by_env() assert set(by_env) == {"a", "b"} and len(by_env["a"]) == 2 and isinstance(by_env["a"], TrainRollouts) - rc.append(mk()) - assert len(rc) == 5 + rc.append(ep(mk())) + assert len(rc) == 5 # sized by rollout, appended by episode def test_to_wandb_distributions(): m = TrainRollouts( - [ - mk(reward=1.0, num_total_tokens=10, num_input_tokens=4), - mk(reward=0.0, num_total_tokens=20, num_input_tokens=6), - ] + solo( + [ + mk(reward=1.0, num_total_tokens=10, num_input_tokens=4), + mk(reward=0.0, num_total_tokens=20, num_input_tokens=6), + ] + ) ).metrics assert m.num_input_tokens.mean() == 5.0 # fluent Stat access out = m.to_wandb(prefix="train/agg", subset="all") @@ -126,15 +139,20 @@ def test_to_wandb_distributions(): def test_episode_and_agent_levels(): # Two proposer-solver episodes: one proposer + two solvers each (the solver fan-out). - rollouts = [ - mk(reward=1.0, num_turns=1, agent_name="proposer", episode_id="e1"), - mk(reward=0.0, num_turns=2, agent_name="solver", episode_id="e1"), - mk(reward=1.0, num_turns=4, agent_name="solver", episode_id="e1"), - mk(reward=0.0, num_turns=3, agent_name="proposer", episode_id="e2"), - mk(reward=1.0, num_turns=6, agent_name="solver", episode_id="e2"), - mk(reward=0.0, num_turns=8, agent_name="solver", episode_id="e2"), - ] - m = TrainRollouts(rollouts).metrics + m = TrainRollouts( + [ + ep( + mk(reward=1.0, num_turns=1, agent_name="proposer"), + mk(reward=0.0, num_turns=2, agent_name="solver"), + mk(reward=1.0, num_turns=4, agent_name="solver"), + ), + ep( + mk(reward=0.0, num_turns=3, agent_name="proposer"), + mk(reward=1.0, num_turns=6, agent_name="solver"), + mk(reward=0.0, num_turns=8, agent_name="solver"), + ), + ] + ).metrics assert m.num_turns.mean() == 12.0 # episode-level sums: 1+2+4 and 3+6+8 assert m.num_total_tokens.values == [30.0, 30.0] # summed across the episode's traces out = m.to_wandb(prefix="train/agg", subset="all") @@ -152,11 +170,11 @@ def test_agent_metrics_are_flat_over_traces(): """Inside a seat the trace is the unit of aggregation, so an uneven fan-out (one solver trace from this episode, three from that) never reweights anything: every agent-level metric is the plain figure over that agent's rollouts.""" - rollouts = [ - mk(agent_name="solver", episode_id="e1", is_truncated=True, reward=1.0), - *[mk(agent_name="solver", episode_id="e2", is_truncated=False, reward=0.0) for _ in range(3)], + episodes = [ + ep(mk(agent_name="solver", is_truncated=True, reward=1.0)), + ep(*[mk(agent_name="solver", is_truncated=False, reward=0.0) for _ in range(3)]), ] - out = TrainRollouts(rollouts).metrics.to_wandb(prefix="train/agg", subset="all") + out = TrainRollouts(episodes).metrics.to_wandb(prefix="train/agg", subset="all") assert out["train/agg/all/solver/is_truncated/mean"] == 0.25 # 1 of 4 traces, not (1.0 + 0.0) / 2 assert out["train/agg/all/solver/is_completed/mean"] == 1.0 assert out["train/agg/all/solver/is_trainable/mean"] == 1.0 # sibling rates agree @@ -164,7 +182,9 @@ def test_agent_metrics_are_flat_over_traces(): def test_boolean_rates_and_error_breakdown_all_only(): - rc = TrainRollouts([mk(is_truncated=True), mk(has_error=True, error_type="ProviderError"), mk(is_filtered=True)]) + rc = TrainRollouts( + solo([mk(is_truncated=True), mk(has_error=True, error_type="ProviderError"), mk(is_filtered=True)]) + ) out = rc.metrics.to_wandb(prefix="train/agg", subset="all") assert out["train/agg/all/agent/is_truncated/mean"] == 1 / 3 assert out["train/agg/all/agent/is_completed/mean"] == 1.0 @@ -202,7 +222,7 @@ def test_nested_metrics_and_rewards(): # scoring failed after seeding: unscored (None) entries count as 0.0 on `all` mk(has_error=True, metrics={"acc": None}, rewards={"correct": None, "format": None}), ] - rc = TrainRollouts(rollouts) + rc = TrainRollouts(solo(rollouts)) m = rc.metrics agent = m.by_agent()["agent"] assert agent.metrics["acc"].mean() == pytest.approx(4 / 3) # nested group access @@ -217,13 +237,15 @@ def test_nested_metrics_and_rewards(): assert eff["train/agg/effective/agent/rewards/format/mean"] == 0.5 # cross-env agg: another env's unscored trace carries different keys, so it can't dilute these other = mk(env_name="other", has_error=True, rewards={"solved": None}) - agg = TrainRollouts(rollouts + [other]).metrics.to_wandb(prefix="train/agg", subset="all") + agg = TrainRollouts(solo(rollouts + [other])).metrics.to_wandb(prefix="train/agg", subset="all") assert agg["train/agg/all/agent/rewards/format/mean"] == pytest.approx(1 / 3) assert agg["train/agg/all/agent/rewards/solved/mean"] == 0.0 def test_nested_timing(): - m = TrainRollouts([mk(setup=1.0, agent=2.0, agent_model=1.5, agent_harness=0.5, finalize=0.5, scoring=0.5)]).metrics + m = TrainRollouts( + solo([mk(setup=1.0, agent=2.0, agent_model=1.5, agent_harness=0.5, finalize=0.5, scoring=0.5)]) + ).metrics timing = m.by_agent()["agent"].timing assert timing.setup.mean() == 1.0 and timing.total.mean() == 4.0 # total sums all four phases assert timing.agent_model.mean() == 1.5 and timing.agent_harness.mean() == 0.5 @@ -244,12 +266,12 @@ def test_train_only_metrics_absent_from_eval(): assert out["train/agg/all/agent/is_filtered/mean"] == 0.5 assert out["train/agg/all/agent/filters/gibberish/mean"] == 0.5 assert "train/agg/all/is_trainable/mean" not in out # pipeline verdicts are per-trace - eval_out = EvalRollouts(rollouts).metrics.to_wandb(prefix="eval/x", subset="all") + eval_out = EvalRollouts(solo(rollouts)).metrics.to_wandb(prefix="eval/x", subset="all") assert not any("is_trainable" in k or "is_filtered" in k or "/filters/" in k for k in eval_out) def test_eval_avg_at_k_and_pass_k(): - binary = EvalRollouts([mk(reward=1.0, group_id="g0"), mk(reward=0.0, group_id="g0")]) + binary = EvalRollouts(solo([mk(reward=1.0, group_id="g0"), mk(reward=0.0, group_id="g0")])) eff = binary.effective.metrics.to_wandb(prefix="eval/x", subset="effective") assert eff["eval/x/effective/agent/avg@2"] == 0.5 # mean reward under avg@ (k from the groups) assert "eval/x/effective/avg@2" not in eff # scores are per-agent, never pooled @@ -257,7 +279,7 @@ def test_eval_avg_at_k_and_pass_k(): all_out = binary.metrics.to_wandb(prefix="eval/x", subset="all") assert all_out["eval/x/all/agent/avg@2"] == 0.5 assert not any("pass@" in k or "pass^" in k for k in all_out) # pass@k effective-only - non_binary = EvalRollouts([mk(reward=0.5, group_id="g0"), mk(reward=1.0, group_id="g0")]) + non_binary = EvalRollouts(solo([mk(reward=0.5, group_id="g0"), mk(reward=1.0, group_id="g0")])) assert not any("pass@" in k for k in non_binary.effective.metrics.to_wandb(prefix="eval/x", subset="effective")) @@ -288,8 +310,8 @@ def test_episode_failure_metrics(): def test_episode_result_carries_traces_xor_failure(): """The envelope makes the phantom trace unrepresentable: an outcome is traces or a failure.""" - rollout = mk() - assert EpisodeResult(kind="train", env_name="env", group_id=uuid4(), policy_version=0, rollouts=[rollout]) + episode = ep(mk()) + assert EpisodeResult(kind="train", env_name="env", group_id=uuid4(), policy_version=0, episode=episode) assert EpisodeResult( kind="train", env_name="env", @@ -305,6 +327,6 @@ def test_episode_result_carries_traces_xor_failure(): env_name="env", group_id=uuid4(), policy_version=0, - rollouts=[rollout], + episode=episode, failure=EpisodeFailure(type="Cancelled", message="Off-policy cancel"), ) From ca0adf601c48beee2d5a88d031933b4c1dbc20b1 Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Mon, 3 Aug 2026 23:12:36 +0000 Subject: [PATCH 24/58] refactor(orchestrator)!: prl Episode extends vf.Episode; failures are traceless episodes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit EmptyEpisode was never a failure mode of its own — vf's run_episode already raises when run() mints no trace, records it on episode.errors and returns the episode — and the code that reported it threw that real error away for a generic string. So there is one shape for every failure that has no trace to ride: an episode with no traces and the reason on errors. EpisodeFailure and EpisodeResult both go; prl's Episode subclasses vf's and adds only the six scheduling facts prime-rl actually contributes. Co-Authored-By: Claude Fable 5 --- src/prime_rl/orchestrator/dispatcher.py | 73 +++++++++------------ src/prime_rl/orchestrator/envs.py | 13 ++-- src/prime_rl/orchestrator/eval_sink.py | 26 +++----- src/prime_rl/orchestrator/metrics.py | 41 ++++++------ src/prime_rl/orchestrator/orchestrator.py | 18 +++--- src/prime_rl/orchestrator/train_sink.py | 25 ++----- src/prime_rl/orchestrator/types.py | 65 +++++++------------ tests/unit/orchestrator/test_metrics.py | 79 +++++++++++------------ 8 files changed, 146 insertions(+), 194 deletions(-) diff --git a/src/prime_rl/orchestrator/dispatcher.py b/src/prime_rl/orchestrator/dispatcher.py index f38a805554..e50475c34c 100644 --- a/src/prime_rl/orchestrator/dispatcher.py +++ b/src/prime_rl/orchestrator/dispatcher.py @@ -6,10 +6,10 @@ single-agent episode). - Optional rate limiting via ``AsyncLimiter(tasks_per_minute, 60)``. - Emit-everything invariant: every dispatched episode eventually reaches ``out_q`` exactly once, - as an ``EpisodeResult``. An env-side failure (env error, empty trajectory) rides the trace that - hit it (``trace.last_error``); a prime-rl-side one (task exception, off-policy cancel) has no - trace to ride, so it arrives as ``EpisodeResult.failure`` rather than as a stand-in rollout. - Sinks decide drop / partial-train policy. + as an ``Episode``. An env-side failure (env error, empty trajectory) rides the trace that hit it + (``trace.last_error``); a failure with no trace to ride — the env produced none, the task raised, + the group was cancelled — is an episode with no traces and the reason on ``episode.errors``, + never a stand-in rollout. Sinks decide drop / partial-train policy. - ``DispatcherMode.PREFER_TRAIN`` / ``PREFER_EVAL`` controls which kind to schedule next. Transitions are level-triggered (driven by the eval source's emptiness), so in-flight rollouts of the opposite kind drain @@ -41,8 +41,7 @@ from prime_rl.orchestrator.eval_source import EvalSource from prime_rl.orchestrator.train_source import TrainSource from prime_rl.orchestrator.types import ( - EpisodeFailure, - EpisodeResult, + Episode, GroupState, InflightRollout, Policy, @@ -158,7 +157,7 @@ def __init__( # Bounded so the dispatcher backpressures on a slow sink. One entry per # episode — the sinks count episodes, never loose traces. - self.out_q: asyncio.Queue[EpisodeResult] = asyncio.Queue(maxsize=max(8, self.max_inflight)) + self.out_q: asyncio.Queue[Episode] = asyncio.Queue(maxsize=max(8, self.max_inflight)) self.mode: DispatcherMode = DispatcherMode.PREFER_TRAIN # Set by the orchestrator after the final train step; pipeline then @@ -513,18 +512,12 @@ async def handle_completed_rollout(self, task: asyncio.Task) -> None: try: result = task.result() - episodes: list[vf.WireEpisode] = result if isinstance(result, list) else [result] + episodes: list[Episode] = result if isinstance(result, list) else [result] except asyncio.CancelledError: return except Exception as exc: get_logger().warning(f"Rollout task failed in group {meta.group_id} ({meta.env_name}): {exc!r}") - await self.emit_failed_episodes(meta, group, EpisodeFailure(type="TaskFailed", message=repr(exc))) - return - if not any(episode.traces for episode in episodes): - get_logger().warning(f"Env returned no traces in group {meta.group_id} ({meta.env_name})") - await self.emit_failed_episodes( - meta, group, EpisodeFailure(type="EmptyEpisode", message="env run returned an episode with no traces") - ) + await self.emit_failed_episodes(meta, group, vf.Error(type="TaskFailed", message=repr(exc))) return for r in (r for episode in episodes for r in episode.traces): @@ -548,13 +541,11 @@ async def emit_episode( self, meta: InflightRollout, group: GroupState | None, - episode: vf.WireEpisode | None, - failure: EpisodeFailure | None = None, + episode: Episode, ) -> None: - """Put one completed episode on ``out_q`` — its traces, or the ``failure`` that stands in - for them. Pops the group from ``self.groups`` once every owed episode has been emitted. - The scheduling facts also ride each trace, where the sinks and the saved records read - them.""" + """Put one completed episode on ``out_q``, stamped with the facts of the dispatch it came + from. Pops the group from ``self.groups`` once every owed episode has been emitted. The + same facts also ride each trace, where the algorithms and the saved records read them.""" eval_step = meta.eval_step policy_version = meta.policy_version if group is not None: @@ -566,7 +557,7 @@ async def emit_episode( if meta.kind == "eval": assert eval_step is not None, "eval episode missing eval_step" - for rollout in episode.traces if episode is not None else []: + for rollout in episode.rollouts: rollout.kind = meta.kind rollout.env_name = meta.env_name rollout.group_id = meta.group_id @@ -574,27 +565,21 @@ async def emit_episode( rollout.off_policy_steps = meta.off_policy_steps if meta.kind == "eval": rollout.eval_step = eval_step - await self.out_q.put( - EpisodeResult( - kind=meta.kind, - env_name=meta.env_name, - group_id=meta.group_id, - policy_version=policy_version, - off_policy_steps=meta.off_policy_steps, - eval_step=eval_step if meta.kind == "eval" else None, - episode=episode, - failure=failure, - ) - ) - - async def emit_failed_episodes( - self, meta: InflightRollout, group: GroupState | None, failure: EpisodeFailure - ) -> None: - """Emit one failed episode per rollout the task owed, so the sink still counts its way to - ``group_size`` — prime-rl's own failure, carried as itself rather than as a stand-in trace.""" + episode.kind = meta.kind + episode.env_name = meta.env_name + episode.group_id = meta.group_id + episode.policy_version = policy_version + episode.off_policy_steps = meta.off_policy_steps + episode.eval_step = eval_step if meta.kind == "eval" else None + await self.out_q.put(episode) + + async def emit_failed_episodes(self, meta: InflightRollout, group: GroupState | None, error: vf.Error) -> None: + """Emit one traceless episode per rollout the task owed, so the sink still counts its way to + ``group_size``. This is the shape vf already gives a run that produced nothing — the reason + on ``errors`` — so prime-rl's own failures need no type of their own.""" for _ in range(meta.rollout_count): self.metrics.record_error(kind=meta.kind, env_name=meta.env_name) - await self.emit_episode(meta, group, None, failure=failure) + await self.emit_episode(meta, group, Episode.model_construct(errors=[error])) async def drop_group(self, group_id: uuid.UUID) -> int: """Cancel remaining in-flight tasks for this group and emit a @@ -619,10 +604,10 @@ async def drop_group(self, group_id: uuid.UUID) -> int: tasks_to_cancel = [task for task, _ in claimed] inflight_cancelled = sum(meta.rollout_count for _, meta in claimed) last_meta: InflightRollout | None = claimed[-1][1] if claimed else None - cancel = EpisodeFailure(type="Cancelled", message="Off-policy cancel") + cancel = vf.Error(type="Cancelled", message="Off-policy cancel") for _, meta in claimed: for _ in range(meta.rollout_count): - await self.emit_episode(meta, group, None, failure=cancel) + await self.emit_episode(meta, group, Episode.model_construct(errors=[cancel])) # For non-group-scoring envs, the group may have rollouts that # were never dispatched (``rollouts_to_schedule > 0``). Emit @@ -643,7 +628,7 @@ async def drop_group(self, group_id: uuid.UUID) -> int: ) unscheduled_cancelled = group.rollouts_to_schedule for _ in range(unscheduled_cancelled): - await self.emit_episode(fallback_meta, group, None, failure=cancel) + await self.emit_episode(fallback_meta, group, Episode.model_construct(errors=[cancel])) cancelled = inflight_cancelled + unscheduled_cancelled if cancelled > 0: diff --git a/src/prime_rl/orchestrator/envs.py b/src/prime_rl/orchestrator/envs.py index c27fa1d42c..01265492ff 100644 --- a/src/prime_rl/orchestrator/envs.py +++ b/src/prime_rl/orchestrator/envs.py @@ -37,15 +37,16 @@ from prime_rl.configs.orchestrator import EnvConfig, EvalSourceConfig, TrainSourceConfig from prime_rl.orchestrator.algo import Algorithm, build_algorithm from prime_rl.orchestrator.sampler import Sampler -from prime_rl.orchestrator.types import Rollout +from prime_rl.orchestrator.types import Episode, Rollout from prime_rl.utils.logger import get_logger # Every wire trace validates into this type. WireTaskData (extra="allow") keeps the env's task # fields without importing the env package — the orchestrator never reads them typed (only # task.idx + task.model_dump). ROLLOUT_TYPE = Rollout[vf.WireTaskData] -# The env server answers a ``WireEpisode``; we keep that envelope and only re-type its traces. -EPISODE_TYPE = vf.WireEpisode +# The env server answers a wire episode; we keep that envelope, re-type its traces, and let the +# dispatcher stamp the scheduling facts onto it. +EPISODE_TYPE = Episode # Max wait for a spawned env server to bind and report its address. A legacy # child loads its dataset before reporting, so this is generous. @@ -208,8 +209,8 @@ async def run( ) -> EPISODE_TYPE: """Run one episode; return it with its traces re-typed as ``Rollout``. A v1 env takes the task itself (``task_data``); the legacy bridge is addressed by dataset row (``task_idx``). - A zero-trace episode comes back as-is for the dispatcher to report as a failure; a - not-``ok`` episode marks its clean traces failed so partial episodes never train. + A zero-trace episode comes back as-is — vf already recorded on it why it produced nothing; + a not-``ok`` episode marks its clean traces failed so partial episodes never train. The episode itself is kept rather than flattened away: it is what groups an env's seats, and downstream reads its aggregates (``by_agent``, ``num_turns``) instead of rebuilding @@ -230,7 +231,7 @@ async def run( ) rollout.errors = [*rollout.errors, error] rollout.ok = False - return episode.model_copy(update={"traces": rollouts}) + return EPISODE_TYPE.model_construct(**{**dict(episode), "traces": rollouts}) async def run_group( self, client: vf.ClientConfig, task_idx: int, model_name: str, group_size: int, cache_salt: str | None diff --git a/src/prime_rl/orchestrator/eval_sink.py b/src/prime_rl/orchestrator/eval_sink.py index 8a6eaae10c..ba8d014640 100644 --- a/src/prime_rl/orchestrator/eval_sink.py +++ b/src/prime_rl/orchestrator/eval_sink.py @@ -8,7 +8,7 @@ 3. ``process_batch`` — at ``num_examples × group_size`` episodes, return an ``EvalBatch`` with the full returned cohort (metrics are computed downstream). -``add()`` takes one ``EpisodeResult`` and returns ``EvalBatch | None``; +``add()`` takes one ``Episode`` and returns ``EvalBatch | None``; all accounting counts episodes, never loose traces. """ @@ -17,11 +17,9 @@ import uuid from collections import defaultdict -import verifiers.v1 as vf - from prime_rl.orchestrator.envs import EvalEnvs from prime_rl.orchestrator.metrics import EvalRollouts -from prime_rl.orchestrator.types import EpisodeFailure, EpisodeResult, EvalBatch, Rollout +from prime_rl.orchestrator.types import Episode, EvalBatch, Rollout from prime_rl.utils.logger import get_logger @@ -30,15 +28,13 @@ class EvalSink: def __init__(self, *, eval_envs: EvalEnvs) -> None: self.eval_envs = eval_envs - self.pending_groups: dict[uuid.UUID, list[vf.WireEpisode]] = defaultdict(list) + self.pending_groups: dict[uuid.UUID, list[Episode]] = defaultdict(list) # Episodes arrived per group / per batch bucket — the finalization counts. self.pending_group_episodes: dict[uuid.UUID, int] = defaultdict(int) - self.pending_batches: dict[tuple[str, int], list[vf.WireEpisode]] = defaultdict(list) + self.pending_batches: dict[tuple[str, int], list[Episode]] = defaultdict(list) self.pending_batch_episodes: dict[tuple[str, int], int] = defaultdict(int) - # Episodes of the epoch that produced no traces at all (cancelled, or the task failed). - self.pending_batch_failures: dict[tuple[str, int], list[EpisodeFailure]] = defaultdict(list) - def add(self, episode: EpisodeResult) -> EvalBatch | None: + def add(self, episode: Episode) -> EvalBatch | None: """Process one episode arrival; finalize the group on the ``group_size``-th episode and the per-env epoch on the ``num_examples × group_size``-th. A failed episode brings no rollouts but still counts toward both.""" @@ -47,10 +43,7 @@ def add(self, episode: EpisodeResult) -> EvalBatch | None: for rollout in episode.rollouts: self.process_rollout(rollout) bkey = (env_name, episode.eval_step) - if episode.failure is not None: - self.pending_batch_failures[bkey].append(episode.failure) - if episode.episode is not None: - self.pending_groups[group_id].append(episode.episode) + self.pending_groups[group_id].append(episode) self.pending_group_episodes[group_id] += 1 if self.pending_group_episodes[group_id] >= self.group_size_for(env_name): self.process_group(group_id) @@ -77,10 +70,10 @@ def batch_progress(self) -> list[tuple[str, int, int, int, int]]: for group_id, episodes in self.pending_groups.items(): if not episodes: continue - env_name = episodes[0].traces[0].env_name + env_name = episodes[0].env_name if self.eval_envs.get(env_name).requires_group_scoring: continue - bkey = (env_name, episodes[0].traces[0].eval_step) + bkey = (env_name, episodes[0].eval_step) buffered[bkey] = buffered.get(bkey, 0) + self.pending_group_episodes.get(group_id, 0) return [ ( @@ -134,5 +127,4 @@ def process_batch(self, key: tuple[str, int]) -> EvalBatch: env_name, step = key episodes = self.pending_batches.pop(key, []) self.pending_batch_episodes.pop(key, None) - failures = self.pending_batch_failures.pop(key, []) - return EvalBatch(env_name=env_name, step=step, rollouts=EvalRollouts(episodes), failures=failures) + return EvalBatch(env_name=env_name, step=step, rollouts=EvalRollouts(episodes)) diff --git a/src/prime_rl/orchestrator/metrics.py b/src/prime_rl/orchestrator/metrics.py index a92b199fa7..11baa6c907 100644 --- a/src/prime_rl/orchestrator/metrics.py +++ b/src/prime_rl/orchestrator/metrics.py @@ -23,28 +23,29 @@ from typing import TYPE_CHECKING, Any, Callable, Iterator, Literal -import verifiers.v1 as vf - from prime_rl.orchestrator.utils import compute_pass_metrics if TYPE_CHECKING: - from prime_rl.orchestrator.types import EpisodeFailure, Rollout + from prime_rl.orchestrator.types import Episode, Rollout Subset = Literal["all", "effective"] -def episode_failure_metrics(failures: list[EpisodeFailure], *, prefix: str, total: int) -> dict[str, float]: +def episode_failure_metrics(episodes: list[Episode], *, prefix: str) -> dict[str, float]: """Episodes that produced no traces, by reason (``{prefix}/episode_failure/``) plus their - share of the window. These are prime-rl's own outcomes — a cancellation or a task that never - reached the env — so they belong to no agent and are counted whole rather than folded into a - seat's error rate. ``total`` is the window's rollout count, for the rate's denominator.""" + share of the window. A cancellation, a task that never reached the env, an env that ran no + agent — all of them are an episode with nothing on it but ``errors``, so one counter covers + every cause. They belong to no agent, so they are counted whole rather than folded into a + seat's error rate.""" + if not episodes: + return {} + failed = [episode for episode in episodes if episode.failed] counts: dict[str, int] = {} - for failure in failures: - counts[failure.type] = counts.get(failure.type, 0) + 1 + for episode in failed: + error = episode.last_error + counts[error.type if error else "Unknown"] = counts.get(error.type if error else "Unknown", 0) + 1 out = {f"{prefix}/episode_failure/{name}": float(n) for name, n in sorted(counts.items())} - denominator = total + len(failures) - if denominator: - out[f"{prefix}/episode_failure/rate"] = len(failures) / denominator + out[f"{prefix}/episode_failure/rate"] = len(failed) / len(episodes) return out @@ -283,7 +284,7 @@ class EpisodeMetrics: ``.mean()`` is the rate) serve the console log lines. ``TrainMetrics`` / ``EvalMetrics`` extend ``to_wandb`` with the train pipeline rates and the eval scores.""" - def __init__(self, episodes: list[vf.WireEpisode]) -> None: + def __init__(self, episodes: list[Episode]) -> None: self.episodes = episodes self.rollouts: list[Rollout] = [trace for episode in episodes for trace in episode.traces] @@ -387,7 +388,7 @@ class EvalMetrics(EpisodeMetrics): per example) is supplied by the container so the ``all`` and ``effective`` subsets — and every agent — share one stable key.""" - def __init__(self, episodes: list[vf.WireEpisode], group_size: int) -> None: + def __init__(self, episodes: list[Episode], group_size: int) -> None: super().__init__(episodes) self.group_size = group_size @@ -405,7 +406,7 @@ def to_wandb(self, *, prefix: str, subset: Subset) -> dict[str, float]: return out -def keep_traces(episode: vf.WireEpisode, keep: Callable[[Rollout], bool]) -> vf.WireEpisode | None: +def keep_traces(episode: Episode, keep: Callable[[Rollout], bool]) -> Episode | None: """The episode narrowed to the traces that pass ``keep``, or ``None`` if none do. A subset view stays a list of episodes so the episode-level aggregates keep describing what survived.""" traces = [trace for trace in episode.traces if keep(trace)] @@ -418,10 +419,10 @@ class TrainRollouts: narrowed to its surviving traces; ``metrics`` builds ``TrainMetrics`` over them. Sized and iterated by *rollout*, since that is the unit almost every consumer wants.""" - def __init__(self, episodes: list[vf.WireEpisode] | None = None) -> None: + def __init__(self, episodes: list[Episode] | None = None) -> None: self.episodes = episodes if episodes is not None else [] - def append(self, episode: vf.WireEpisode) -> None: + def append(self, episode: Episode) -> None: self.episodes.append(episode) @property @@ -443,9 +444,9 @@ def effective(self) -> TrainRollouts: return TrainRollouts([episode for episode in kept if episode is not None]) def by_env(self) -> dict[str, TrainRollouts]: - grouped: dict[str, list[vf.WireEpisode]] = {} + grouped: dict[str, list[Episode]] = {} for episode in self.episodes: - grouped.setdefault(episode.traces[0].env_name, []).append(episode) + grouped.setdefault(episode.env_name, []).append(episode) return {env: TrainRollouts(episodes) for env, episodes in grouped.items()} @property @@ -459,7 +460,7 @@ class EvalRollouts: ``group_size`` (rollouts per example, the ``avg@k`` k) is derived from the full epoch and carried onto ``effective`` so both subsets share one stable key; ``metrics`` builds ``EvalMetrics``.""" - def __init__(self, episodes: list[vf.WireEpisode] | None = None, group_size: int | None = None) -> None: + def __init__(self, episodes: list[Episode] | None = None, group_size: int | None = None) -> None: self.episodes = episodes if episodes is not None else [] self._group_size = group_size diff --git a/src/prime_rl/orchestrator/orchestrator.py b/src/prime_rl/orchestrator/orchestrator.py index 19314f2838..6a06250b7a 100644 --- a/src/prime_rl/orchestrator/orchestrator.py +++ b/src/prime_rl/orchestrator/orchestrator.py @@ -3,8 +3,8 @@ ``Orchestrator`` owns the shared state (policy, progress, ckpt, monitor) and drives the pipeline. Components are single-purpose: -- ``RolloutDispatcher`` schedules rollouts; emits ``EpisodeResult`` (train/eval - discriminated by ``kind``) on its queue — its traces, or the failure that replaced them. +- ``RolloutDispatcher`` schedules rollouts; emits ``Episode`` (train/eval discriminated by + ``kind``) on its queue — the env's own episode, or a traceless one saying why there is none. - ``TrainSink`` ingests train rollouts (tokenize → advantages → filters) and returns a ``TrainBatch`` when the threshold is met. - ``EvalSink`` ingests eval rollouts and returns an ``EvalBatch`` (the full @@ -58,7 +58,7 @@ from prime_rl.orchestrator.train_sink import TrainSink from prime_rl.orchestrator.train_source import TrainSource from prime_rl.orchestrator.types import ( - EpisodeResult, + Episode, EvalBatch, Policy, Progress, @@ -511,7 +511,7 @@ async def start(self) -> None: trim_process_memory() async def main_loop(self) -> None: - """Consume ``EpisodeResult``\\ s from the dispatcher and route them + """Consume ``Episode``\\ s from the dispatcher and route them to the train / eval sink. Both sinks return a finalized batch (or ``None``) from ``add()``; we just dispatch on the result.""" while not self.stopped.is_set(): @@ -521,7 +521,7 @@ async def main_loop(self) -> None: break try: - episode: EpisodeResult = await asyncio.wait_for(self.dispatcher.out_q.get(), timeout=0.5) + episode: Episode = await asyncio.wait_for(self.dispatcher.out_q.get(), timeout=0.5) except asyncio.TimeoutError: continue @@ -529,7 +529,7 @@ async def main_loop(self) -> None: # ``all`` trace file the moment it arrives, so it survives crashes and drains. # Train rollouts belong to the batch window currently collecting (``progress.step``), # eval rollouts to the step whose eval triggered them. A failed episode has no trace - # to write; it is counted in the batch's failure tally instead. + # to write; it is counted by the window's episode-failure metrics instead. kind = episode.kind step = episode.eval_step if kind == "eval" else self.progress.step assert step is not None @@ -546,7 +546,7 @@ async def main_loop(self) -> None: episode_id=rollout.episode_id, policy_version=rollout.policy_version, ) - if episode.rollouts: + if not episode.failed: await asyncio.to_thread( save_rollouts, [rollout.to_record() for rollout in episode.rollouts], @@ -665,7 +665,7 @@ async def finalize_train_batch(self, batch: TrainBatch) -> None: metrics |= env_pool.metrics.to_wandb(prefix=f"train/{env_name}", subset=subset) # Episodes that produced no trace at all are an episode-level fact — they belong to no # agent, so they are counted here rather than folded into any seat's error rate. - metrics |= episode_failure_metrics(batch.failures, prefix="train/agg", total=len(batch.rollouts)) + metrics |= episode_failure_metrics(batch.rollouts.episodes, prefix="train/agg") # Progress / timing / env-share / pre-filter accounting (assembled here, not in the metrics # objects). ``num_tokens`` is over the full arrival window; the input/output breakdown is over @@ -882,7 +882,7 @@ async def finalize_eval_batch(self, batch: EvalBatch) -> None: metrics: dict[str, float] = {} for subset, pool in (("all", rollouts), ("effective", effective)): metrics |= pool.metrics.to_wandb(prefix=f"eval/{batch.env_name}", subset=subset) - metrics |= episode_failure_metrics(batch.failures, prefix=f"eval/{batch.env_name}", total=len(rollouts)) + metrics |= episode_failure_metrics(rollouts.episodes, prefix=f"eval/{batch.env_name}") metrics[f"eval/{batch.env_name}/policy_version"] = float(policy_version) metrics["step"] = float(batch.step) self.monitor.log(metrics, step=batch.step) diff --git a/src/prime_rl/orchestrator/train_sink.py b/src/prime_rl/orchestrator/train_sink.py index 35d0c56f77..a154cf4f00 100644 --- a/src/prime_rl/orchestrator/train_sink.py +++ b/src/prime_rl/orchestrator/train_sink.py @@ -10,7 +10,7 @@ 3. ``process_batch`` — applies post-batch filter annotations and assembles the trainer-bound ``TrainingSample`` list. Returns a ``TrainBatch``. -``add()`` takes one ``EpisodeResult`` and returns +``add()`` takes one ``Episode`` and returns ``TrainBatch | None``; group accounting counts episodes, never loose traces. I/O concerns (ship to trainer, save_rollouts, monitor.log) live on the orchestrator. @@ -22,14 +22,12 @@ import uuid from collections import defaultdict -import verifiers.v1 as vf - from prime_rl.configs.orchestrator import OrchestratorConfig from prime_rl.orchestrator.envs import TrainEnvs from prime_rl.orchestrator.filters import RolloutFilter, apply_filters from prime_rl.orchestrator.metrics import TrainRollouts from prime_rl.orchestrator.trajectories import trace_to_samples -from prime_rl.orchestrator.types import EpisodeFailure, EpisodeResult, Rollout, TrainBatch +from prime_rl.orchestrator.types import Episode, Rollout, TrainBatch from prime_rl.transport import TrainingSample from prime_rl.utils.logger import get_logger @@ -78,14 +76,10 @@ def __init__( # finalized since the last ship (errored + filtered + survivors). # In-progress groups stay out until they finalize. self.pending_rollouts: TrainRollouts = TrainRollouts() - # Episodes in that window that produced no traces at all (cancelled, or the task - # itself failed). They have no rollout to be counted through, so they are tallied - # here and reported alongside the window's metrics. - self.pending_failures: list[EpisodeFailure] = [] # Keyed by the dispatcher's group UUID. ``(env_name, task_idx)`` # isn't unique — the same task can be re-sampled while an # earlier group is still in flight - self.pending_groups: dict[uuid.UUID, list[vf.WireEpisode]] = defaultdict(list) + self.pending_groups: dict[uuid.UUID, list[Episode]] = defaultdict(list) # Episodes arrived per group — the finalization count (an episode may # add several traces to ``pending_groups`` but counts once here). self.pending_group_episodes: dict[uuid.UUID, int] = defaultdict(int) @@ -119,7 +113,7 @@ def buffered_count(self) -> int: return sum( self.pending_group_episodes.get(group_id, 0) for group_id, episodes in self.pending_groups.items() - if episodes and not self.train_envs.get(episodes[0].traces[0].env_name).requires_group_scoring + if episodes and not self.train_envs.get(episodes[0].env_name).requires_group_scoring ) def pending_batch_by_env(self) -> dict[str, int]: @@ -130,7 +124,7 @@ def pending_batch_by_env(self) -> dict[str, int]: counts[r.env_name] += 1 return dict(counts) - async def add(self, episode: EpisodeResult) -> TrainBatch | None: + async def add(self, episode: Episode) -> TrainBatch | None: """Process one episode arrival; finalize the group on the ``group_size``-th episode; return a ``TrainBatch`` if the finalization pushed (or left) the batch over its threshold. Arrivals into @@ -138,12 +132,9 @@ async def add(self, episode: EpisodeResult) -> TrainBatch | None: rollouts, but still counts toward the group so finalization triggers.""" group_id = episode.group_id env_name = episode.env_name - if episode.failure is not None: - self.pending_failures.append(episode.failure) for rollout in episode.rollouts: await self.process_rollout(rollout) - if episode.episode is not None: - self.pending_groups[group_id].append(episode.episode) + self.pending_groups[group_id].append(episode) self.pending_group_episodes[group_id] += 1 if self.pending_group_episodes[group_id] < self.group_size_for(env_name): return None @@ -297,11 +288,9 @@ def process_batch(self) -> TrainBatch: # samples) — an empty batch is dropped unlogged by the orchestrator, so keep accumulating its # finalized groups (and any overflow) into the next shipped batch's window. rollouts = self.pending_rollouts - failures = self.pending_failures if samples: self.pending_rollouts = TrainRollouts() - self.pending_failures = [] - return TrainBatch(rollouts=rollouts, samples=samples, failures=failures) + return TrainBatch(rollouts=rollouts, samples=samples) def reset_pre_filter_stats(self) -> None: self.pre_filter_seen = 0 diff --git a/src/prime_rl/orchestrator/types.py b/src/prime_rl/orchestrator/types.py index 54ad33c0af..cc134a4dae 100644 --- a/src/prime_rl/orchestrator/types.py +++ b/src/prime_rl/orchestrator/types.py @@ -3,8 +3,8 @@ from __future__ import annotations import uuid -from dataclasses import dataclass, field -from typing import TYPE_CHECKING, Generic, Literal, Protocol +from dataclasses import dataclass +from typing import TYPE_CHECKING, Generic, Literal, Protocol, cast import verifiers.v1 as vf from pydantic import ConfigDict, Field @@ -140,45 +140,35 @@ def is_trainable(self) -> bool: return bool(self.advantages) and any(a != 0.0 for a in self.advantages) -@dataclass -class EpisodeFailure: - """Why a dispatched episode carries no traces of its own. This is prime-rl's own verdict — - the episode never reached the env, or was pulled back before it finished — as opposed to an - env-side error, which rides on the trace that hit it (``Trace.last_error``).""" - - type: Literal["Cancelled", "TaskFailed", "EmptyEpisode"] - message: str +class Episode(vf.WireEpisode): + """The env's own ``vf.Episode`` extended with prime-rl's scheduling facts — the only thing + prime-rl genuinely adds, so the episode itself travels rather than a wrapper around it. Those + fields are ``exclude=True``, so dumping an Episode yields a plain wire episode. + An episode that produced no traces is not a special case and needs no stand-in rollout: vf + already records why on ``errors`` (its ``run_episode`` puts the exception there and returns the + episode with ``ok`` false), and prime-rl's own outcomes — an off-policy cancel, a task that + raised before reaching the env — are minted the same way. So ``failed`` is simply "no traces", + and ``last_error`` says why in one vocabulary for every cause.""" -@dataclass -class EpisodeResult: - """One dispatched episode as it arrives from the dispatcher: the scheduling facts that - describe the whole episode, plus either the env's traces or the reason there are none. + model_config = ConfigDict(arbitrary_types_allowed=True) # traces are ``Rollout``s - The scheduling facts live here rather than on each trace because that is what they describe — - a group is dispatched, cancelled, and counted whole. ``episode`` is the env's own - ``vf.Episode``, kept whole so its aggregates (``by_agent``, ``num_turns``) are read rather than - rebuilt; it is set exactly when ``failure`` is not, since prime-rl never mints a stand-in trace - to carry an outcome the env did not produce.""" - - kind: RolloutKind - env_name: str - group_id: uuid.UUID - policy_version: int - off_policy_steps: int = 0 - eval_step: int | None = None - episode: vf.WireEpisode | None = None - failure: EpisodeFailure | None = None - - def __post_init__(self) -> None: - assert (self.episode is not None) != (self.failure is not None), ( - "an episode carries either its traces or a failure, never both or neither" - ) + kind: RolloutKind = Field(default="train", exclude=True) + env_name: str = Field(default="", exclude=True) + group_id: uuid.UUID = Field(default_factory=uuid.uuid4, exclude=True) + policy_version: int = Field(default=0, exclude=True) + off_policy_steps: int = Field(default=0, exclude=True) + eval_step: int | None = Field(default=None, exclude=True) @property def rollouts(self) -> list[Rollout]: - """The episode's traces; empty when it failed before producing any.""" - return list(self.episode.traces) if self.episode is not None else [] + """The episode's traces, typed as the rollouts prime-rl works with.""" + return cast(list[Rollout], self.traces) + + @property + def failed(self) -> bool: + """Whether the episode produced nothing — ``last_error`` carries the reason.""" + return not self.traces @dataclass @@ -192,9 +182,6 @@ class TrainBatch: rollouts: TrainRollouts samples: list[TrainingSample] - failures: list[EpisodeFailure] = field(default_factory=list) - """Episodes in the window that produced no traces at all — they appear nowhere in - ``rollouts``, so their count would otherwise be invisible.""" @dataclass @@ -205,8 +192,6 @@ class EvalBatch: env_name: str step: int rollouts: EvalRollouts - failures: list[EpisodeFailure] = field(default_factory=list) - """Episodes in the epoch that produced no traces at all.""" class VersionObserver(Protocol): diff --git a/tests/unit/orchestrator/test_metrics.py b/tests/unit/orchestrator/test_metrics.py index ffe3903d35..5de63262a1 100644 --- a/tests/unit/orchestrator/test_metrics.py +++ b/tests/unit/orchestrator/test_metrics.py @@ -1,13 +1,12 @@ import math from itertools import count from types import SimpleNamespace -from uuid import uuid4 import pytest import verifiers.v1 as vf from prime_rl.orchestrator.metrics import EvalRollouts, Stat, TrainRollouts, episode_failure_metrics -from prime_rl.orchestrator.types import EpisodeFailure, EpisodeResult +from prime_rl.orchestrator.types import Episode from prime_rl.orchestrator.utils import compute_pass_metrics _ids = count() @@ -81,10 +80,10 @@ def mk( ) -def ep(*rollouts): +def ep(*rollouts, env_name: str = "env", errors=()): """One episode over these traces. ``model_construct`` skips validation so the duck-typed stand-ins above can stand in for real ones.""" - return vf.WireEpisode.model_construct(id=f"e{next(_ids)}", traces=list(rollouts)) + return Episode.model_construct(id=f"e{next(_ids)}", traces=list(rollouts), env_name=env_name, errors=list(errors)) def solo(rollouts): @@ -106,7 +105,12 @@ def test_stat(): def test_container_effective_by_env_and_listlike(): rc = TrainRollouts( - solo([mk(env_name="a"), mk(env_name="a", has_error=True), mk(env_name="b", is_filtered=True), mk(env_name="b")]) + [ + ep(mk(env_name="a"), env_name="a"), + ep(mk(env_name="a", has_error=True), env_name="a"), + ep(mk(env_name="b", is_filtered=True), env_name="b"), + ep(mk(env_name="b"), env_name="b"), + ] ) assert len(rc) == 4 and [r.env_name for r in rc] == ["a", "a", "b", "b"] # sized + iterable eff = rc.effective @@ -237,7 +241,9 @@ def test_nested_metrics_and_rewards(): assert eff["train/agg/effective/agent/rewards/format/mean"] == 0.5 # cross-env agg: another env's unscored trace carries different keys, so it can't dilute these other = mk(env_name="other", has_error=True, rewards={"solved": None}) - agg = TrainRollouts(solo(rollouts + [other])).metrics.to_wandb(prefix="train/agg", subset="all") + agg = TrainRollouts(solo(rollouts) + [ep(other, env_name="other")]).metrics.to_wandb( + prefix="train/agg", subset="all" + ) assert agg["train/agg/all/agent/rewards/format/mean"] == pytest.approx(1 / 3) assert agg["train/agg/all/agent/rewards/solved/mean"] == 0.0 @@ -292,41 +298,34 @@ def test_compute_pass_metrics_matches_closed_form(): def test_episode_failure_metrics(): - """Episodes that never produced a trace are counted whole, by reason. They belong to no agent, - so they must never appear under an agent subtree — the phantom-seat bug this replaced.""" - failures = [ - EpisodeFailure(type="Cancelled", message="Off-policy cancel"), - EpisodeFailure(type="Cancelled", message="Off-policy cancel"), - EpisodeFailure(type="TaskFailed", message="boom"), + """An episode that produced no traces is counted whole, by the reason vf left on it. Every + cause — a cancel, a task that raised, an env that ran no agent — reads the same way, and none + of them is ever attributed to a seat.""" + cancelled = vf.Error(type="Cancelled", message="Off-policy cancel") + episodes = [ + ep(mk()), + ep(mk()), + ep(errors=[cancelled]), + ep(errors=[cancelled]), + ep(errors=[vf.Error(type="TaskFailed", message="boom")]), ] - out = episode_failure_metrics(failures, prefix="train/agg", total=9) + out = episode_failure_metrics(episodes, prefix="train/agg") assert out["train/agg/episode_failure/Cancelled"] == 2.0 assert out["train/agg/episode_failure/TaskFailed"] == 1.0 - assert out["train/agg/episode_failure/rate"] == 0.25 # 3 of the window's 12 episodes + assert out["train/agg/episode_failure/rate"] == 0.6 # 3 of 5 episodes assert not any("/agent/" in k for k in out) # never attributed to a seat - assert episode_failure_metrics([], prefix="train/agg", total=4) == {"train/agg/episode_failure/rate": 0.0} - assert episode_failure_metrics([], prefix="train/agg", total=0) == {} # nothing arrived at all - - -def test_episode_result_carries_traces_xor_failure(): - """The envelope makes the phantom trace unrepresentable: an outcome is traces or a failure.""" - episode = ep(mk()) - assert EpisodeResult(kind="train", env_name="env", group_id=uuid4(), policy_version=0, episode=episode) - assert EpisodeResult( - kind="train", - env_name="env", - group_id=uuid4(), - policy_version=0, - failure=EpisodeFailure(type="Cancelled", message="Off-policy cancel"), - ) - with pytest.raises(AssertionError): # neither - EpisodeResult(kind="train", env_name="env", group_id=uuid4(), policy_version=0) - with pytest.raises(AssertionError): # both - EpisodeResult( - kind="train", - env_name="env", - group_id=uuid4(), - policy_version=0, - episode=episode, - failure=EpisodeFailure(type="Cancelled", message="Off-policy cancel"), - ) + assert episode_failure_metrics([ep(mk())], prefix="train/agg") == {"train/agg/episode_failure/rate": 0.0} + assert episode_failure_metrics([], prefix="train/agg") == {} # nothing arrived at all + + +def test_traceless_episode_keeps_its_reason(): + """The failure needs no type of its own: ``failed`` is just "no traces", and the reason is the + error vf (or the dispatcher) already put on the episode.""" + episode = ep(errors=[vf.Error(type="Cancelled", message="Off-policy cancel")]) + assert episode.failed and episode.rollouts == [] + assert episode.last_error is not None and episode.last_error.type == "Cancelled" + assert not ep(mk()).failed + # A traceless episode still counts as an episode, but contributes no rollouts. + pool = TrainRollouts([ep(mk()), episode]) + assert len(pool) == 1 and len(pool.episodes) == 2 + assert len(pool.effective.episodes) == 1 # it survives nothing, so the subset drops it From 56d2c208c5bb686d3ea1dcb82715f35aa50009f0 Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Mon, 3 Aug 2026 23:23:00 +0000 Subject: [PATCH 25/58] refactor(orchestrator)!: train/eval are types, not a kind field MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Episode splits into TrainEpisode and EvalEpisode, Rollout into Rollout and TrainRollout, so each carries only what its path means: eval_step is non-optional on the episode that always has one, off_policy_steps exists only where staleness can happen, and the trainer-bound samples, advantages and filter verdicts are gone from eval traces that never used them. The kind discriminator becomes the class, with KIND left as a ClassVar for the run record and the dispatcher's counters. The dispatch facts leave the traces entirely — they describe the episode, which now carries them — and the sinks read a group's env and step off an episode, which a fully cancelled group still has when it has no trace. Co-Authored-By: Claude Fable 5 --- src/prime_rl/orchestrator/algo/__init__.py | 6 +- src/prime_rl/orchestrator/algo/base.py | 14 ++-- src/prime_rl/orchestrator/algo/echo.py | 6 +- src/prime_rl/orchestrator/algo/grpo.py | 4 +- .../orchestrator/algo/hierarchical_grpo.py | 6 +- src/prime_rl/orchestrator/algo/max_rl.py | 4 +- src/prime_rl/orchestrator/algo/opd.py | 4 +- src/prime_rl/orchestrator/algo/opsd.py | 6 +- src/prime_rl/orchestrator/algo/rae.py | 4 +- src/prime_rl/orchestrator/algo/routing.py | 4 +- src/prime_rl/orchestrator/dispatcher.py | 35 +++++---- src/prime_rl/orchestrator/envs.py | 12 +-- src/prime_rl/orchestrator/eval_sink.py | 18 +++-- src/prime_rl/orchestrator/filters.py | 14 ++-- src/prime_rl/orchestrator/metrics.py | 16 ++-- src/prime_rl/orchestrator/orchestrator.py | 28 ++++--- src/prime_rl/orchestrator/train_sink.py | 20 ++--- src/prime_rl/orchestrator/types.py | 77 +++++++++++++------ tests/unit/orchestrator/test_advantage.py | 22 +++--- tests/unit/orchestrator/test_algorithms.py | 12 +-- tests/unit/orchestrator/test_filters.py | 10 +-- tests/unit/orchestrator/test_metrics.py | 33 ++++++-- tests/unit/utils/test_prime_monitor.py | 10 +-- 23 files changed, 206 insertions(+), 159 deletions(-) diff --git a/src/prime_rl/orchestrator/algo/__init__.py b/src/prime_rl/orchestrator/algo/__init__.py index 745a8ebc79..24b6be1fe9 100644 --- a/src/prime_rl/orchestrator/algo/__init__.py +++ b/src/prime_rl/orchestrator/algo/__init__.py @@ -17,7 +17,7 @@ ``finalize_rollout`` / ``finalize_group`` methods the pipeline drives. Advantages are per-token everywhere they are stored or shipped — there is no scalar advantage in the pipeline. An algorithm assigns credit in its scoring - hook via ``Rollout.assign_advantages``: a scalar that is *broadcast* over the + hook via ``TrainRollout.assign_advantages``: a scalar that is *broadcast* over the rollout's completion tokens (uniform credit, the common case), or an explicit full-length-N per-token list aligned to the concatenated sample token_ids (0.0 off-mask). @@ -39,7 +39,7 @@ from prime_rl.orchestrator.algo.rae import RAEAlgorithm from prime_rl.orchestrator.algo.routing import stamp_advantages, stamp_loss_routing from prime_rl.orchestrator.algo.sft import SFTDistillAlgorithm -from prime_rl.orchestrator.types import Rollout +from prime_rl.orchestrator.types import TrainRollout if TYPE_CHECKING: from prime_rl.configs.algorithm import AlgoConfig @@ -79,7 +79,7 @@ def build_algorithm(config: AlgoConfig, policy_pool: InferencePool) -> Algorithm "OPDAlgorithm", "OPSDAlgorithm", "RAEAlgorithm", - "Rollout", + "TrainRollout", "SFTDistillAlgorithm", "build_algorithm", "connect_frozen_pool", diff --git a/src/prime_rl/orchestrator/algo/base.py b/src/prime_rl/orchestrator/algo/base.py index b12d17df62..11c63ce724 100644 --- a/src/prime_rl/orchestrator/algo/base.py +++ b/src/prime_rl/orchestrator/algo/base.py @@ -49,7 +49,7 @@ if TYPE_CHECKING: from renderers import RendererConfig - from prime_rl.orchestrator.types import Rollout + from prime_rl.orchestrator.types import TrainRollout from prime_rl.utils.client import InferencePool @@ -91,9 +91,9 @@ class Algorithm: (``action_loss_type``); - lifecycle — :meth:`setup` connects client pools to the frozen models the algorithm declares, resolving each reference via :meth:`connect`; - - the two scoring hooks, each ``async`` and given the :class:`Rollout` + - the two scoring hooks, each ``async`` and given the :class:`TrainRollout` directly — read the trace, write credit via - :meth:`Rollout.assign_advantages`. They are + :meth:`TrainRollout.assign_advantages`. They are async so either stage may do I/O — e.g. a process-reward model or a teacher at arrival, or a judge at group time whose signal a pre-batch filter then reads; a hook that only does advantage math simply never @@ -137,24 +137,24 @@ async def connect(self, reference: FrozenModelConfig) -> InferencePool: self.connected_pools.append(pool) return pool - async def score_rollout(self, rollout: Rollout) -> None: + async def score_rollout(self, rollout: TrainRollout) -> None: """Arrival phase, one rollout, before its group is complete: write rollout-local credit (``rollout.assign_advantages``), observation ce weights (echo), or per-token results from a model — an inference pool connected in :meth:`setup`, or the live policy (opsd). No siblings, no group stats.""" - async def score_group(self, group: list[Rollout]) -> None: + async def score_group(self, group: list[TrainRollout]) -> None: """Group phase, the finalized cohort, before filtering: write group-relative credit.""" - async def finalize_rollout(self, rollout: Rollout) -> None: + async def finalize_rollout(self, rollout: TrainRollout) -> None: """Arrival phase (non-virtual): rollout-local scoring as each rollout is tokenized.""" if rollout.samples: await self.score_rollout(rollout) - async def finalize_group(self, rollouts: list[Rollout]) -> None: + async def finalize_group(self, rollouts: list[TrainRollout]) -> None: """Group phase (non-virtual): group-relative scoring, then stamp each sample's wire fields (the advantage stream + loss routing). After this the records are frozen — groups die at stamping.""" diff --git a/src/prime_rl/orchestrator/algo/echo.py b/src/prime_rl/orchestrator/algo/echo.py index d4ecf74fa3..c7c5e15dda 100644 --- a/src/prime_rl/orchestrator/algo/echo.py +++ b/src/prime_rl/orchestrator/algo/echo.py @@ -11,7 +11,7 @@ if TYPE_CHECKING: import verifiers.v1 as vf - from prime_rl.orchestrator.types import Rollout + from prime_rl.orchestrator.types import TrainRollout from prime_rl.utils.client import InferencePool @@ -34,12 +34,12 @@ def __init__(self, config: EchoAlgoConfig, policy_pool: InferencePool): if config.filter is not None: self.filter_fn = partial(import_object(config.filter.import_path), **config.filter.kwargs) - async def score_rollout(self, rollout: Rollout) -> None: + async def score_rollout(self, rollout: TrainRollout) -> None: # Observation weighting is rollout-local; the group-relative GRPO # baseline is inherited unchanged as ``score_group``. self._weight_observations(rollout) - def _weight_observations(self, rollout: Rollout) -> None: + def _weight_observations(self, rollout: TrainRollout) -> None: """Write each sample's ``ce_weights`` stream over the env-provided observation tokens of later turns. Provenance is structural under v1: within a branch, the non-sampled nodes that follow the first model diff --git a/src/prime_rl/orchestrator/algo/grpo.py b/src/prime_rl/orchestrator/algo/grpo.py index 0b87fa7cee..d538fda9bd 100644 --- a/src/prime_rl/orchestrator/algo/grpo.py +++ b/src/prime_rl/orchestrator/algo/grpo.py @@ -8,7 +8,7 @@ from prime_rl.orchestrator.algo.base import Algorithm if TYPE_CHECKING: - from prime_rl.orchestrator.types import Rollout + from prime_rl.orchestrator.types import TrainRollout from prime_rl.utils.client import InferencePool @@ -21,7 +21,7 @@ def __init__(self, config: GRPOAlgoConfig, policy_pool: InferencePool): super().__init__(config, policy_pool) self.length_penalty = config.length_penalty - async def score_group(self, group: list[Rollout]) -> None: + async def score_group(self, group: list[TrainRollout]) -> None: rewards = torch.tensor([rollout.reward for rollout in group], dtype=torch.float32) length_penalty = self.length_penalty if length_penalty is None: diff --git a/src/prime_rl/orchestrator/algo/hierarchical_grpo.py b/src/prime_rl/orchestrator/algo/hierarchical_grpo.py index 2945e341de..a1b4912ca7 100644 --- a/src/prime_rl/orchestrator/algo/hierarchical_grpo.py +++ b/src/prime_rl/orchestrator/algo/hierarchical_grpo.py @@ -7,7 +7,7 @@ from prime_rl.orchestrator.algo.base import Algorithm if TYPE_CHECKING: - from prime_rl.orchestrator.types import Rollout + from prime_rl.orchestrator.types import TrainRollout from prime_rl.utils.client import InferencePool @@ -28,8 +28,8 @@ def __init__(self, config: HierarchicalGRPOAlgoConfig, policy_pool: InferencePoo super().__init__(config, policy_pool) self.episode_agents = set(config.episode_agents) - async def score_group(self, group: list[Rollout]) -> None: - peers: dict[tuple[str, str | None], list[Rollout]] = defaultdict(list) + async def score_group(self, group: list[TrainRollout]) -> None: + peers: dict[tuple[str, str | None], list[TrainRollout]] = defaultdict(list) for rollout in group: episode_scoped = rollout.agent.name in self.episode_agents key = (rollout.agent.name, rollout.episode_id if episode_scoped else None) diff --git a/src/prime_rl/orchestrator/algo/max_rl.py b/src/prime_rl/orchestrator/algo/max_rl.py index 9a3978108d..58c233e8d4 100644 --- a/src/prime_rl/orchestrator/algo/max_rl.py +++ b/src/prime_rl/orchestrator/algo/max_rl.py @@ -7,7 +7,7 @@ from prime_rl.orchestrator.algo.base import Algorithm if TYPE_CHECKING: - from prime_rl.orchestrator.types import Rollout + from prime_rl.orchestrator.types import TrainRollout class MaxRLAlgorithm(Algorithm): @@ -23,7 +23,7 @@ class MaxRLAlgorithm(Algorithm): <= 0 carries no signal and gets zero advantages (the zero-advantage filter drops it, matching the paper's no-success convention).""" - async def score_group(self, group: list[Rollout]) -> None: + async def score_group(self, group: list[TrainRollout]) -> 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 diff --git a/src/prime_rl/orchestrator/algo/opd.py b/src/prime_rl/orchestrator/algo/opd.py index 3135f2a9b2..96ea134738 100644 --- a/src/prime_rl/orchestrator/algo/opd.py +++ b/src/prime_rl/orchestrator/algo/opd.py @@ -8,7 +8,7 @@ from prime_rl.utils.client import StaticInferencePool if TYPE_CHECKING: - from prime_rl.orchestrator.types import Rollout + from prime_rl.orchestrator.types import TrainRollout from prime_rl.transport import TrainingSample from prime_rl.utils.client import InferencePool @@ -37,7 +37,7 @@ async def setup(self) -> None: raise TypeError("opd teacher must be a static endpoint — prefill scoring needs fixed endpoints") self.teacher_pool = pool - async def score_rollout(self, rollout: Rollout) -> None: + async def score_rollout(self, rollout: TrainRollout) -> None: pool = self.teacher_pool assert pool is not None, "teacher pool not connected — Algorithm.setup() must run first" diff --git a/src/prime_rl/orchestrator/algo/opsd.py b/src/prime_rl/orchestrator/algo/opsd.py index 737666bea9..888a2b7391 100644 --- a/src/prime_rl/orchestrator/algo/opsd.py +++ b/src/prime_rl/orchestrator/algo/opsd.py @@ -9,7 +9,7 @@ if TYPE_CHECKING: from renderers.base import Renderer - from prime_rl.orchestrator.types import Rollout + from prime_rl.orchestrator.types import TrainRollout from prime_rl.transport import TrainingSample from prime_rl.utils.client import InferencePool @@ -50,7 +50,7 @@ async def setup(self) -> None: self.renderer = create_renderer(load_tokenizer(self.policy_pool.model_name), self.renderer_config) - def _demonstration(self, rollout: Rollout) -> str: + def _demonstration(self, rollout: TrainRollout) -> str: demonstration = rollout.info.get(self.demo_key) if demonstration is None: demonstration = getattr(rollout.task.data, self.demo_key, None) @@ -61,7 +61,7 @@ def _demonstration(self, rollout: Rollout) -> str: ) return demonstration - async def score_rollout(self, rollout: Rollout) -> None: + async def score_rollout(self, rollout: TrainRollout) -> None: pool = self.teacher_pool renderer = self.renderer assert renderer is not None, "renderer not built — Algorithm.setup() must run first" diff --git a/src/prime_rl/orchestrator/algo/rae.py b/src/prime_rl/orchestrator/algo/rae.py index 264ae48f84..db4240f30e 100644 --- a/src/prime_rl/orchestrator/algo/rae.py +++ b/src/prime_rl/orchestrator/algo/rae.py @@ -7,7 +7,7 @@ from prime_rl.orchestrator.algo.base import Algorithm if TYPE_CHECKING: - from prime_rl.orchestrator.types import Rollout + from prime_rl.orchestrator.types import TrainRollout from prime_rl.utils.client import InferencePool @@ -34,7 +34,7 @@ def __init__(self, config: RAEAlgoConfig, policy_pool: InferencePool): self.decay = config.decay self.baselines: dict[str, float] = defaultdict(float) - async def score_group(self, group: list[Rollout]) -> None: + async def score_group(self, group: list[TrainRollout]) -> None: for rollout in group: baseline = self.baselines[rollout.agent.name] rollout.assign_advantages(rollout.reward - baseline) diff --git a/src/prime_rl/orchestrator/algo/routing.py b/src/prime_rl/orchestrator/algo/routing.py index 0337fa696c..c21283e001 100644 --- a/src/prime_rl/orchestrator/algo/routing.py +++ b/src/prime_rl/orchestrator/algo/routing.py @@ -17,7 +17,7 @@ from prime_rl.transport import TrainingSample if TYPE_CHECKING: - from prime_rl.orchestrator.types import Rollout + from prime_rl.orchestrator.types import TrainRollout def stamp_loss_routing(sample: TrainingSample, action_loss_type: ActionLossType) -> None: @@ -51,7 +51,7 @@ def stamp_loss_routing(sample: TrainingSample, action_loss_type: ActionLossType) sample.ref_kl_weights = action_weights -def stamp_advantages(rollout: Rollout) -> None: +def stamp_advantages(rollout: TrainRollout) -> None: """Stamp the rollout's per-token advantage stream onto its samples' wire fields. The stream is full-length-N — aligned to the samples' ``token_ids`` concatenated in order, 0.0 on non-trainable positions — and sliced across diff --git a/src/prime_rl/orchestrator/dispatcher.py b/src/prime_rl/orchestrator/dispatcher.py index e50475c34c..68f9248edf 100644 --- a/src/prime_rl/orchestrator/dispatcher.py +++ b/src/prime_rl/orchestrator/dispatcher.py @@ -42,6 +42,7 @@ from prime_rl.orchestrator.train_source import TrainSource from prime_rl.orchestrator.types import ( Episode, + EvalEpisode, GroupState, InflightRollout, Policy, @@ -512,7 +513,7 @@ async def handle_completed_rollout(self, task: asyncio.Task) -> None: try: result = task.result() - episodes: list[Episode] = result if isinstance(result, list) else [result] + episodes: list[vf.WireEpisode] = result if isinstance(result, list) else [result] except asyncio.CancelledError: return except Exception as exc: @@ -541,7 +542,7 @@ async def emit_episode( self, meta: InflightRollout, group: GroupState | None, - episode: Episode, + episode: vf.WireEpisode, ) -> None: """Put one completed episode on ``out_q``, stamped with the facts of the dispatch it came from. Pops the group from ``self.groups`` once every owed episode has been emitted. The @@ -557,21 +558,21 @@ async def emit_episode( if meta.kind == "eval": assert eval_step is not None, "eval episode missing eval_step" - for rollout in episode.rollouts: - rollout.kind = meta.kind + for rollout in episode.traces: rollout.env_name = meta.env_name rollout.group_id = meta.group_id - rollout.policy_version = policy_version - rollout.off_policy_steps = meta.off_policy_steps - if meta.kind == "eval": - rollout.eval_step = eval_step - episode.kind = meta.kind - episode.env_name = meta.env_name - episode.group_id = meta.group_id - episode.policy_version = policy_version - episode.off_policy_steps = meta.off_policy_steps - episode.eval_step = eval_step if meta.kind == "eval" else None - await self.out_q.put(episode) + shared = { + **dict(episode), + "env_name": meta.env_name, + "group_id": meta.group_id, + "policy_version": policy_version, + } + dispatched: Episode = ( + EvalEpisode.model_construct(**shared, eval_step=eval_step) + if meta.kind == "eval" + else TrainEpisode.model_construct(**shared, off_policy_steps=meta.off_policy_steps) + ) + await self.out_q.put(dispatched) async def emit_failed_episodes(self, meta: InflightRollout, group: GroupState | None, error: vf.Error) -> None: """Emit one traceless episode per rollout the task owed, so the sink still counts its way to @@ -607,7 +608,7 @@ async def drop_group(self, group_id: uuid.UUID) -> int: cancel = vf.Error(type="Cancelled", message="Off-policy cancel") for _, meta in claimed: for _ in range(meta.rollout_count): - await self.emit_episode(meta, group, Episode.model_construct(errors=[cancel])) + await self.emit_episode(meta, group, vf.WireEpisode.model_construct(errors=[cancel])) # For non-group-scoring envs, the group may have rollouts that # were never dispatched (``rollouts_to_schedule > 0``). Emit @@ -628,7 +629,7 @@ async def drop_group(self, group_id: uuid.UUID) -> int: ) unscheduled_cancelled = group.rollouts_to_schedule for _ in range(unscheduled_cancelled): - await self.emit_episode(fallback_meta, group, Episode.model_construct(errors=[cancel])) + await self.emit_episode(fallback_meta, group, vf.WireEpisode.model_construct(errors=[cancel])) cancelled = inflight_cancelled + unscheduled_cancelled if cancelled > 0: diff --git a/src/prime_rl/orchestrator/envs.py b/src/prime_rl/orchestrator/envs.py index 01265492ff..f796697d31 100644 --- a/src/prime_rl/orchestrator/envs.py +++ b/src/prime_rl/orchestrator/envs.py @@ -37,16 +37,16 @@ from prime_rl.configs.orchestrator import EnvConfig, EvalSourceConfig, TrainSourceConfig from prime_rl.orchestrator.algo import Algorithm, build_algorithm from prime_rl.orchestrator.sampler import Sampler -from prime_rl.orchestrator.types import Episode, Rollout +from prime_rl.orchestrator.types import TrainRollout from prime_rl.utils.logger import get_logger # Every wire trace validates into this type. WireTaskData (extra="allow") keeps the env's task # fields without importing the env package — the orchestrator never reads them typed (only # task.idx + task.model_dump). -ROLLOUT_TYPE = Rollout[vf.WireTaskData] -# The env server answers a wire episode; we keep that envelope, re-type its traces, and let the -# dispatcher stamp the scheduling facts onto it. -EPISODE_TYPE = Episode +ROLLOUT_TYPE = TrainRollout[vf.WireTaskData] +# The env server answers a wire episode; we keep that envelope and only re-type its traces. The +# dispatcher then mints the Train/Eval episode that carries the dispatch's own facts. +EPISODE_TYPE = vf.WireEpisode # Max wait for a spawned env server to bind and report its address. A legacy # child loads its dataset before reporting, so this is generous. @@ -235,7 +235,7 @@ async def run( async def run_group( self, client: vf.ClientConfig, task_idx: int, model_name: str, group_size: int, cache_salt: str | None - ) -> list[Rollout]: + ) -> list[TrainRollout]: """Run a group of rollouts for ``task_idx`` (group-scoring envs); return typed Traces.""" wires = await self.env_client.run_group( task_idx=task_idx, diff --git a/src/prime_rl/orchestrator/eval_sink.py b/src/prime_rl/orchestrator/eval_sink.py index ba8d014640..7b33c2c453 100644 --- a/src/prime_rl/orchestrator/eval_sink.py +++ b/src/prime_rl/orchestrator/eval_sink.py @@ -19,7 +19,7 @@ from prime_rl.orchestrator.envs import EvalEnvs from prime_rl.orchestrator.metrics import EvalRollouts -from prime_rl.orchestrator.types import Episode, EvalBatch, Rollout +from prime_rl.orchestrator.types import EvalBatch, EvalEpisode, Rollout from prime_rl.utils.logger import get_logger @@ -28,13 +28,13 @@ class EvalSink: def __init__(self, *, eval_envs: EvalEnvs) -> None: self.eval_envs = eval_envs - self.pending_groups: dict[uuid.UUID, list[Episode]] = defaultdict(list) + self.pending_groups: dict[uuid.UUID, list[EvalEpisode]] = defaultdict(list) # Episodes arrived per group / per batch bucket — the finalization counts. self.pending_group_episodes: dict[uuid.UUID, int] = defaultdict(int) - self.pending_batches: dict[tuple[str, int], list[Episode]] = defaultdict(list) + self.pending_batches: dict[tuple[str, int], list[EvalEpisode]] = defaultdict(list) self.pending_batch_episodes: dict[tuple[str, int], int] = defaultdict(int) - def add(self, episode: Episode) -> EvalBatch | None: + def add(self, episode: EvalEpisode) -> EvalBatch | None: """Process one episode arrival; finalize the group on the ``group_size``-th episode and the per-env epoch on the ``num_examples × group_size``-th. A failed episode brings no rollouts but still counts toward both.""" @@ -102,10 +102,12 @@ def process_group(self, group_id: uuid.UUID) -> None: episodes = self.pending_group_episodes.pop(group_id, 0) if not finished: return - group = [trace for episode in finished for trace in episode.traces] - env_name = group[0].env_name - task_idx = group[0].task.data.idx - eval_step = group[0].eval_step + # Read the group's facts off an episode, not a trace: every episode in it may have + # produced none (a whole group cancelled off-policy). + env_name = finished[0].env_name + eval_step = finished[0].eval_step + group = [trace for episode in finished for trace in episode.rollouts] + task_idx = group[0].task.data.idx if group else -1 bucket = self.pending_batches[(env_name, eval_step)] bucket.extend(finished) self.pending_batch_episodes[(env_name, eval_step)] += episodes diff --git a/src/prime_rl/orchestrator/filters.py b/src/prime_rl/orchestrator/filters.py index ad023fd928..01f241c1c0 100644 --- a/src/prime_rl/orchestrator/filters.py +++ b/src/prime_rl/orchestrator/filters.py @@ -16,7 +16,7 @@ from prime_rl.utils.logger import get_logger if TYPE_CHECKING: - from prime_rl.orchestrator.types import Rollout + from prime_rl.orchestrator.types import TrainRollout @dataclass @@ -28,7 +28,7 @@ class RolloutFilter(Protocol): name: str enforce: bool - def check(self, rollout: Rollout) -> FilterResult: ... + def check(self, rollout: TrainRollout) -> FilterResult: ... @dataclass @@ -48,7 +48,7 @@ class GibberishFilter: logprob_threshold: float enforce: bool = False - def check(self, rollout: Rollout) -> FilterResult: + def check(self, rollout: TrainRollout) -> FilterResult: for branch in rollout.branches: # branch.{token_ids,logprobs,sampled_mask} are flat and mutually aligned; the raw # node arrays are not (node.logprobs covers only the sampled suffix, not the @@ -78,7 +78,7 @@ class RepetitionFilter: logprob_threshold: float enforce: bool = False - def check(self, rollout: Rollout) -> FilterResult: + def check(self, rollout: TrainRollout) -> FilterResult: for branch in rollout.branches: # Aligned branch streams (see GibberishFilter), and reset the streak per branch: # flat rollout.nodes interleaves distinct root->leaf paths (compaction/subagents), @@ -104,7 +104,7 @@ class ZeroAdvantageFilter: name: str enforce: bool = True - def check(self, rollout: Rollout) -> FilterResult: + def check(self, rollout: TrainRollout) -> FilterResult: if rollout.advantages is not None and all(a == 0.0 for a in rollout.advantages): return FilterResult(detected=True) return FilterResult(detected=False) @@ -146,8 +146,8 @@ def setup_filters(configs: list[FilterConfig], vocab_size: int, *, kind: str) -> return filters -def apply_filters(filters: list[RolloutFilter], rollouts: list[Rollout]) -> None: - """Flag ``Rollout``\\ s in place with per-filter detection + drop decision. +def apply_filters(filters: list[RolloutFilter], rollouts: list[TrainRollout]) -> None: + """Flag ``TrainRollout``\\ s in place with per-filter detection + drop decision. Each rollout's ``filter_results`` dict records per-filter detection bools; ``is_filtered`` is True iff an enforcing filter detected it. First matching diff --git a/src/prime_rl/orchestrator/metrics.py b/src/prime_rl/orchestrator/metrics.py index 11baa6c907..1f3d1f5043 100644 --- a/src/prime_rl/orchestrator/metrics.py +++ b/src/prime_rl/orchestrator/metrics.py @@ -26,7 +26,7 @@ from prime_rl.orchestrator.utils import compute_pass_metrics if TYPE_CHECKING: - from prime_rl.orchestrator.types import Episode, Rollout + from prime_rl.orchestrator.types import Episode, EvalEpisode, Rollout, TrainEpisode, TrainRollout Subset = Literal["all", "effective"] @@ -419,20 +419,20 @@ class TrainRollouts: narrowed to its surviving traces; ``metrics`` builds ``TrainMetrics`` over them. Sized and iterated by *rollout*, since that is the unit almost every consumer wants.""" - def __init__(self, episodes: list[Episode] | None = None) -> None: + def __init__(self, episodes: list[TrainEpisode] | None = None) -> None: self.episodes = episodes if episodes is not None else [] - def append(self, episode: Episode) -> None: + def append(self, episode: TrainEpisode) -> None: self.episodes.append(episode) @property - def rollouts(self) -> list[Rollout]: - return [trace for episode in self.episodes for trace in episode.traces] + def rollouts(self) -> list[TrainRollout]: + return [trace for episode in self.episodes for trace in episode.rollouts] def __len__(self) -> int: return sum(len(episode.traces) for episode in self.episodes) - def __iter__(self) -> Iterator[Rollout]: + def __iter__(self) -> Iterator[TrainRollout]: return iter(self.rollouts) @property @@ -444,7 +444,7 @@ def effective(self) -> TrainRollouts: return TrainRollouts([episode for episode in kept if episode is not None]) def by_env(self) -> dict[str, TrainRollouts]: - grouped: dict[str, list[Episode]] = {} + grouped: dict[str, list[TrainEpisode]] = {} for episode in self.episodes: grouped.setdefault(episode.env_name, []).append(episode) return {env: TrainRollouts(episodes) for env, episodes in grouped.items()} @@ -460,7 +460,7 @@ class EvalRollouts: ``group_size`` (rollouts per example, the ``avg@k`` k) is derived from the full epoch and carried onto ``effective`` so both subsets share one stable key; ``metrics`` builds ``EvalMetrics``.""" - def __init__(self, episodes: list[Episode] | None = None, group_size: int | None = None) -> None: + def __init__(self, episodes: list[EvalEpisode] | None = None, group_size: int | None = None) -> None: self.episodes = episodes if episodes is not None else [] self._group_size = group_size diff --git a/src/prime_rl/orchestrator/orchestrator.py b/src/prime_rl/orchestrator/orchestrator.py index 6a06250b7a..1c820d1c1b 100644 --- a/src/prime_rl/orchestrator/orchestrator.py +++ b/src/prime_rl/orchestrator/orchestrator.py @@ -60,6 +60,7 @@ from prime_rl.orchestrator.types import ( Episode, EvalBatch, + EvalEpisode, Policy, Progress, TrainBatch, @@ -530,13 +531,10 @@ async def main_loop(self) -> None: # Train rollouts belong to the batch window currently collecting (``progress.step``), # eval rollouts to the step whose eval triggered them. A failed episode has no trace # to write; it is counted by the window's episode-failure metrics instead. - kind = episode.kind - step = episode.eval_step if kind == "eval" else self.progress.step - assert step is not None + is_eval = isinstance(episode, EvalEpisode) + step = episode.eval_step if isinstance(episode, EvalEpisode) else self.progress.step run: vf.RunInfo = ( - vf.EvalRunInfo(id=self.run_id, step=step) - if kind == "eval" - else vf.TrainRunInfo(id=self.run_id, step=step) + vf.EvalRunInfo(id=self.run_id, step=step) if is_eval else vf.TrainRunInfo(id=self.run_id, step=step) ) for rollout in episode.rollouts: rollout.record_run( @@ -544,16 +542,16 @@ async def main_loop(self) -> None: env_name=rollout.env_name, group_id=str(rollout.group_id), episode_id=rollout.episode_id, - policy_version=rollout.policy_version, + policy_version=episode.policy_version, ) if not episode.failed: await asyncio.to_thread( save_rollouts, [rollout.to_record() for rollout in episode.rollouts], - get_trace_path(self.config.output_dir, step, kind, "all"), + get_trace_path(self.config.output_dir, step, episode.KIND, "all"), ) - if kind == "eval": + if isinstance(episode, EvalEpisode): assert self.eval_sink is not None # eval rollouts only emitted when eval is configured eval_batch = self.eval_sink.add(episode) if eval_batch is not None: @@ -638,9 +636,9 @@ async def finalize_train_batch(self, batch: TrainBatch) -> None: # off-policy — queue time included, unlike the dispatcher's in-flight # counter, which only sees weight updates during generation. Frozen- # sourced rollouts stay 0 (their sampler doesn't follow the policy). - for r in batch.rollouts: - if self.train_envs.get(r.env_name).sampler.samples_from_live_policy: - r.off_policy_steps = (step - 1) - r.policy_version + for train_episode in batch.rollouts.episodes: + if self.train_envs.get(train_episode.env_name).sampler.samples_from_live_policy: + train_episode.off_policy_steps = (step - 1) - train_episode.policy_version # The effective (clean, trained-on) subset lands in the per-step ``effective`` trace file # at ship time; the full arrival window already streamed into ``all`` on arrival. @@ -824,7 +822,7 @@ def log_train_batch(self, batch: TrainBatch, *, step: int, step_time: float) -> n_effective = len(effective) n_trainable = sum(1 for r in effective if r.is_trainable) trainable_rate = (n_trainable / n_effective) if n_effective else 0.0 - max_off_policy = max((r.off_policy_steps for r in effective), default=0) + max_off_policy = max((e.off_policy_steps for e in effective.episodes), default=0) head = ( f"Step {step} | {format_time(step_time):>7} | Reward {eff.reward.mean():.4f} | " @@ -848,7 +846,7 @@ def log_train_batch(self, batch: TrainBatch, *, step: int, step_time: float) -> lines.append( f"╰─ {env_name:<{name_width}} | Ratio {ratio:.1%} | Reward {env_eff.reward.mean():.4f} | " f"Turns {env_eff.num_turns.mean():.1f} | Branches {env_eff.num_branches.mean():.1f} | " - f"Max Off-Policy {max((r.off_policy_steps for r in env_eff_pool), default=0)} | " + f"Max Off-Policy {max((e.off_policy_steps for e in env_eff_pool.episodes), default=0)} | " f"Error {pool.metrics.has_error.mean():.1%} | Truncation {env_eff.is_truncated.mean():.1%}" ) get_logger().success("\n\t\t ".join(lines)) @@ -869,7 +867,7 @@ async def finalize_eval_batch(self, batch: EvalBatch) -> None: save_rollouts, records, get_trace_path(self.config.output_dir, batch.step, "eval", "effective") ) self.monitor.log_eval_samples(batch.rollouts, env_name=batch.env_name, step=batch.step) - policy_versions = {r.policy_version for r in batch.rollouts} + policy_versions = {e.policy_version for e in batch.rollouts.episodes} policy_version = min(policy_versions) if len(policy_versions) > 1: get_logger().warning( diff --git a/src/prime_rl/orchestrator/train_sink.py b/src/prime_rl/orchestrator/train_sink.py index a154cf4f00..60deeec462 100644 --- a/src/prime_rl/orchestrator/train_sink.py +++ b/src/prime_rl/orchestrator/train_sink.py @@ -27,12 +27,12 @@ from prime_rl.orchestrator.filters import RolloutFilter, apply_filters from prime_rl.orchestrator.metrics import TrainRollouts from prime_rl.orchestrator.trajectories import trace_to_samples -from prime_rl.orchestrator.types import Episode, Rollout, TrainBatch +from prime_rl.orchestrator.types import TrainBatch, TrainEpisode, TrainRollout from prime_rl.transport import TrainingSample from prime_rl.utils.logger import get_logger -def payload_tokens(rollout: Rollout) -> int: +def payload_tokens(rollout: TrainRollout) -> int: """Token cost of the rollout's trainer-bound payload — the samples built by ``process_rollout``. This is what actually ships: forked traces can drop branches with no trainable tokens, so ``Trace.num_total_tokens`` (which sums @@ -79,11 +79,11 @@ def __init__( # Keyed by the dispatcher's group UUID. ``(env_name, task_idx)`` # isn't unique — the same task can be re-sampled while an # earlier group is still in flight - self.pending_groups: dict[uuid.UUID, list[Episode]] = defaultdict(list) + self.pending_groups: dict[uuid.UUID, list[TrainEpisode]] = defaultdict(list) # Episodes arrived per group — the finalization count (an episode may # add several traces to ``pending_groups`` but counts once here). self.pending_group_episodes: dict[uuid.UUID, int] = defaultdict(int) - self.pending_batch: list[Rollout] = [] + self.pending_batch: list[TrainRollout] = [] # Running payload-token total of ``pending_batch`` (token-batched # runs), kept in sync on append/pop so the readiness check never # re-sums per arrival. @@ -124,7 +124,7 @@ def pending_batch_by_env(self) -> dict[str, int]: counts[r.env_name] += 1 return dict(counts) - async def add(self, episode: Episode) -> TrainBatch | None: + async def add(self, episode: TrainEpisode) -> TrainBatch | None: """Process one episode arrival; finalize the group on the ``group_size``-th episode; return a ``TrainBatch`` if the finalization pushed (or left) the batch over its threshold. Arrivals into @@ -151,7 +151,7 @@ async def add(self, episode: Episode) -> TrainBatch | None: return self.process_batch() return None - async def process_rollout(self, rollout: Rollout) -> None: + async def process_rollout(self, rollout: TrainRollout) -> None: """Build training samples from the rollout's Trace (one per branch), walking the message graph. Training is renderer-only across all modes (RL/OPD student, SFT teacher), so every node already carries its tokens. Errored rollouts are dropped at the group @@ -178,15 +178,17 @@ async def process_group(self, group_id: uuid.UUID) -> None: self.pending_group_episodes.pop(group_id, None) if not episodes: return - group = [trace for episode in episodes for trace in episode.traces] + # Read the group's facts off an episode, not a trace: every episode in it may have + # produced none (a whole group cancelled off-policy). + env_name = episodes[0].env_name + group = [trace for episode in episodes for trace in episode.rollouts] # Window membership follows group finalization, not arrival: a rollout # only becomes observable (metrics / persistence) once its whole group # is finalized, so a batch's window never claims rollouts of a group # that ships later. Dropped groups still land here — they were observed. for episode in episodes: self.pending_rollouts.append(episode) - env_name = group[0].env_name - task_idx = group[0].task.data.idx + task_idx = group[0].task.data.idx if group else -1 survivors = [r for r in group if not r.has_error] num_errored = len(group) - len(survivors) diff --git a/src/prime_rl/orchestrator/types.py b/src/prime_rl/orchestrator/types.py index cc134a4dae..9077946282 100644 --- a/src/prime_rl/orchestrator/types.py +++ b/src/prime_rl/orchestrator/types.py @@ -4,7 +4,7 @@ import uuid from dataclasses import dataclass -from typing import TYPE_CHECKING, Generic, Literal, Protocol, cast +from typing import TYPE_CHECKING, ClassVar, Generic, Literal, Protocol, cast import verifiers.v1 as vf from pydantic import ConfigDict, Field @@ -73,29 +73,30 @@ class GroupState: class Rollout(vf.Trace[DataT], Generic[DataT]): - """A completed rollout: the env's typed ``vf.Trace`` *is* the rollout — prime-rl's - orchestration metadata lives on it directly (set by the dispatcher once the rollout - returns), so there's no wrapper. Train vs eval is the ``kind`` discriminator. All metadata - fields are ``exclude=True``, so dumping a Rollout yields a plain trace on the wire; the - orchestrator mirrors them onto the trace's own fields via ``vf.Trace.record_run`` (``kind`` as - ``run.type``, the run id, the step, and ``env_name``/``group_id``/``policy_version`` into - ``info``) when a rollout arrives, so the on-disk records stay fully placeable. - - It is also the single currency the scoring hooks receive: a hook reads the trace - directly (``rollout.reward``, ``rollout.nodes``, ``rollout.num_turns``) and writes - credit through :meth:`assign_advantages` (scalar broadcast or per-token), which - spreads over the samples' trainable (mask-True) tokens.""" + """A completed rollout: the env's typed ``vf.Trace`` *is* the rollout, carrying only the links + prime-rl needs to place a loose trace back among its peers — its episode, its comparison group, + its env. Everything about the dispatch itself lives on the ``Episode``, which is the thing that + was dispatched. All added fields are ``exclude=True``, so dumping a Rollout yields a plain + trace on the wire; ``vf.Trace.record_run`` mirrors them into ``info`` on arrival so the on-disk + records stay fully placeable. - model_config = ConfigDict(arbitrary_types_allowed=True) # ``samples`` holds msgspec structs + It is also the currency the scoring hooks receive: a hook reads the trace directly + (``rollout.reward``, ``rollout.nodes``, ``rollout.num_turns``).""" - kind: RolloutKind = Field(default="train", exclude=True) env_name: str = Field(default="", exclude=True) group_id: uuid.UUID = Field(default_factory=uuid.uuid4, exclude=True) # Links the traces of one episode; stamped into ``info`` on arrival so # saved records keep their grouping. episode_id: str = Field(default="", exclude=True) - policy_version: int = Field(default=0, exclude=True) - off_policy_steps: int = Field(default=0, exclude=True) + + +class TrainRollout(Rollout[DataT], Generic[DataT]): + """A rollout on the training path, which alone carries training state: the trainer-bound + samples built from its branches, the credit assigned over them, and the filter verdicts. Eval + rollouts have none of this, so they are plain ``Rollout``\\ s and can't be asked for it.""" + + model_config = ConfigDict(arbitrary_types_allowed=True) # ``samples`` holds msgspec structs + samples: list[TrainingSample] = Field(default_factory=list, exclude=True) # Per-token rl advantage stream, full-length-N (= len(token_ids)) per # sample, concatenated across the rollout's samples in order; 0.0 on @@ -104,7 +105,6 @@ class Rollout(vf.Trace[DataT], Generic[DataT]): advantages: list[float] | None = Field(default=None, exclude=True) is_filtered: bool = Field(default=False, exclude=True) filter_results: dict[str, bool] = Field(default_factory=dict, exclude=True) - eval_step: int | None = Field(default=None, exclude=True) def assign_advantages(self, values: float | list[float]) -> None: """Write the rl advantage stream: a scalar broadcast over the @@ -141,24 +141,29 @@ def is_trainable(self) -> bool: class Episode(vf.WireEpisode): - """The env's own ``vf.Episode`` extended with prime-rl's scheduling facts — the only thing - prime-rl genuinely adds, so the episode itself travels rather than a wrapper around it. Those - fields are ``exclude=True``, so dumping an Episode yields a plain wire episode. + """The env's own ``vf.Episode`` extended with the facts of the dispatch it came from — the only + thing prime-rl genuinely adds, so the episode itself travels rather than a wrapper around it. + Those fields are ``exclude=True``, so dumping an Episode yields a plain wire episode. An episode that produced no traces is not a special case and needs no stand-in rollout: vf already records why on ``errors`` (its ``run_episode`` puts the exception there and returns the episode with ``ok`` false), and prime-rl's own outcomes — an off-policy cancel, a task that raised before reaching the env — are minted the same way. So ``failed`` is simply "no traces", - and ``last_error`` says why in one vocabulary for every cause.""" + and ``last_error`` says why in one vocabulary for every cause. + + Train and eval are the two subclasses rather than a discriminator field, so each carries only + what its path means: an eval episode has a step it belongs to, a train episode has the policy + it was generated from.""" model_config = ConfigDict(arbitrary_types_allowed=True) # traces are ``Rollout``s - kind: RolloutKind = Field(default="train", exclude=True) + KIND: ClassVar[RolloutKind] + """Which path this episode is on, for the run record and the dispatcher's counters.""" + env_name: str = Field(default="", exclude=True) group_id: uuid.UUID = Field(default_factory=uuid.uuid4, exclude=True) policy_version: int = Field(default=0, exclude=True) - off_policy_steps: int = Field(default=0, exclude=True) - eval_step: int | None = Field(default=None, exclude=True) + """The policy that generated it — the thing being trained on one path, measured on the other.""" @property def rollouts(self) -> list[Rollout]: @@ -171,6 +176,28 @@ def failed(self) -> bool: return not self.traces +class TrainEpisode(Episode): + """An episode collected for training, which alone can go stale relative to the live policy.""" + + KIND: ClassVar[RolloutKind] = "train" + + off_policy_steps: int = Field(default=0, exclude=True) + """How stale it was by the time it shipped — meaningless for eval, which never trains on it.""" + + @property + def rollouts(self) -> list[TrainRollout]: + return cast(list[TrainRollout], self.traces) + + +class EvalEpisode(Episode): + """An episode collected for one eval epoch. ``eval_step`` is the step whose eval triggered it — + always known, unlike on the train path, so it is not optional here.""" + + KIND: ClassVar[RolloutKind] = "eval" + + eval_step: int = Field(default=0, exclude=True) + + @dataclass class TrainBatch: """``rollouts`` is the observation window since the last ship — every rollout of every group diff --git a/tests/unit/orchestrator/test_advantage.py b/tests/unit/orchestrator/test_advantage.py index 9da1e56e21..d2ae909afe 100644 --- a/tests/unit/orchestrator/test_advantage.py +++ b/tests/unit/orchestrator/test_advantage.py @@ -11,7 +11,7 @@ from prime_rl.orchestrator.algo.grpo import GRPOAlgorithm from prime_rl.orchestrator.algo.max_rl import MaxRLAlgorithm from prime_rl.orchestrator.trajectories import trace_to_samples -from prime_rl.orchestrator.types import Rollout +from prime_rl.orchestrator.types import TrainRollout def _build_rollout( @@ -21,8 +21,8 @@ def _build_rollout( obs_lengths: list[int] | None = None, env_name: str = "test", metrics: dict | None = None, -) -> Rollout: - """Build a ``Rollout`` (a ``vf.Trace``) as an alternating message graph. +) -> TrainRollout: + """Build a ``TrainRollout`` (a ``vf.Trace``) as an alternating message graph. ``sampled_lengths`` gives the token count of each model turn (a sampled ``AssistantMessage`` node); ``obs_lengths`` (one shorter, if given) gives the @@ -97,7 +97,7 @@ def _take(n: int) -> list[int]: ) parent = len(nodes) - 1 - rollout = Rollout[vf.TaskData]( + rollout = TrainRollout[vf.TaskData]( task=vf.TraceTask(type="Task", data=vf.TaskData(idx=0, prompt=None)), agent=vf.AgentInfo(config=vf.AgentConfig()), nodes=nodes, @@ -116,8 +116,8 @@ def _make_rollout( num_turns: int = 1, env_name: str = "test", metrics: dict | None = None, -) -> Rollout: - """Build a ``Rollout`` carrying ``completion_len`` model-sampled tokens split +) -> TrainRollout: + """Build a ``TrainRollout`` carrying ``completion_len`` model-sampled tokens split across ``num_turns`` sampled turns. Always carries at least one trainable token so credit broadcasts somewhere.""" num_turns = max(num_turns, 1) @@ -127,8 +127,8 @@ def _make_rollout( return _build_rollout(reward, sampled_lengths=sampled_lengths, env_name=env_name, metrics=metrics) -def _make_group(rewards, completion_lengths=None, num_turns=None) -> list[Rollout]: - """Build one group of ``Rollout``\\ s from 1D arrays of rewards/lengths/turns — +def _make_group(rewards, completion_lengths=None, num_turns=None) -> list[TrainRollout]: + """Build one group of ``TrainRollout``\\ s from 1D arrays of rewards/lengths/turns — exactly what ``score_group`` sees.""" rollouts = [] for i, reward in enumerate(rewards): @@ -138,21 +138,21 @@ def _make_group(rewards, completion_lengths=None, num_turns=None) -> list[Rollou return rollouts -def _scalar(rollout: Rollout) -> float: +def _scalar(rollout: TrainRollout) -> float: """The per-rollout advantage scalar an algorithm assigned — broadcast over the rollout's trainable (mask-True) tokens, so any trainable position holds it.""" mask = [m for sample in rollout.samples for m in sample.mask] return rollout.advantages[mask.index(True)] -def _grpo(group: list[Rollout], length_penalty=None) -> list[float]: +def _grpo(group: list[TrainRollout], length_penalty=None) -> list[float]: """Drive ``GRPOAlgorithm.score_group`` and read back each per-rollout scalar.""" algo = GRPOAlgorithm(GRPOAlgoConfig(length_penalty=length_penalty), policy_pool=None) asyncio.run(algo.score_group(group)) return [_scalar(rollout) for rollout in group] -def _max_rl(group: list[Rollout]) -> list[float]: +def _max_rl(group: list[TrainRollout]) -> list[float]: """Drive ``MaxRLAlgorithm.score_group`` and read back each per-rollout scalar.""" algo = MaxRLAlgorithm(MaxRLAlgoConfig(), policy_pool=None) asyncio.run(algo.score_group(group)) diff --git a/tests/unit/orchestrator/test_algorithms.py b/tests/unit/orchestrator/test_algorithms.py index b2dbb01d05..e17d8bd158 100644 --- a/tests/unit/orchestrator/test_algorithms.py +++ b/tests/unit/orchestrator/test_algorithms.py @@ -10,7 +10,7 @@ from prime_rl.configs.algorithm import AlgoConfig, FrozenModelConfig from prime_rl.orchestrator.algo import EchoAlgorithm, stamp_advantages, stamp_loss_routing from prime_rl.orchestrator.trajectories import trace_to_samples -from prime_rl.orchestrator.types import Rollout +from prime_rl.orchestrator.types import TrainRollout from prime_rl.transport.types import TrainingSample FROZEN = {"name": "org/ref-model", "base_url": ["http://ref:8001/v1"]} @@ -160,8 +160,8 @@ def test_stamp_loss_routing_merges_action_weights_into_ce_stream(): def _make_rollout( samples: list[TrainingSample], advantages: list[float] | None = None, -) -> Rollout: - rollout = Rollout( +) -> TrainRollout: + rollout = TrainRollout( task=vf.TraceTask(type="Task", data=vf.TaskData(idx=0, prompt=None)), agent=vf.AgentInfo(config=vf.AgentConfig()), nodes=[], @@ -244,7 +244,7 @@ def _node(message, *, parent, sampled, token_ids, logprobs=None, is_content=None ) -def _two_turn_rollout(observation_role: str = "tool") -> Rollout: +def _two_turn_rollout(observation_role: str = "tool") -> TrainRollout: """A single linear branch: user prompt, an assistant response, an env-provided observation (tool output / user feedback), then a second assistant response. Tokens: prompt [1,2], action [3,4], observation @@ -259,7 +259,7 @@ def _two_turn_rollout(observation_role: str = "tool") -> Rollout: _node(obs_message, parent=1, sampled=False, token_ids=[5, 6]), _node(AssistantMessage(content="A2"), parent=2, sampled=True, token_ids=[7, 8], logprobs=[-0.3, -0.4]), ] - rollout = Rollout( + rollout = TrainRollout( task=vf.TraceTask(type="Task", data=vf.TaskData(idx=0, prompt=None)), agent=vf.AgentInfo(config=vf.AgentConfig()), nodes=nodes, @@ -311,7 +311,7 @@ def test_echo_weights_only_content_tokens_when_is_content_present(): ), _node(AssistantMessage(content="A2"), parent=2, sampled=True, token_ids=[7, 8], logprobs=[-0.3, -0.4]), ] - rollout = Rollout( + rollout = TrainRollout( task=vf.TraceTask(type="Task", data=vf.TaskData(idx=0, prompt=None)), agent=vf.AgentInfo(config=vf.AgentConfig()), nodes=nodes, diff --git a/tests/unit/orchestrator/test_filters.py b/tests/unit/orchestrator/test_filters.py index 69ce76d029..c344de9010 100644 --- a/tests/unit/orchestrator/test_filters.py +++ b/tests/unit/orchestrator/test_filters.py @@ -11,7 +11,7 @@ setup_filter, setup_filters, ) -from prime_rl.orchestrator.types import Rollout +from prime_rl.orchestrator.types import TrainRollout def _assistant_node(token_ids: list[int], logprobs: list[float]) -> vf.MessageNode: @@ -46,8 +46,8 @@ def _make_rollout( *, reward: float = 1.0, multi_step: bool = False, -) -> Rollout: - """Build a ``Rollout`` (a message-graph trace) carrying the completion tokens — enough for +) -> TrainRollout: + """Build a ``TrainRollout`` (a message-graph trace) carrying the completion tokens — enough for the filters to inspect each node's sampled tokens / logprobs.""" if multi_step: mid = len(completion_ids) // 2 @@ -57,7 +57,7 @@ def _make_rollout( ] else: nodes = [_assistant_node(completion_ids, completion_logprobs)] - rollout = Rollout[vf.TaskData]( + rollout = TrainRollout[vf.TaskData]( task=vf.TraceTask(type="Task", data=vf.TaskData(idx=0, prompt="")), agent=vf.AgentInfo(config=vf.AgentConfig()), nodes=nodes, @@ -140,7 +140,7 @@ def test_gibberish_aligns_logprobs_under_generation_prompt_scaffold(): it; reading the aligned branch streams detects it.""" gibberish_filter = _make_gibberish_filter() - rollout = Rollout[vf.TaskData]( + rollout = TrainRollout[vf.TaskData]( task=vf.TraceTask(type="Task", data=vf.TaskData(idx=0, prompt="")), agent=vf.AgentInfo(config=vf.AgentConfig()), nodes=[_scaffold_assistant_node([50, 80, 120_000], [-1.0, -0.5, gibberish_filter.logprob_threshold - 1.0])], diff --git a/tests/unit/orchestrator/test_metrics.py b/tests/unit/orchestrator/test_metrics.py index 5de63262a1..8c6cd611d3 100644 --- a/tests/unit/orchestrator/test_metrics.py +++ b/tests/unit/orchestrator/test_metrics.py @@ -6,7 +6,7 @@ import verifiers.v1 as vf from prime_rl.orchestrator.metrics import EvalRollouts, Stat, TrainRollouts, episode_failure_metrics -from prime_rl.orchestrator.types import Episode +from prime_rl.orchestrator.types import EvalEpisode, Rollout, TrainEpisode, TrainRollout from prime_rl.orchestrator.utils import compute_pass_metrics _ids = count() @@ -80,15 +80,15 @@ def mk( ) -def ep(*rollouts, env_name: str = "env", errors=()): +def ep(*rollouts, env_name: str = "env", errors=(), cls=TrainEpisode): """One episode over these traces. ``model_construct`` skips validation so the duck-typed stand-ins above can stand in for real ones.""" - return Episode.model_construct(id=f"e{next(_ids)}", traces=list(rollouts), env_name=env_name, errors=list(errors)) + return cls.model_construct(id=f"e{next(_ids)}", traces=list(rollouts), env_name=env_name, errors=list(errors)) -def solo(rollouts): +def solo(rollouts, cls=TrainEpisode): """Each rollout as its own single-trace episode — the single-agent shape.""" - return [ep(r) for r in rollouts] + return [ep(r, cls=cls) for r in rollouts] def train_wandb(rollouts, subset: str = "all") -> dict: @@ -272,12 +272,12 @@ def test_train_only_metrics_absent_from_eval(): assert out["train/agg/all/agent/is_filtered/mean"] == 0.5 assert out["train/agg/all/agent/filters/gibberish/mean"] == 0.5 assert "train/agg/all/is_trainable/mean" not in out # pipeline verdicts are per-trace - eval_out = EvalRollouts(solo(rollouts)).metrics.to_wandb(prefix="eval/x", subset="all") + eval_out = EvalRollouts(solo(rollouts, cls=EvalEpisode)).metrics.to_wandb(prefix="eval/x", subset="all") assert not any("is_trainable" in k or "is_filtered" in k or "/filters/" in k for k in eval_out) def test_eval_avg_at_k_and_pass_k(): - binary = EvalRollouts(solo([mk(reward=1.0, group_id="g0"), mk(reward=0.0, group_id="g0")])) + binary = EvalRollouts(solo([mk(reward=1.0, group_id="g0"), mk(reward=0.0, group_id="g0")], cls=EvalEpisode)) eff = binary.effective.metrics.to_wandb(prefix="eval/x", subset="effective") assert eff["eval/x/effective/agent/avg@2"] == 0.5 # mean reward under avg@ (k from the groups) assert "eval/x/effective/avg@2" not in eff # scores are per-agent, never pooled @@ -285,7 +285,7 @@ def test_eval_avg_at_k_and_pass_k(): all_out = binary.metrics.to_wandb(prefix="eval/x", subset="all") assert all_out["eval/x/all/agent/avg@2"] == 0.5 assert not any("pass@" in k or "pass^" in k for k in all_out) # pass@k effective-only - non_binary = EvalRollouts(solo([mk(reward=0.5, group_id="g0"), mk(reward=1.0, group_id="g0")])) + non_binary = EvalRollouts(solo([mk(reward=0.5, group_id="g0"), mk(reward=1.0, group_id="g0")], cls=EvalEpisode)) assert not any("pass@" in k for k in non_binary.effective.metrics.to_wandb(prefix="eval/x", subset="effective")) @@ -329,3 +329,20 @@ def test_traceless_episode_keeps_its_reason(): pool = TrainRollouts([ep(mk()), episode]) assert len(pool) == 1 and len(pool.episodes) == 2 assert len(pool.effective.episodes) == 1 # it survives nothing, so the subset drops it + + +def test_train_and_eval_episodes_carry_only_their_own_facts(): + """The path is the type, not a discriminator field: an eval episode always knows its step, and + only a train episode can go stale — neither can be asked for the other's fact.""" + assert (TrainEpisode.KIND, EvalEpisode.KIND) == ("train", "eval") + train, evaluation = TrainEpisode.model_construct(), EvalEpisode.model_construct() + assert train.off_policy_steps == 0 and "off_policy_steps" not in EvalEpisode.model_fields + assert evaluation.eval_step == 0 and "eval_step" not in TrainEpisode.model_fields + assert train.policy_version == evaluation.policy_version == 0 # both measure a policy + + +def test_training_state_is_train_only(): + """Only a train rollout carries trainer-bound state; an eval trace has no field for it.""" + assert {"samples", "advantages", "is_filtered", "filter_results"} <= set(TrainRollout.model_fields) + assert not {"samples", "advantages", "is_filtered", "filter_results"} & set(Rollout.model_fields) + assert {"env_name", "group_id", "episode_id"} <= set(Rollout.model_fields) # links stay shared diff --git a/tests/unit/utils/test_prime_monitor.py b/tests/unit/utils/test_prime_monitor.py index 798e041885..7be5468949 100644 --- a/tests/unit/utils/test_prime_monitor.py +++ b/tests/unit/utils/test_prime_monitor.py @@ -5,7 +5,7 @@ import pyarrow.parquet as pq import verifiers.v1 as vf -from prime_rl.orchestrator.types import Rollout +from prime_rl.orchestrator.types import TrainRollout from prime_rl.utils.monitor.prime import PrimeMonitor @@ -15,8 +15,8 @@ def _new_monitor() -> PrimeMonitor: return monitor -def _build_rollout(*, example_id: int, reward: float, task: str) -> Rollout: - """Build a v1 ``Rollout`` (message-graph trace). The user node carries the prompt and the +def _build_rollout(*, example_id: int, reward: float, task: str) -> TrainRollout: + """Build a v1 ``TrainRollout`` (message-graph trace). The user node carries the prompt and the assistant node the completion; ``_rollouts_to_parquet_bytes`` reads the conversation off the branches (its ``completion`` column is the last branch's messages, ``trajectory`` is one message list per branch).""" @@ -35,7 +35,7 @@ def _build_rollout(*, example_id: int, reward: float, task: str) -> Rollout: sampled=True, ), ] - rollout = Rollout[vf.TaskData]( + rollout = TrainRollout[vf.TaskData]( task=vf.TraceTask(type="Task", data=vf.TaskData(idx=example_id, prompt=f"prompt-{example_id}")), agent=vf.AgentInfo(config=vf.AgentConfig()), nodes=nodes, @@ -81,7 +81,7 @@ def test_rollouts_to_parquet_bytes_skips_rollouts_without_trajectory(): monitor.run_id = "run-456" rollout_with_branches = _build_rollout(example_id=1, reward=1.0, task="task-a") - rollout_without_branches = Rollout[vf.TaskData]( + rollout_without_branches = TrainRollout[vf.TaskData]( task=vf.TraceTask(type="Task", data=vf.TaskData(idx=2, prompt="missing-trajectory")), agent=vf.AgentInfo(config=vf.AgentConfig()), ) From 5e5ba98214e6e96ab179f18bc349847c63ceb152 Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Mon, 3 Aug 2026 23:27:15 +0000 Subject: [PATCH 26/58] refactor(orchestrator): EvalEpisode.eval_step is just step On the episode whose whole existence is one eval epoch, the qualifier says nothing the type doesn't. The dispatcher's GroupState and InflightRollout keep eval_step, where it is an optional field that both kinds pass through. Co-Authored-By: Claude Fable 5 --- src/prime_rl/orchestrator/dispatcher.py | 2 +- src/prime_rl/orchestrator/eval_sink.py | 6 +++--- src/prime_rl/orchestrator/orchestrator.py | 2 +- src/prime_rl/orchestrator/types.py | 7 ++++--- tests/unit/orchestrator/test_metrics.py | 2 +- 5 files changed, 10 insertions(+), 9 deletions(-) diff --git a/src/prime_rl/orchestrator/dispatcher.py b/src/prime_rl/orchestrator/dispatcher.py index 68f9248edf..80a2f10d24 100644 --- a/src/prime_rl/orchestrator/dispatcher.py +++ b/src/prime_rl/orchestrator/dispatcher.py @@ -568,7 +568,7 @@ async def emit_episode( "policy_version": policy_version, } dispatched: Episode = ( - EvalEpisode.model_construct(**shared, eval_step=eval_step) + EvalEpisode.model_construct(**shared, step=eval_step) if meta.kind == "eval" else TrainEpisode.model_construct(**shared, off_policy_steps=meta.off_policy_steps) ) diff --git a/src/prime_rl/orchestrator/eval_sink.py b/src/prime_rl/orchestrator/eval_sink.py index 7b33c2c453..9c1936a11a 100644 --- a/src/prime_rl/orchestrator/eval_sink.py +++ b/src/prime_rl/orchestrator/eval_sink.py @@ -42,7 +42,7 @@ def add(self, episode: EvalEpisode) -> EvalBatch | None: group_id = episode.group_id for rollout in episode.rollouts: self.process_rollout(rollout) - bkey = (env_name, episode.eval_step) + bkey = (env_name, episode.step) self.pending_groups[group_id].append(episode) self.pending_group_episodes[group_id] += 1 if self.pending_group_episodes[group_id] >= self.group_size_for(env_name): @@ -73,7 +73,7 @@ def batch_progress(self) -> list[tuple[str, int, int, int, int]]: env_name = episodes[0].env_name if self.eval_envs.get(env_name).requires_group_scoring: continue - bkey = (env_name, episodes[0].eval_step) + bkey = (env_name, episodes[0].step) buffered[bkey] = buffered.get(bkey, 0) + self.pending_group_episodes.get(group_id, 0) return [ ( @@ -105,7 +105,7 @@ def process_group(self, group_id: uuid.UUID) -> None: # Read the group's facts off an episode, not a trace: every episode in it may have # produced none (a whole group cancelled off-policy). env_name = finished[0].env_name - eval_step = finished[0].eval_step + eval_step = finished[0].step group = [trace for episode in finished for trace in episode.rollouts] task_idx = group[0].task.data.idx if group else -1 bucket = self.pending_batches[(env_name, eval_step)] diff --git a/src/prime_rl/orchestrator/orchestrator.py b/src/prime_rl/orchestrator/orchestrator.py index 1c820d1c1b..8f43ce93cd 100644 --- a/src/prime_rl/orchestrator/orchestrator.py +++ b/src/prime_rl/orchestrator/orchestrator.py @@ -532,7 +532,7 @@ async def main_loop(self) -> None: # eval rollouts to the step whose eval triggered them. A failed episode has no trace # to write; it is counted by the window's episode-failure metrics instead. is_eval = isinstance(episode, EvalEpisode) - step = episode.eval_step if isinstance(episode, EvalEpisode) else self.progress.step + step = episode.step if isinstance(episode, EvalEpisode) else self.progress.step run: vf.RunInfo = ( vf.EvalRunInfo(id=self.run_id, step=step) if is_eval else vf.TrainRunInfo(id=self.run_id, step=step) ) diff --git a/src/prime_rl/orchestrator/types.py b/src/prime_rl/orchestrator/types.py index 9077946282..86adeb5edb 100644 --- a/src/prime_rl/orchestrator/types.py +++ b/src/prime_rl/orchestrator/types.py @@ -190,12 +190,13 @@ def rollouts(self) -> list[TrainRollout]: class EvalEpisode(Episode): - """An episode collected for one eval epoch. ``eval_step`` is the step whose eval triggered it — - always known, unlike on the train path, so it is not optional here.""" + """An episode collected for one eval epoch. ``step`` is the training step whose eval triggered + it — always known, unlike on the train path, so it is not optional here. It is the source for + the ``run.step`` each of its traces gets stamped with on arrival.""" KIND: ClassVar[RolloutKind] = "eval" - eval_step: int = Field(default=0, exclude=True) + step: int = Field(default=0, exclude=True) @dataclass diff --git a/tests/unit/orchestrator/test_metrics.py b/tests/unit/orchestrator/test_metrics.py index 8c6cd611d3..139f3cf493 100644 --- a/tests/unit/orchestrator/test_metrics.py +++ b/tests/unit/orchestrator/test_metrics.py @@ -337,7 +337,7 @@ def test_train_and_eval_episodes_carry_only_their_own_facts(): assert (TrainEpisode.KIND, EvalEpisode.KIND) == ("train", "eval") train, evaluation = TrainEpisode.model_construct(), EvalEpisode.model_construct() assert train.off_policy_steps == 0 and "off_policy_steps" not in EvalEpisode.model_fields - assert evaluation.eval_step == 0 and "eval_step" not in TrainEpisode.model_fields + assert evaluation.step == 0 and "step" not in TrainEpisode.model_fields assert train.policy_version == evaluation.policy_version == 0 # both measure a policy From 3c03c9dc7fc0c6ed48a2101c6009fe4e006ddb35 Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Mon, 3 Aug 2026 23:50:01 +0000 Subject: [PATCH 27/58] refactor(orchestrator): InflightEpisode stamps the Episode that lands MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The dispatcher's per-task record was called InflightRollout while its own invariant said one permit is one episode. It becomes InflightEpisode, and the dispatch-to-landed transition — reconciling the group's values, picking the train or eval class — moves out of emit_episode into stamp(), the one place where the facts of a dispatch become facts of an episode. Co-Authored-By: Claude Fable 5 --- src/prime_rl/orchestrator/dispatcher.py | 71 ++++++++++--------------- src/prime_rl/orchestrator/types.py | 47 ++++++++++------ tests/unit/orchestrator/test_batch.py | 8 +-- tests/unit/orchestrator/test_metrics.py | 21 +++++++- 4 files changed, 84 insertions(+), 63 deletions(-) diff --git a/src/prime_rl/orchestrator/dispatcher.py b/src/prime_rl/orchestrator/dispatcher.py index 80a2f10d24..9e2f51850b 100644 --- a/src/prime_rl/orchestrator/dispatcher.py +++ b/src/prime_rl/orchestrator/dispatcher.py @@ -41,10 +41,8 @@ from prime_rl.orchestrator.eval_source import EvalSource from prime_rl.orchestrator.train_source import TrainSource from prime_rl.orchestrator.types import ( - Episode, - EvalEpisode, GroupState, - InflightRollout, + InflightEpisode, Policy, RolloutKind, ) @@ -153,7 +151,7 @@ def __init__( AsyncLimiter(tasks_per_minute, time_period=60) if tasks_per_minute else None ) - self.inflight: dict[asyncio.Task, InflightRollout] = {} + self.inflight: dict[asyncio.Task, InflightEpisode] = {} self.groups: dict[uuid.UUID, GroupState] = {} # Bounded so the dispatcher backpressures on a slow sink. One entry per @@ -185,11 +183,11 @@ def _train_pool_for(self, env_name: str) -> tuple[InferencePool, str, bool]: @property def inflight_train_count(self) -> int: - return sum(m.rollout_count for m in self.inflight.values() if m.kind == "train") + return sum(m.episodes_owed for m in self.inflight.values() if m.kind == "train") @property def inflight_eval_count(self) -> int: - return sum(m.rollout_count for m in self.inflight.values() if m.kind == "eval") + return sum(m.episodes_owed for m in self.inflight.values() if m.kind == "eval") @property def available_permits(self) -> int: @@ -199,7 +197,7 @@ def available_permits(self) -> int: def inflight_by_env(self) -> dict[tuple[RolloutKind, str], int]: counts: dict[tuple[RolloutKind, str], int] = defaultdict(int) for meta in self.inflight.values(): - counts[(meta.kind, meta.env_name)] += meta.rollout_count + counts[(meta.kind, meta.env_name)] += meta.episodes_owed return dict(counts) @property @@ -474,12 +472,12 @@ async def schedule_group_rollout(self, group_id: uuid.UUID, group: GroupState) - ) ) - self.inflight[task] = InflightRollout( + self.inflight[task] = InflightEpisode( kind=group.kind, env_name=group.env_name, group_id=group_id, policy_version=group.policy_version_at_start, - rollout_count=permits, + episodes_owed=permits, client_config=client, eval_step=group.eval_step, ) @@ -498,8 +496,8 @@ def release(self, n: int) -> None: async def handle_completed_rollout(self, task: asyncio.Task) -> None: """Emit every dispatched episode exactly once to ``out_q``: a ``run`` - result as one episode, a legacy ``run_group`` result as ``rollout_count`` - single-trace episodes. Task exceptions synthesize ``rollout_count`` + result as one episode, a legacy ``run_group`` result as ``episodes_owed`` + single-trace episodes. Task exceptions synthesize ``episodes_owed`` error-marker episodes so the sink's count-to-``group_size`` finalization still triggers. Cancelled tasks (popped by ``drop_group``) raise ``CancelledError`` and are discarded — ``drop_group`` already emitted @@ -508,7 +506,7 @@ async def handle_completed_rollout(self, task: asyncio.Task) -> None: meta = self.inflight.pop(task, None) if meta is None: return # already handled by drop_group / cancel_inflight_rollouts - self.release(meta.rollout_count) + self.release(meta.episodes_owed) group = self.groups.get(meta.group_id) try: @@ -540,7 +538,7 @@ async def handle_completed_rollout(self, task: asyncio.Task) -> None: async def emit_episode( self, - meta: InflightRollout, + meta: InflightEpisode, group: GroupState | None, episode: vf.WireEpisode, ) -> None: @@ -555,30 +553,17 @@ async def emit_episode( group.emitted += 1 if group.emitted >= group.target_rollouts: self.groups.pop(meta.group_id, None) - if meta.kind == "eval": - assert eval_step is not None, "eval episode missing eval_step" for rollout in episode.traces: rollout.env_name = meta.env_name rollout.group_id = meta.group_id - shared = { - **dict(episode), - "env_name": meta.env_name, - "group_id": meta.group_id, - "policy_version": policy_version, - } - dispatched: Episode = ( - EvalEpisode.model_construct(**shared, step=eval_step) - if meta.kind == "eval" - else TrainEpisode.model_construct(**shared, off_policy_steps=meta.off_policy_steps) - ) - await self.out_q.put(dispatched) + await self.out_q.put(meta.stamp(episode, policy_version=policy_version, eval_step=eval_step)) - async def emit_failed_episodes(self, meta: InflightRollout, group: GroupState | None, error: vf.Error) -> None: + async def emit_failed_episodes(self, meta: InflightEpisode, group: GroupState | None, error: vf.Error) -> None: """Emit one traceless episode per rollout the task owed, so the sink still counts its way to ``group_size``. This is the shape vf already gives a run that produced nothing — the reason on ``errors`` — so prime-rl's own failures need no type of their own.""" - for _ in range(meta.rollout_count): + for _ in range(meta.episodes_owed): self.metrics.record_error(kind=meta.kind, env_name=meta.env_name) await self.emit_episode(meta, group, Episode.model_construct(errors=[error])) @@ -594,20 +579,20 @@ async def drop_group(self, group_id: uuid.UUID) -> int: # the dropped tasks are no longer reachable from ``self.inflight``, # so ``handle_completed_rollout``'s existing None-guard makes the # subsequent async emit phase race-free. - claimed: list[tuple[asyncio.Task, InflightRollout]] = [] + claimed: list[tuple[asyncio.Task, InflightEpisode]] = [] for task, meta in list(self.inflight.items()): if meta.group_id != group_id: continue del self.inflight[task] - self.release(meta.rollout_count) + self.release(meta.episodes_owed) claimed.append((task, meta)) tasks_to_cancel = [task for task, _ in claimed] - inflight_cancelled = sum(meta.rollout_count for _, meta in claimed) - last_meta: InflightRollout | None = claimed[-1][1] if claimed else None + inflight_cancelled = sum(meta.episodes_owed for _, meta in claimed) + last_meta: InflightEpisode | None = claimed[-1][1] if claimed else None cancel = vf.Error(type="Cancelled", message="Off-policy cancel") for _, meta in claimed: - for _ in range(meta.rollout_count): + for _ in range(meta.episodes_owed): await self.emit_episode(meta, group, vf.WireEpisode.model_construct(errors=[cancel])) # For non-group-scoring envs, the group may have rollouts that @@ -619,12 +604,12 @@ async def drop_group(self, group_id: uuid.UUID) -> int: # and us reaching it — synthesize a stand-in from the group state unscheduled_cancelled = 0 if group is not None and group.rollouts_to_schedule > 0: - fallback_meta = last_meta or InflightRollout( + fallback_meta = last_meta or InflightEpisode( kind=group.kind, env_name=group.env_name, group_id=group_id, policy_version=group.policy_version_at_start, - rollout_count=1, + episodes_owed=1, eval_step=group.eval_step, ) unscheduled_cancelled = group.rollouts_to_schedule @@ -634,12 +619,12 @@ async def drop_group(self, group_id: uuid.UUID) -> int: cancelled = inflight_cancelled + unscheduled_cancelled if cancelled > 0: meta_for_log = last_meta or ( - InflightRollout( + InflightEpisode( kind=group.kind, env_name=group.env_name, group_id=group_id, policy_version=group.policy_version_at_start if group else 0, - rollout_count=1, + episodes_owed=1, eval_step=group.eval_step, ) if group is not None @@ -660,8 +645,8 @@ async def cancel_inflight_rollouts(self) -> None: """Cancel all in-flight rollouts. Used on shutdown — doesn't emit markers since the sinks are being torn down anyway.""" for meta in self.inflight.values(): - self.metrics.record_cancellation(kind=meta.kind, env_name=meta.env_name, n=meta.rollout_count) - self.release(meta.rollout_count) + self.metrics.record_cancellation(kind=meta.kind, env_name=meta.env_name, n=meta.episodes_owed) + self.release(meta.episodes_owed) tasks = list(self.inflight.keys()) self.inflight.clear() self.groups.clear() @@ -679,9 +664,9 @@ async def cancel_inflight_train_rollouts(self) -> int: if meta.kind != "train": continue self.inflight.pop(task, None) - self.release(meta.rollout_count) - self.metrics.record_cancellation(kind="train", env_name=meta.env_name, n=meta.rollout_count) - cancelled += meta.rollout_count + self.release(meta.episodes_owed) + self.metrics.record_cancellation(kind="train", env_name=meta.env_name, n=meta.episodes_owed) + cancelled += meta.episodes_owed train_tasks.append(task) train_group_ids.add(meta.group_id) for gid in train_group_ids: diff --git a/src/prime_rl/orchestrator/types.py b/src/prime_rl/orchestrator/types.py index 86adeb5edb..445070eacd 100644 --- a/src/prime_rl/orchestrator/types.py +++ b/src/prime_rl/orchestrator/types.py @@ -38,21 +38,6 @@ class Progress: RolloutKind = Literal["train", "eval"] -@dataclass -class InflightRollout: - """Per-task scheduling state in the dispatcher; one entry per in-flight - ``run`` / ``run_group`` task.""" - - kind: RolloutKind - env_name: str - group_id: uuid.UUID - policy_version: int - rollout_count: int - client_config: vf.ClientConfig | None = None - off_policy_steps: int = 0 - eval_step: int | None = None - - @dataclass class GroupState: """Per-group dispatcher state: what's left to schedule + the pinned @@ -199,6 +184,38 @@ class EvalEpisode(Episode): step: int = Field(default=0, exclude=True) +@dataclass +class InflightEpisode: + """One episode in flight, and the facts of the dispatch that will be stamped onto it when it + lands. The pair is the whole lifecycle: an ``InflightEpisode`` going out, an ``Episode`` coming + back — so nothing downstream has to know how a rollout was scheduled.""" + + kind: RolloutKind + env_name: str + group_id: uuid.UUID + policy_version: int + episodes_owed: int + """How many episodes this dispatch owes the sink — one, except on the legacy group path.""" + client_config: vf.ClientConfig | None = None + off_policy_steps: int = 0 + eval_step: int | None = None + + def stamp(self, episode: vf.WireEpisode, *, policy_version: int, eval_step: int | None) -> Episode: + """Mint the landed episode: the env's own, carrying the dispatch it came from. The group's + values win over this dispatch's when it is still alive, so they are passed in rather than + read off ``self``.""" + common = { + **dict(episode), + "env_name": self.env_name, + "group_id": self.group_id, + "policy_version": policy_version, + } + if self.kind == "eval": + assert eval_step is not None, "eval episode missing its step" + return EvalEpisode.model_construct(**common, step=eval_step) + return TrainEpisode.model_construct(**common, off_policy_steps=self.off_policy_steps) + + @dataclass class TrainBatch: """``rollouts`` is the observation window since the last ship — every rollout of every group diff --git a/tests/unit/orchestrator/test_batch.py b/tests/unit/orchestrator/test_batch.py index 8b4b31f7ed..6bf6a22aaa 100644 --- a/tests/unit/orchestrator/test_batch.py +++ b/tests/unit/orchestrator/test_batch.py @@ -93,18 +93,18 @@ def test_training_sample_requires_env_name(): @pytest.mark.parametrize( - ("rollout_count", "num_train_workers", "expected_batches_per_worker"), [(4, 2, 2), (5, 2, 3), (7, 1, 7), (11, 4, 3)] + ("episodes_owed", "num_train_workers", "expected_batches_per_worker"), [(4, 2, 2), (5, 2, 3), (7, 1, 7), (11, 4, 3)] ) def test_prepare_batch_balances_micro_batches_across_workers( - make_training_example, rollout_count, num_train_workers, expected_batches_per_worker + make_training_example, episodes_owed, num_train_workers, expected_batches_per_worker ): - examples = [make_training_example() for i in range(rollout_count)] + examples = [make_training_example() for i in range(episodes_owed)] batches_per_gpu = prepare_batch( rollouts=examples, seq_len=4, num_train_workers=num_train_workers, - idxs=[0] * rollout_count, + idxs=[0] * episodes_owed, num_loras=1, bin_cost=build_bin_cost(None), ) diff --git a/tests/unit/orchestrator/test_metrics.py b/tests/unit/orchestrator/test_metrics.py index 139f3cf493..ee908154f2 100644 --- a/tests/unit/orchestrator/test_metrics.py +++ b/tests/unit/orchestrator/test_metrics.py @@ -1,12 +1,14 @@ import math +from dataclasses import replace from itertools import count from types import SimpleNamespace +from uuid import uuid4 import pytest import verifiers.v1 as vf from prime_rl.orchestrator.metrics import EvalRollouts, Stat, TrainRollouts, episode_failure_metrics -from prime_rl.orchestrator.types import EvalEpisode, Rollout, TrainEpisode, TrainRollout +from prime_rl.orchestrator.types import EvalEpisode, InflightEpisode, Rollout, TrainEpisode, TrainRollout from prime_rl.orchestrator.utils import compute_pass_metrics _ids = count() @@ -346,3 +348,20 @@ def test_training_state_is_train_only(): assert {"samples", "advantages", "is_filtered", "filter_results"} <= set(TrainRollout.model_fields) assert not {"samples", "advantages", "is_filtered", "filter_results"} & set(Rollout.model_fields) assert {"env_name", "group_id", "episode_id"} <= set(Rollout.model_fields) # links stay shared + + +def test_inflight_episode_stamps_what_lands(): + """The dispatch and the landed episode are one pair: `stamp` is the only place the facts of a + dispatch become facts of an episode, and it mints the class the path calls for.""" + wire = vf.WireEpisode.model_construct(id="e", traces=[]) + inflight = InflightEpisode( + kind="train", env_name="rt", group_id=uuid4(), policy_version=3, episodes_owed=1, off_policy_steps=2 + ) + train = inflight.stamp(wire, policy_version=7, eval_step=None) + assert isinstance(train, TrainEpisode) and train.KIND == "train" + assert (train.env_name, train.policy_version, train.off_policy_steps) == ("rt", 7, 2) # group's version wins + + evaluation = replace(inflight, kind="eval").stamp(wire, policy_version=7, eval_step=12) + assert isinstance(evaluation, EvalEpisode) and evaluation.step == 12 + with pytest.raises(AssertionError): # an eval episode without its step is not representable + replace(inflight, kind="eval").stamp(wire, policy_version=7, eval_step=None) From f875000a7cc203317344ff8c65405582fd094faa Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Mon, 3 Aug 2026 23:55:26 +0000 Subject: [PATCH 28/58] style(orchestrator): short loop names in comprehensions Co-Authored-By: Claude Fable 5 --- src/prime_rl/orchestrator/dispatcher.py | 2 +- src/prime_rl/orchestrator/eval_sink.py | 2 +- src/prime_rl/orchestrator/metrics.py | 33 +++++++++++------------ src/prime_rl/orchestrator/orchestrator.py | 2 +- src/prime_rl/orchestrator/train_sink.py | 2 +- 5 files changed, 20 insertions(+), 21 deletions(-) diff --git a/src/prime_rl/orchestrator/dispatcher.py b/src/prime_rl/orchestrator/dispatcher.py index 9e2f51850b..7ff2ca1630 100644 --- a/src/prime_rl/orchestrator/dispatcher.py +++ b/src/prime_rl/orchestrator/dispatcher.py @@ -519,7 +519,7 @@ async def handle_completed_rollout(self, task: asyncio.Task) -> None: await self.emit_failed_episodes(meta, group, vf.Error(type="TaskFailed", message=repr(exc))) return - for r in (r for episode in episodes for r in episode.traces): + for r in (r for e in episodes for r in e.traces): if not r.has_error and r.num_turns == 0: # Empty trajectory: promote to an explicit error so the sink # treats it like any other failure (``has_error`` reads ``ok``) diff --git a/src/prime_rl/orchestrator/eval_sink.py b/src/prime_rl/orchestrator/eval_sink.py index 9c1936a11a..1a892a1f05 100644 --- a/src/prime_rl/orchestrator/eval_sink.py +++ b/src/prime_rl/orchestrator/eval_sink.py @@ -106,7 +106,7 @@ def process_group(self, group_id: uuid.UUID) -> None: # produced none (a whole group cancelled off-policy). env_name = finished[0].env_name eval_step = finished[0].step - group = [trace for episode in finished for trace in episode.rollouts] + group = [t for e in finished for t in e.rollouts] task_idx = group[0].task.data.idx if group else -1 bucket = self.pending_batches[(env_name, eval_step)] bucket.extend(finished) diff --git a/src/prime_rl/orchestrator/metrics.py b/src/prime_rl/orchestrator/metrics.py index 1f3d1f5043..c2726bcc0a 100644 --- a/src/prime_rl/orchestrator/metrics.py +++ b/src/prime_rl/orchestrator/metrics.py @@ -39,7 +39,7 @@ def episode_failure_metrics(episodes: list[Episode], *, prefix: str) -> dict[str seat's error rate.""" if not episodes: return {} - failed = [episode for episode in episodes if episode.failed] + failed = [e for e in episodes if e.failed] counts: dict[str, int] = {} for episode in failed: error = episode.last_error @@ -286,7 +286,7 @@ class EpisodeMetrics: def __init__(self, episodes: list[Episode]) -> None: self.episodes = episodes - self.rollouts: list[Rollout] = [trace for episode in episodes for trace in episode.traces] + self.rollouts: list[Rollout] = [t for e in episodes for t in e.traces] def by_agent(self) -> dict[str, TraceMetrics]: """Per-agent metric views, merging each episode's own ``vf.Episode.by_agent`` grouping.""" @@ -299,23 +299,23 @@ def by_agent(self) -> dict[str, TraceMetrics]: # Episode-level count metrics, one value per episode — ``vf.Episode``'s own aggregates. @property def num_total_tokens(self) -> Stat: - return Stat([float(episode.num_total_tokens) for episode in self.episodes]) + return Stat([float(e.num_total_tokens) for e in self.episodes]) @property def num_input_tokens(self) -> Stat: - return Stat([float(episode.num_input_tokens) for episode in self.episodes]) + return Stat([float(e.num_input_tokens) for e in self.episodes]) @property def num_output_tokens(self) -> Stat: - return Stat([float(episode.num_output_tokens) for episode in self.episodes]) + return Stat([float(e.num_output_tokens) for e in self.episodes]) @property def num_turns(self) -> Stat: - return Stat([float(episode.num_turns) for episode in self.episodes]) + return Stat([float(e.num_turns) for e in self.episodes]) @property def num_branches(self) -> Stat: - return Stat([float(sum(trace.num_branches for trace in episode.traces)) for episode in self.episodes]) + return Stat([float(sum(t.num_branches for t in e.traces)) for e in self.episodes]) # Boolean rate metrics for the console log lines (0/1 distributions — ``.mean()`` is the # rate); to_wandb emits their per-agent counterparts instead. @@ -409,7 +409,7 @@ def to_wandb(self, *, prefix: str, subset: Subset) -> dict[str, float]: def keep_traces(episode: Episode, keep: Callable[[Rollout], bool]) -> Episode | None: """The episode narrowed to the traces that pass ``keep``, or ``None`` if none do. A subset view stays a list of episodes so the episode-level aggregates keep describing what survived.""" - traces = [trace for trace in episode.traces if keep(trace)] + traces = [t for t in episode.traces if keep(t)] return episode.model_copy(update={"traces": traces}) if traces else None @@ -427,10 +427,10 @@ def append(self, episode: TrainEpisode) -> None: @property def rollouts(self) -> list[TrainRollout]: - return [trace for episode in self.episodes for trace in episode.rollouts] + return [t for e in self.episodes for t in e.rollouts] def __len__(self) -> int: - return sum(len(episode.traces) for episode in self.episodes) + return sum(len(e.traces) for e in self.episodes) def __iter__(self) -> Iterator[TrainRollout]: return iter(self.rollouts) @@ -438,10 +438,9 @@ def __iter__(self) -> Iterator[TrainRollout]: @property def effective(self) -> TrainRollouts: kept = ( - keep_traces(episode, lambda r: not r.has_error and not r.is_filtered and r.agent.trainable) - for episode in self.episodes + keep_traces(e, lambda r: not r.has_error and not r.is_filtered and r.agent.trainable) for e in self.episodes ) - return TrainRollouts([episode for episode in kept if episode is not None]) + return TrainRollouts([e for e in kept if e is not None]) def by_env(self) -> dict[str, TrainRollouts]: grouped: dict[str, list[TrainEpisode]] = {} @@ -466,10 +465,10 @@ def __init__(self, episodes: list[EvalEpisode] | None = None, group_size: int | @property def rollouts(self) -> list[Rollout]: - return [trace for episode in self.episodes for trace in episode.traces] + return [t for e in self.episodes for t in e.traces] def __len__(self) -> int: - return sum(len(episode.traces) for episode in self.episodes) + return sum(len(e.traces) for e in self.episodes) def __iter__(self) -> Iterator[Rollout]: return iter(self.rollouts) @@ -490,8 +489,8 @@ def group_size(self) -> int: @property def effective(self) -> EvalRollouts: - kept = (keep_traces(episode, lambda r: not r.has_error and r.agent.trainable) for episode in self.episodes) - return EvalRollouts([episode for episode in kept if episode is not None], group_size=self.group_size) + kept = (keep_traces(e, lambda r: not r.has_error and r.agent.trainable) for e in self.episodes) + return EvalRollouts([e for e in kept if e is not None], group_size=self.group_size) @property def metrics(self) -> EvalMetrics: diff --git a/src/prime_rl/orchestrator/orchestrator.py b/src/prime_rl/orchestrator/orchestrator.py index 8f43ce93cd..80bb354ff0 100644 --- a/src/prime_rl/orchestrator/orchestrator.py +++ b/src/prime_rl/orchestrator/orchestrator.py @@ -547,7 +547,7 @@ async def main_loop(self) -> None: if not episode.failed: await asyncio.to_thread( save_rollouts, - [rollout.to_record() for rollout in episode.rollouts], + [r.to_record() for r in episode.rollouts], get_trace_path(self.config.output_dir, step, episode.KIND, "all"), ) diff --git a/src/prime_rl/orchestrator/train_sink.py b/src/prime_rl/orchestrator/train_sink.py index 60deeec462..723a63154a 100644 --- a/src/prime_rl/orchestrator/train_sink.py +++ b/src/prime_rl/orchestrator/train_sink.py @@ -181,7 +181,7 @@ async def process_group(self, group_id: uuid.UUID) -> None: # Read the group's facts off an episode, not a trace: every episode in it may have # produced none (a whole group cancelled off-policy). env_name = episodes[0].env_name - group = [trace for episode in episodes for trace in episode.rollouts] + group = [t for e in episodes for t in e.rollouts] # Window membership follows group finalization, not arrival: a rollout # only becomes observable (metrics / persistence) once its whole group # is finalized, so a batch's window never claims rollouts of a group From e4880074c8268a7d39958114c78a70fe3179b3f3 Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Mon, 3 Aug 2026 23:57:05 +0000 Subject: [PATCH 29/58] fix(orchestrator): Episode.failed was really 'nothing came back' MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit An episode can error and still have traces — vf keeps the completed subset and marks its clean siblings failed — so a predicate returning False for it had no business being called failed. It is is_empty, and vf's ok stays the success sentinel. Co-Authored-By: Claude Fable 5 --- src/prime_rl/orchestrator/metrics.py | 6 +++--- src/prime_rl/orchestrator/orchestrator.py | 2 +- src/prime_rl/orchestrator/types.py | 11 +++++++---- tests/unit/orchestrator/test_metrics.py | 12 ++++++++++-- 4 files changed, 21 insertions(+), 10 deletions(-) diff --git a/src/prime_rl/orchestrator/metrics.py b/src/prime_rl/orchestrator/metrics.py index c2726bcc0a..e199e59788 100644 --- a/src/prime_rl/orchestrator/metrics.py +++ b/src/prime_rl/orchestrator/metrics.py @@ -39,13 +39,13 @@ def episode_failure_metrics(episodes: list[Episode], *, prefix: str) -> dict[str seat's error rate.""" if not episodes: return {} - failed = [e for e in episodes if e.failed] + empty = [e for e in episodes if e.is_empty] counts: dict[str, int] = {} - for episode in failed: + for episode in empty: error = episode.last_error counts[error.type if error else "Unknown"] = counts.get(error.type if error else "Unknown", 0) + 1 out = {f"{prefix}/episode_failure/{name}": float(n) for name, n in sorted(counts.items())} - out[f"{prefix}/episode_failure/rate"] = len(failed) / len(episodes) + out[f"{prefix}/episode_failure/rate"] = len(empty) / len(episodes) return out diff --git a/src/prime_rl/orchestrator/orchestrator.py b/src/prime_rl/orchestrator/orchestrator.py index 80bb354ff0..1c623634ed 100644 --- a/src/prime_rl/orchestrator/orchestrator.py +++ b/src/prime_rl/orchestrator/orchestrator.py @@ -544,7 +544,7 @@ async def main_loop(self) -> None: episode_id=rollout.episode_id, policy_version=episode.policy_version, ) - if not episode.failed: + if episode.traces: await asyncio.to_thread( save_rollouts, [r.to_record() for r in episode.rollouts], diff --git a/src/prime_rl/orchestrator/types.py b/src/prime_rl/orchestrator/types.py index 445070eacd..8a32638527 100644 --- a/src/prime_rl/orchestrator/types.py +++ b/src/prime_rl/orchestrator/types.py @@ -133,8 +133,8 @@ class Episode(vf.WireEpisode): An episode that produced no traces is not a special case and needs no stand-in rollout: vf already records why on ``errors`` (its ``run_episode`` puts the exception there and returns the episode with ``ok`` false), and prime-rl's own outcomes — an off-policy cancel, a task that - raised before reaching the env — are minted the same way. So ``failed`` is simply "no traces", - and ``last_error`` says why in one vocabulary for every cause. + raised before reaching the env — are minted the same way. So ``is_empty`` is simply "no + traces", and ``last_error`` says why in one vocabulary for every cause. Train and eval are the two subclasses rather than a discriminator field, so each carries only what its path means: an eval episode has a step it belongs to, a train episode has the policy @@ -156,8 +156,11 @@ def rollouts(self) -> list[Rollout]: return cast(list[Rollout], self.traces) @property - def failed(self) -> bool: - """Whether the episode produced nothing — ``last_error`` carries the reason.""" + def is_empty(self) -> bool: + """Whether nothing came back at all — ``last_error`` then carries the reason. Not the + same as failing: an episode can error and still have traces (vf keeps the completed subset + and marks its clean siblings failed), and that failure is accounted for through those + traces. vf's ``ok`` is the success sentinel; this is only "there is nothing here".""" return not self.traces diff --git a/tests/unit/orchestrator/test_metrics.py b/tests/unit/orchestrator/test_metrics.py index ee908154f2..dc5b1716a2 100644 --- a/tests/unit/orchestrator/test_metrics.py +++ b/tests/unit/orchestrator/test_metrics.py @@ -324,9 +324,9 @@ def test_traceless_episode_keeps_its_reason(): """The failure needs no type of its own: ``failed`` is just "no traces", and the reason is the error vf (or the dispatcher) already put on the episode.""" episode = ep(errors=[vf.Error(type="Cancelled", message="Off-policy cancel")]) - assert episode.failed and episode.rollouts == [] + assert episode.is_empty and episode.rollouts == [] assert episode.last_error is not None and episode.last_error.type == "Cancelled" - assert not ep(mk()).failed + assert not ep(mk()).is_empty # A traceless episode still counts as an episode, but contributes no rollouts. pool = TrainRollouts([ep(mk()), episode]) assert len(pool) == 1 and len(pool.episodes) == 2 @@ -365,3 +365,11 @@ def test_inflight_episode_stamps_what_lands(): assert isinstance(evaluation, EvalEpisode) and evaluation.step == 12 with pytest.raises(AssertionError): # an eval episode without its step is not representable replace(inflight, kind="eval").stamp(wire, policy_version=7, eval_step=None) + + +def test_empty_is_not_the_same_as_failed(): + """An episode can error and still have traces — that failure rides the traces and must not be + counted as an episode that produced nothing.""" + errored = ep(mk(has_error=True), errors=[vf.Error(type="EnvError", message="boom")]) + assert not errored.is_empty and not errored.ok # failed, but something came back + assert episode_failure_metrics([errored], prefix="t") == {"t/episode_failure/rate": 0.0} From ca4e5ff5a5e4ad4bd3cfe4a6058fb43ad138f5b9 Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Tue, 4 Aug 2026 17:22:33 +0000 Subject: [PATCH 30/58] feat(orchestrator)!: traces.jsonl stores one episode per line MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Verifiers writes an episode per line and its read_episodes expects that shape; prime-rl wrote a trace per line, so its output was the legacy form its own dependency reads only by sniffing. Episode.to_record mirrors Trace.to_record, and an episode that produced no traces is now written too — its errors are the record of why nothing came back. Co-Authored-By: Claude Fable 5 --- src/prime_rl/orchestrator/orchestrator.py | 35 +++++++++++------------ src/prime_rl/orchestrator/train_sink.py | 2 +- src/prime_rl/orchestrator/types.py | 9 +++++- src/prime_rl/orchestrator/utils.py | 13 +++++---- 4 files changed, 33 insertions(+), 26 deletions(-) diff --git a/src/prime_rl/orchestrator/orchestrator.py b/src/prime_rl/orchestrator/orchestrator.py index 1c623634ed..18f0c763a0 100644 --- a/src/prime_rl/orchestrator/orchestrator.py +++ b/src/prime_rl/orchestrator/orchestrator.py @@ -68,7 +68,7 @@ from prime_rl.orchestrator.utils import ( get_weight_dir, intercept_vf_logging, - save_rollouts, + save_episodes, set_default_executor, setup_policy_inference_pool, trim_process_memory, @@ -526,11 +526,11 @@ async def main_loop(self) -> None: except asyncio.TimeoutError: continue - # Every completed rollout — errored, filtered, or never batched — lands in the - # ``all`` trace file the moment it arrives, so it survives crashes and drains. - # Train rollouts belong to the batch window currently collecting (``progress.step``), - # eval rollouts to the step whose eval triggered them. A failed episode has no trace - # to write; it is counted by the window's episode-failure metrics instead. + # Every completed episode — errored, filtered, or never batched — lands in the + # ``all`` trace file the moment it arrives, so it survives crashes and drains. One + # episode per line, so an episode that produced no traces still records why. + # Train episodes belong to the batch window currently collecting (``progress.step``), + # eval ones to the step whose eval triggered them. is_eval = isinstance(episode, EvalEpisode) step = episode.step if isinstance(episode, EvalEpisode) else self.progress.step run: vf.RunInfo = ( @@ -544,12 +544,11 @@ async def main_loop(self) -> None: episode_id=rollout.episode_id, policy_version=episode.policy_version, ) - if episode.traces: - await asyncio.to_thread( - save_rollouts, - [r.to_record() for r in episode.rollouts], - get_trace_path(self.config.output_dir, step, episode.KIND, "all"), - ) + await asyncio.to_thread( + save_episodes, + [episode.to_record()], + get_trace_path(self.config.output_dir, step, episode.KIND, "all"), + ) if isinstance(episode, EvalEpisode): assert self.eval_sink is not None # eval rollouts only emitted when eval is configured @@ -566,7 +565,7 @@ async def main_loop(self) -> None: async def finalize_train_batch(self, batch: TrainBatch) -> None: """Ship one ``TrainBatch`` out to the trainer and handle the I/O - side-effects (ckpt, save_rollouts, reference scoring, sender.send, + side-effects (ckpt, save_episodes, reference scoring, sender.send, metrics, heartbeat, progress, eval trigger). The sink has already done all data-transformation work.""" config = self.config @@ -644,8 +643,8 @@ async def finalize_train_batch(self, batch: TrainBatch) -> None: # at ship time; the full arrival window already streamed into ``all`` on arrival. # to_record drops the per-node training tensors — they're for training, not the rollout # record, and can't round-trip json (raw numpy bytes). - records = [r.to_record() for r in effective] - await asyncio.to_thread(save_rollouts, records, get_trace_path(config.output_dir, step, "train", "effective")) + records = [e.to_record() for e in effective.episodes] + await asyncio.to_thread(save_episodes, records, get_trace_path(config.output_dir, step, "train", "effective")) await self.sender.send(TrainingBatch(examples=batch.samples, step=step)) self.progress.step += 1 @@ -852,7 +851,7 @@ def log_train_batch(self, batch: TrainBatch, *, step: int, step_time: float) -> get_logger().success("\n\t\t ".join(lines)) async def finalize_eval_batch(self, batch: EvalBatch) -> None: - """Persist + log one completed eval epoch (save_rollouts, + """Persist + log one completed eval epoch (save_episodes, monitor.log_eval_samples, monitor.log).""" if not batch.rollouts: get_logger().warning(f"Eval @ step={batch.step} env={batch.env_name}: no rollouts returned, skipping log") @@ -862,9 +861,9 @@ async def finalize_eval_batch(self, batch: EvalBatch) -> None: # completion (multiple eval envs share the step file — each epoch appends its cohort # once, and every record carries ``env_name``); the full returned cohort already # streamed into ``all`` on arrival. - records = [r.to_record() for r in batch.rollouts.effective] + records = [e.to_record() for e in batch.rollouts.effective.episodes] await asyncio.to_thread( - save_rollouts, records, get_trace_path(self.config.output_dir, batch.step, "eval", "effective") + save_episodes, records, get_trace_path(self.config.output_dir, batch.step, "eval", "effective") ) self.monitor.log_eval_samples(batch.rollouts, env_name=batch.env_name, step=batch.step) policy_versions = {e.policy_version for e in batch.rollouts.episodes} diff --git a/src/prime_rl/orchestrator/train_sink.py b/src/prime_rl/orchestrator/train_sink.py index 723a63154a..5cc2e13805 100644 --- a/src/prime_rl/orchestrator/train_sink.py +++ b/src/prime_rl/orchestrator/train_sink.py @@ -12,7 +12,7 @@ ``add()`` takes one ``Episode`` and returns ``TrainBatch | None``; group accounting counts episodes, never loose traces. -I/O concerns (ship to trainer, save_rollouts, monitor.log) live on the +I/O concerns (ship to trainer, save_episodes, monitor.log) live on the orchestrator. """ diff --git a/src/prime_rl/orchestrator/types.py b/src/prime_rl/orchestrator/types.py index 8a32638527..93287cfeda 100644 --- a/src/prime_rl/orchestrator/types.py +++ b/src/prime_rl/orchestrator/types.py @@ -4,11 +4,12 @@ import uuid from dataclasses import dataclass -from typing import TYPE_CHECKING, ClassVar, Generic, Literal, Protocol, cast +from typing import TYPE_CHECKING, Any, ClassVar, Generic, Literal, Protocol, cast import verifiers.v1 as vf from pydantic import ConfigDict, Field from verifiers.v1.task import DataT +from verifiers.v1.trace import EXCLUDE_FIELDS from prime_rl.transport import TrainingSample @@ -155,6 +156,12 @@ def rollouts(self) -> list[Rollout]: """The episode's traces, typed as the rollouts prime-rl works with.""" return cast(list[Rollout], self.traces) + def to_record(self) -> dict[str, Any]: + """JSON record without raw tensors — the episode form of ``Trace.to_record``, and the unit + ``traces.jsonl`` stores: one episode per line, matching what verifiers writes and what its + ``read_episodes`` expects.""" + return self.model_dump(mode="json", exclude={"traces": {"__all__": EXCLUDE_FIELDS}}) + @property def is_empty(self) -> bool: """Whether nothing came back at all — ``last_error`` then carries the reason. Not the diff --git a/src/prime_rl/orchestrator/utils.py b/src/prime_rl/orchestrator/utils.py index 374a1070f0..62ca24b872 100644 --- a/src/prime_rl/orchestrator/utils.py +++ b/src/prime_rl/orchestrator/utils.py @@ -49,15 +49,16 @@ async def setup_policy_inference_pool(*, config: OrchestratorConfig, tokenizer): return renderer, inference_pool -def save_rollouts(rollouts: list[dict], path: Path) -> None: - """Append rollouts (Trace record dicts, already JSON-serializable) to a JSONL file. - The trace streams are append-only: ``all`` grows one rollout at a time as they - complete, ``effective`` one batch at a time on finalize.""" +def save_episodes(episodes: list[dict], path: Path) -> None: + """Append episodes (Episode record dicts, already JSON-serializable) to a JSONL file — one + episode per line, the shape verifiers writes and its ``read_episodes`` reads. The streams are + append-only: ``all`` grows one episode at a time as they complete, ``effective`` one batch at a + time on finalize.""" path.parent.mkdir(parents=True, exist_ok=True) opts = orjson.OPT_APPEND_NEWLINE | orjson.OPT_SERIALIZE_NUMPY with open(path, "ab") as f: - for rollout in rollouts: - f.write(orjson.dumps(rollout, default=str, option=opts)) + for episode in episodes: + f.write(orjson.dumps(episode, default=str, option=opts)) def intercept_vf_logging(logger: str = "verifiers", level: str = "DEBUG", prefix: str | None = None): From 1234eda8a7355754977a7b294930693be322758b Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Tue, 4 Aug 2026 17:38:19 +0000 Subject: [PATCH 31/58] chore: bump verifiers to the run-on-episode merge vf#2244 landed, so the run is stamped once on the episode instead of looped onto every trace. Merges main for the ServeConfig rename that came with the same vf range. Co-Authored-By: Claude Fable 5 --- deps/verifiers | 2 +- pyproject.toml | 2 +- src/prime_rl/orchestrator/orchestrator.py | 14 ++++++-------- 3 files changed, 8 insertions(+), 10 deletions(-) diff --git a/deps/verifiers b/deps/verifiers index d30a3f48e5..f14b41cf1a 160000 --- a/deps/verifiers +++ b/deps/verifiers @@ -1 +1 @@ -Subproject commit d30a3f48e5f14b06b3081b2102ec32cc3149b849 +Subproject commit f14b41cf1ad6fe9446d2473aa805a08d2b0b48c0 diff --git a/pyproject.toml b/pyproject.toml index 852fef7ce9..149695db84 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -33,7 +33,7 @@ dependencies = [ "setproctitle>=1.3.0", "uvloop>=0.21.0", "torchtitan", - "verifiers[harbor]>=0.2.2.dev76", + "verifiers[harbor]>=0.2.2.dev81", "renderers", "dion", "tilelang>=0.1.8", diff --git a/src/prime_rl/orchestrator/orchestrator.py b/src/prime_rl/orchestrator/orchestrator.py index 18f0c763a0..ff81b7233e 100644 --- a/src/prime_rl/orchestrator/orchestrator.py +++ b/src/prime_rl/orchestrator/orchestrator.py @@ -536,14 +536,12 @@ async def main_loop(self) -> None: run: vf.RunInfo = ( vf.EvalRunInfo(id=self.run_id, step=step) if is_eval else vf.TrainRunInfo(id=self.run_id, step=step) ) - for rollout in episode.rollouts: - rollout.record_run( - run, - env_name=rollout.env_name, - group_id=str(rollout.group_id), - episode_id=rollout.episode_id, - policy_version=episode.policy_version, - ) + episode.record_run( + run, + env_name=episode.env_name, + group_id=str(episode.group_id), + policy_version=episode.policy_version, + ) await asyncio.to_thread( save_episodes, [episode.to_record()], From 3c3982be981f521296c8e250dd30f554b842eaa5 Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Tue, 4 Aug 2026 17:40:15 +0000 Subject: [PATCH 32/58] refactor(orchestrator): count empty episodes on the dispatcher counters episode_failure/ was a new namespace for something the dispatcher already reports: a cancel lands on dispatcher/cancelled, a task failure on dispatcher/errored. Only an env that returned no traces slipped past both, so record that too and drop the metric. Co-Authored-By: Claude Fable 5 --- src/prime_rl/orchestrator/dispatcher.py | 6 ++++++ src/prime_rl/orchestrator/metrics.py | 20 +----------------- src/prime_rl/orchestrator/orchestrator.py | 5 ----- tests/unit/orchestrator/test_metrics.py | 25 ++--------------------- 4 files changed, 9 insertions(+), 47 deletions(-) diff --git a/src/prime_rl/orchestrator/dispatcher.py b/src/prime_rl/orchestrator/dispatcher.py index 7ff2ca1630..7bb8e6e8b9 100644 --- a/src/prime_rl/orchestrator/dispatcher.py +++ b/src/prime_rl/orchestrator/dispatcher.py @@ -534,6 +534,12 @@ async def handle_completed_rollout(self, task: asyncio.Task) -> None: ) # A ``run`` task answers one episode; a legacy ``run_group`` task one per rollout. for episode in episodes: + if not episode.traces: + # No trace to carry the failure, so the episode's own error is all there is. + error = episode.last_error + detail = f"{error.type}: {error.message}" if error is not None else "no error recorded" + get_logger().warning(f"Env returned no traces in group {meta.group_id} ({meta.env_name}) — {detail}") + self.metrics.record_error(kind=meta.kind, env_name=meta.env_name) await self.emit_episode(meta, group, episode) async def emit_episode( diff --git a/src/prime_rl/orchestrator/metrics.py b/src/prime_rl/orchestrator/metrics.py index e199e59788..a2d58ef7b9 100644 --- a/src/prime_rl/orchestrator/metrics.py +++ b/src/prime_rl/orchestrator/metrics.py @@ -26,29 +26,11 @@ from prime_rl.orchestrator.utils import compute_pass_metrics if TYPE_CHECKING: - from prime_rl.orchestrator.types import Episode, EvalEpisode, Rollout, TrainEpisode, TrainRollout + from prime_rl.orchestrator.types import EvalEpisode, Rollout, TrainEpisode, TrainRollout Subset = Literal["all", "effective"] -def episode_failure_metrics(episodes: list[Episode], *, prefix: str) -> dict[str, float]: - """Episodes that produced no traces, by reason (``{prefix}/episode_failure/``) plus their - share of the window. A cancellation, a task that never reached the env, an env that ran no - agent — all of them are an episode with nothing on it but ``errors``, so one counter covers - every cause. They belong to no agent, so they are counted whole rather than folded into a - seat's error rate.""" - if not episodes: - return {} - empty = [e for e in episodes if e.is_empty] - counts: dict[str, int] = {} - for episode in empty: - error = episode.last_error - counts[error.type if error else "Unknown"] = counts.get(error.type if error else "Unknown", 0) + 1 - out = {f"{prefix}/episode_failure/{name}": float(n) for name, n in sorted(counts.items())} - out[f"{prefix}/episode_failure/rate"] = len(empty) / len(episodes) - return out - - class Stat: """A distribution of per-rollout values with mean/max/min and p10/p90 accessors.""" diff --git a/src/prime_rl/orchestrator/orchestrator.py b/src/prime_rl/orchestrator/orchestrator.py index ff81b7233e..d3527b5f86 100644 --- a/src/prime_rl/orchestrator/orchestrator.py +++ b/src/prime_rl/orchestrator/orchestrator.py @@ -49,7 +49,6 @@ from prime_rl.orchestrator.eval_source import EvalSource from prime_rl.orchestrator.filters import setup_filters from prime_rl.orchestrator.inference_metrics import InferenceMetricsCollector -from prime_rl.orchestrator.metrics import episode_failure_metrics from prime_rl.orchestrator.patches import ( monkey_patch_chat_completion_logprobs, monkey_patch_oai_iterable_types, @@ -658,9 +657,6 @@ async def finalize_train_batch(self, batch: TrainBatch) -> None: metrics |= pool.metrics.to_wandb(prefix="train/agg", subset=subset) for env_name, env_pool in pool.by_env().items(): metrics |= env_pool.metrics.to_wandb(prefix=f"train/{env_name}", subset=subset) - # Episodes that produced no trace at all are an episode-level fact — they belong to no - # agent, so they are counted here rather than folded into any seat's error rate. - metrics |= episode_failure_metrics(batch.rollouts.episodes, prefix="train/agg") # Progress / timing / env-share / pre-filter accounting (assembled here, not in the metrics # objects). ``num_tokens`` is over the full arrival window; the input/output breakdown is over @@ -877,7 +873,6 @@ async def finalize_eval_batch(self, batch: EvalBatch) -> None: metrics: dict[str, float] = {} for subset, pool in (("all", rollouts), ("effective", effective)): metrics |= pool.metrics.to_wandb(prefix=f"eval/{batch.env_name}", subset=subset) - metrics |= episode_failure_metrics(rollouts.episodes, prefix=f"eval/{batch.env_name}") metrics[f"eval/{batch.env_name}/policy_version"] = float(policy_version) metrics["step"] = float(batch.step) self.monitor.log(metrics, step=batch.step) diff --git a/tests/unit/orchestrator/test_metrics.py b/tests/unit/orchestrator/test_metrics.py index dc5b1716a2..c427439e5b 100644 --- a/tests/unit/orchestrator/test_metrics.py +++ b/tests/unit/orchestrator/test_metrics.py @@ -7,7 +7,7 @@ import pytest import verifiers.v1 as vf -from prime_rl.orchestrator.metrics import EvalRollouts, Stat, TrainRollouts, episode_failure_metrics +from prime_rl.orchestrator.metrics import EvalRollouts, Stat, TrainRollouts from prime_rl.orchestrator.types import EvalEpisode, InflightEpisode, Rollout, TrainEpisode, TrainRollout from prime_rl.orchestrator.utils import compute_pass_metrics @@ -299,27 +299,6 @@ def test_compute_pass_metrics_matches_closed_form(): assert set(out) == {"pass@1", "pass@2", "pass@4", "pass^1", "pass^2", "pass^4"} -def test_episode_failure_metrics(): - """An episode that produced no traces is counted whole, by the reason vf left on it. Every - cause — a cancel, a task that raised, an env that ran no agent — reads the same way, and none - of them is ever attributed to a seat.""" - cancelled = vf.Error(type="Cancelled", message="Off-policy cancel") - episodes = [ - ep(mk()), - ep(mk()), - ep(errors=[cancelled]), - ep(errors=[cancelled]), - ep(errors=[vf.Error(type="TaskFailed", message="boom")]), - ] - out = episode_failure_metrics(episodes, prefix="train/agg") - assert out["train/agg/episode_failure/Cancelled"] == 2.0 - assert out["train/agg/episode_failure/TaskFailed"] == 1.0 - assert out["train/agg/episode_failure/rate"] == 0.6 # 3 of 5 episodes - assert not any("/agent/" in k for k in out) # never attributed to a seat - assert episode_failure_metrics([ep(mk())], prefix="train/agg") == {"train/agg/episode_failure/rate": 0.0} - assert episode_failure_metrics([], prefix="train/agg") == {} # nothing arrived at all - - def test_traceless_episode_keeps_its_reason(): """The failure needs no type of its own: ``failed`` is just "no traces", and the reason is the error vf (or the dispatcher) already put on the episode.""" @@ -372,4 +351,4 @@ def test_empty_is_not_the_same_as_failed(): counted as an episode that produced nothing.""" errored = ep(mk(has_error=True), errors=[vf.Error(type="EnvError", message="boom")]) assert not errored.is_empty and not errored.ok # failed, but something came back - assert episode_failure_metrics([errored], prefix="t") == {"t/episode_failure/rate": 0.0} + assert errored.rollouts[0].has_error # the failure rides the trace, where the seat's rate sees it From c8913f9884e848b58b5f1a08514f499e55ec7957 Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Tue, 4 Aug 2026 18:25:28 +0000 Subject: [PATCH 33/58] feat(orchestrator)!: credit lives on the graph's nodes assign_advantages writes each node's trainable tokens, so branches sharing a node cannot disagree about its credit, and stamping is a copy of branch.advantages instead of slicing one flat stream back apart by offset. Unassigned stays distinct from assigned-zero all the way to the sample. Co-Authored-By: Claude Fable 5 --- deps/verifiers | 2 +- src/prime_rl/orchestrator/algo/routing.py | 26 +++------- src/prime_rl/orchestrator/types.py | 50 +++++++++---------- tests/unit/orchestrator/test_advantage.py | 38 +++++++++------ tests/unit/orchestrator/test_algorithms.py | 57 +++++++++------------- tests/unit/orchestrator/test_metrics.py | 8 +-- tests/unit/utils/test_prime_monitor.py | 4 +- 7 files changed, 81 insertions(+), 104 deletions(-) diff --git a/deps/verifiers b/deps/verifiers index f14b41cf1a..7eba3e65d2 160000 --- a/deps/verifiers +++ b/deps/verifiers @@ -1 +1 @@ -Subproject commit f14b41cf1ad6fe9446d2473aa805a08d2b0b48c0 +Subproject commit 7eba3e65d2e5df2249519b387f2e7cc6658abf7b diff --git a/src/prime_rl/orchestrator/algo/routing.py b/src/prime_rl/orchestrator/algo/routing.py index c21283e001..3d3192ba63 100644 --- a/src/prime_rl/orchestrator/algo/routing.py +++ b/src/prime_rl/orchestrator/algo/routing.py @@ -14,6 +14,7 @@ from typing import TYPE_CHECKING from prime_rl.configs.algorithm import ActionLossType +from prime_rl.orchestrator.trajectories import iter_trainable_branches from prime_rl.transport import TrainingSample if TYPE_CHECKING: @@ -52,23 +53,8 @@ def stamp_loss_routing(sample: TrainingSample, action_loss_type: ActionLossType) def stamp_advantages(rollout: TrainRollout) -> None: - """Stamp the rollout's per-token advantage stream onto its samples' wire - fields. The stream is full-length-N — aligned to the samples' ``token_ids`` - concatenated in order, 0.0 on non-trainable positions — and sliced across - them. Rollouts with no credit assigned (``advantages=None``, e.g. opd/opsd) - ship no advantage stream. - """ - advantages = rollout.advantages - if advantages is None: - return - total = sum(len(sample.token_ids) for sample in rollout.samples) - if len(advantages) != total: - raise ValueError( - f"advantage stream must align with the rollout's tokens: " - f"got {len(advantages)}, expected {total} (env '{rollout.env_name}')." - ) - offset = 0 - for sample in rollout.samples: - num_tokens = len(sample.token_ids) - sample.advantages = list(advantages[offset : offset + num_tokens]) - offset += num_tokens + """Copy each trainable branch's per-token credit onto the sample built from it. The branch + spreads its nodes' values across its own tokens, so the two align by construction; a rollout + that was never scored (opd/opsd) leaves ``None`` and ships no advantage stream.""" + for sample, (branch, _) in zip(rollout.samples, iter_trainable_branches(rollout), strict=True): + sample.advantages = branch.advantages diff --git a/src/prime_rl/orchestrator/types.py b/src/prime_rl/orchestrator/types.py index 93287cfeda..8fd6a100e3 100644 --- a/src/prime_rl/orchestrator/types.py +++ b/src/prime_rl/orchestrator/types.py @@ -84,46 +84,40 @@ class TrainRollout(Rollout[DataT], Generic[DataT]): model_config = ConfigDict(arbitrary_types_allowed=True) # ``samples`` holds msgspec structs samples: list[TrainingSample] = Field(default_factory=list, exclude=True) - # Per-token rl advantage stream, full-length-N (= len(token_ids)) per - # sample, concatenated across the rollout's samples in order; 0.0 on - # non-trainable positions. None = no credit assigned (advantage-based - # filters skip it; the wire ships no advantage stream). - advantages: list[float] | None = Field(default=None, exclude=True) is_filtered: bool = Field(default=False, exclude=True) filter_results: dict[str, bool] = Field(default_factory=dict, exclude=True) - def assign_advantages(self, values: float | list[float]) -> None: - """Write the rl advantage stream: a scalar broadcast over the - rollout's trainable (mask-True) tokens (0.0 elsewhere), or a per-token - list already aligned full-length to the samples' concatenated - ``token_ids``. A rollout never assigned ships no advantage stream.""" - total = sum(len(sample.token_ids) for sample in self.samples) - if isinstance(values, (int, float)): - self.advantages = [ - float(values) if trainable else 0.0 for sample in self.samples for trainable in sample.mask - ] - return - if len(values) != total: - raise ValueError( - f"per-token advantages must align with the rollout's tokens: " - f"got {len(values)}, expected {total} (env '{self.env_name}')." - ) - self.advantages = [float(v) for v in values] + def assign_advantages(self, value: float) -> None: + """Write ``value`` as the credit for every trainable token, node by node. Credit lives on + the nodes so branches sharing one cannot disagree about it, and so the trainer reads it + aligned to the tokens (``Branch.advantages``) rather than re-sliced by offset.""" + for node in self.nodes: + trainable = sum(node.mask) + node.advantages = [value] * trainable if trainable else None + + @property + def advantages(self) -> list[float] | None: + """Every assigned credit on this trace, or ``None`` if it was never scored — which a + trace assigned all zeros is not.""" + if all(node.advantages is None for node in self.nodes): + return None + return [a for node in self.nodes for a in (node.advantages or [])] def scalar_advantage(self) -> float | None: - """Scalar view of the per-token advantage stream for monitoring: the - mean over assigned (non-zero) positions — exact for the uniform GRPO - case, 0.0 for a zero-advantage group, None when no credit was assigned.""" - if not self.advantages: + """Scalar view of the credit for monitoring: the mean over assigned (non-zero) positions — + exact for the uniform GRPO case, 0.0 for a zero-advantage group, None when unscored.""" + advantages = self.advantages + if not advantages: return None - nonzero = [a for a in self.advantages if a != 0.0] + nonzero = [a for a in advantages if a != 0.0] return sum(nonzero) / len(nonzero) if nonzero else 0.0 @property def is_trainable(self) -> bool: """Whether the rollout carries a training signal — a nonzero advantage on some token. A uniform-reward GRPO group (all-zero advantages) or an unscored rollout has no gradient.""" - return bool(self.advantages) and any(a != 0.0 for a in self.advantages) + advantages = self.advantages + return bool(advantages) and any(a != 0.0 for a in advantages) class Episode(vf.WireEpisode): diff --git a/tests/unit/orchestrator/test_advantage.py b/tests/unit/orchestrator/test_advantage.py index d2ae909afe..a2bcb1adfb 100644 --- a/tests/unit/orchestrator/test_advantage.py +++ b/tests/unit/orchestrator/test_advantage.py @@ -10,7 +10,7 @@ ) from prime_rl.orchestrator.algo.grpo import GRPOAlgorithm from prime_rl.orchestrator.algo.max_rl import MaxRLAlgorithm -from prime_rl.orchestrator.trajectories import trace_to_samples +from prime_rl.orchestrator.trajectories import iter_trainable_branches, trace_to_samples from prime_rl.orchestrator.types import TrainRollout @@ -139,10 +139,10 @@ def _make_group(rewards, completion_lengths=None, num_turns=None) -> list[TrainR def _scalar(rollout: TrainRollout) -> float: - """The per-rollout advantage scalar an algorithm assigned — broadcast over - the rollout's trainable (mask-True) tokens, so any trainable position holds it.""" - mask = [m for sample in rollout.samples for m in sample.mask] - return rollout.advantages[mask.index(True)] + """The per-rollout advantage scalar an algorithm assigned — broadcast over the rollout's + trainable tokens, so every assigned position holds it.""" + assert rollout.advantages is not None + return rollout.advantages[0] def _grpo(group: list[TrainRollout], length_penalty=None) -> list[float]: @@ -240,23 +240,29 @@ def test_linear_turns_term_penalizes_more_turns(): def test_assign_advantages_broadcasts_scalar(): - """A scalar broadcasts uniformly over the rollout's trainable (mask-True) tokens.""" - rollout = _build_rollout(0.0, sampled_lengths=[2]) - # one user prompt token (masked) + 2 sampled tokens (trainable) + """A scalar broadcasts over the trainable tokens, and only those — the credit is stored per + node against `mask`, so untrainable positions hold nothing rather than a zero.""" + rollout = _build_rollout(0.0, sampled_lengths=[2]) # 1 masked prompt token + 2 trainable rollout.assign_advantages(0.7) - assert rollout.advantages == [0.0, 0.7, 0.7] + assert rollout.advantages == [0.7, 0.7] + # The branch view widens it back to the token sequence the trainer indexes. + (branch, _), *_ = iter_trainable_branches(rollout) + assert branch.advantages == [0.0, 0.7, 0.7] def test_assign_advantages_zeros_non_trainable(): - """Non-trainable (mask=False) positions stay 0.0 under scalar broadcast.""" + """A non-trainable (mask=False) position reads 0.0 in the branch view.""" # prompt(1, masked) + sampled(1) + obs(1, masked): mask is [F, T, F] rollout = _build_rollout(0.0, sampled_lengths=[1], obs_lengths=[1]) rollout.assign_advantages(0.7) - assert rollout.advantages == [0.0, 0.7, 0.0] + (branch, _), *_ = iter_trainable_branches(rollout) + assert branch.advantages == [0.0, 0.7, 0.0] -def test_assign_advantages_rejects_misaligned(): - rollout = _build_rollout(0.0, sampled_lengths=[2]) - # full length is 3 (prompt + 2 sampled); a 1-element list must be rejected - with pytest.raises(ValueError, match="align"): - rollout.assign_advantages([0.5]) +def test_unassigned_credit_is_not_zero_credit(): + """An unscored rollout is distinguishable from one scored zero, all the way to the sample.""" + unscored = _build_rollout(0.0, sampled_lengths=[2]) + assert unscored.advantages is None and not unscored.is_trainable + scored = _build_rollout(0.0, sampled_lengths=[2]) + scored.assign_advantages(0.0) + assert scored.advantages == [0.0, 0.0] and not scored.is_trainable # scored, but no gradient diff --git a/tests/unit/orchestrator/test_algorithms.py b/tests/unit/orchestrator/test_algorithms.py index e17d8bd158..2e219360b0 100644 --- a/tests/unit/orchestrator/test_algorithms.py +++ b/tests/unit/orchestrator/test_algorithms.py @@ -9,7 +9,7 @@ from prime_rl.configs.algorithm import AlgoConfig, FrozenModelConfig from prime_rl.orchestrator.algo import EchoAlgorithm, stamp_advantages, stamp_loss_routing -from prime_rl.orchestrator.trajectories import trace_to_samples +from prime_rl.orchestrator.trajectories import iter_trainable_branches, trace_to_samples from prime_rl.orchestrator.types import TrainRollout from prime_rl.transport.types import TrainingSample @@ -157,36 +157,39 @@ def test_stamp_loss_routing_merges_action_weights_into_ce_stream(): assert sample.ref_kl_weights is None -def _make_rollout( - samples: list[TrainingSample], - advantages: list[float] | None = None, -) -> TrainRollout: +def _make_rollout(samples: list[TrainingSample]) -> TrainRollout: + """A rollout whose graph mirrors ``_make_sample``: one node per sample, carrying that + sample's tokens and mask, so branches and samples line up as they do in a real trace.""" + nodes = [ + MessageNode( + parent=None, + message=AssistantMessage(role="assistant", content=""), + sampled=True, + token_ids=list(sample.token_ids), + mask=list(sample.mask), + logprobs=[lp for lp, m in zip(sample.logprobs, sample.mask, strict=True) if m], + ) + for sample in samples + ] rollout = TrainRollout( task=vf.TraceTask(type="Task", data=vf.TaskData(idx=0, prompt=None)), agent=vf.AgentInfo(config=vf.AgentConfig()), - nodes=[], + nodes=nodes, rewards={}, env_name="test-env", ) rollout.samples = samples - rollout.advantages = advantages return rollout -def test_stamp_advantages_full_length_stream(): - # The advantage stream is full-length-N: 0.0 on prompt + non-trainable - # positions, the rl credit on trainable (mask True) tokens. - rollout = _make_rollout([_make_sample()], advantages=[0.0, 0.0, 0.5, -0.5, 0.0, 1.0]) - stamp_advantages(rollout) - assert rollout.samples[0].advantages == [0.0, 0.0, 0.5, -0.5, 0.0, 1.0] - - -def test_stamp_advantages_slices_across_samples(): - samples = [_make_sample(), _make_sample()] - rollout = _make_rollout(samples, advantages=[1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0, 12.0]) +def test_stamp_advantages_copies_the_branch_view(): + """Each sample takes its branch's spread credit — no offsets, so the two align by + construction. Covered end to end in test_advantage.py, which builds real nodes.""" + rollout = _make_rollout([_make_sample()]) + rollout.assign_advantages(1.0) stamp_advantages(rollout) - assert rollout.samples[0].advantages == [1.0, 2.0, 3.0, 4.0, 5.0, 6.0] - assert rollout.samples[1].advantages == [7.0, 8.0, 9.0, 10.0, 11.0, 12.0] + (branch, _), *_ = iter_trainable_branches(rollout) + assert rollout.samples[0].advantages == branch.advantages def test_stamp_advantages_no_credit_ships_none(): @@ -195,22 +198,10 @@ def test_stamp_advantages_no_credit_ships_none(): assert rollout.samples[0].advantages is None -def test_stamp_advantages_rejects_misaligned(): - rollout = _make_rollout([_make_sample()], advantages=[0.5]) - with pytest.raises(ValueError, match="align"): - stamp_advantages(rollout) - - def test_assign_advantages_scalar_broadcasts_over_mask(): rollout = _make_rollout([_make_sample()]) rollout.assign_advantages(1.0) - assert rollout.advantages == [0.0, 0.0, 1.0, 1.0, 0.0, 1.0] - - -def test_assign_advantages_list_rejects_misaligned(): - rollout = _make_rollout([_make_sample()]) - with pytest.raises(ValueError, match="align"): - rollout.assign_advantages([0.5]) + assert rollout.advantages == [1.0, 1.0, 1.0] # one per trainable token # -------------------------------------------------------------------------- diff --git a/tests/unit/orchestrator/test_metrics.py b/tests/unit/orchestrator/test_metrics.py index c427439e5b..457e205cbd 100644 --- a/tests/unit/orchestrator/test_metrics.py +++ b/tests/unit/orchestrator/test_metrics.py @@ -323,9 +323,11 @@ def test_train_and_eval_episodes_carry_only_their_own_facts(): def test_training_state_is_train_only(): - """Only a train rollout carries trainer-bound state; an eval trace has no field for it.""" - assert {"samples", "advantages", "is_filtered", "filter_results"} <= set(TrainRollout.model_fields) - assert not {"samples", "advantages", "is_filtered", "filter_results"} & set(Rollout.model_fields) + """Only a train rollout carries trainer-bound state; an eval trace has no field for it. + Credit is not among them — it lives on the graph's nodes, which every trace has.""" + assert {"samples", "is_filtered", "filter_results"} <= set(TrainRollout.model_fields) + assert not {"samples", "is_filtered", "filter_results"} & set(Rollout.model_fields) + assert "advantages" not in TrainRollout.model_fields # derived from the nodes assert {"env_name", "group_id", "episode_id"} <= set(Rollout.model_fields) # links stay shared diff --git a/tests/unit/utils/test_prime_monitor.py b/tests/unit/utils/test_prime_monitor.py index 7be5468949..799493f549 100644 --- a/tests/unit/utils/test_prime_monitor.py +++ b/tests/unit/utils/test_prime_monitor.py @@ -42,9 +42,7 @@ def _build_rollout(*, example_id: int, reward: float, task: str) -> TrainRollout rewards={"reward": vf.Reward(score=reward)}, ) rollout.env_name = task - # Per-token advantage stream (full-length-N): 0.0 on the 3 prompt tokens, - # reward/2 on the 2 completion (mask-True) tokens. - rollout.advantages = [0.0, 0.0, 0.0, reward / 2, reward / 2] + rollout.assign_advantages(reward / 2) # over the 2 completion (mask-True) tokens return rollout From 1ada1077b08dd9c72d440b7cdcf9dca55287c59c Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Tue, 4 Aug 2026 18:33:36 +0000 Subject: [PATCH 34/58] feat(orchestrator)!: algorithms score episodes, not loose traces MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit score_group/finalize_group take the group's TrainEpisodes, so an algorithm can compare within an episode as well as across them — hierarchical GRPO keys its solver baselines off episode.id instead of a foreign key copied onto every trace. group_rollouts() flattens the cohort for the algorithms that only ever compare across. The sink narrows each episode to its trainable survivors (Episode.narrow, moved off metrics.py) rather than flattening the group. Co-Authored-By: Claude Fable 5 --- docs/algorithms.md | 16 +++++----- src/prime_rl/orchestrator/algo/__init__.py | 7 ++--- src/prime_rl/orchestrator/algo/base.py | 30 +++++++++++-------- src/prime_rl/orchestrator/algo/grpo.py | 16 +++++----- .../orchestrator/algo/hierarchical_grpo.py | 16 +++++----- src/prime_rl/orchestrator/algo/max_rl.py | 10 ++++--- src/prime_rl/orchestrator/algo/rae.py | 7 +++-- src/prime_rl/orchestrator/metrics.py | 13 ++------ src/prime_rl/orchestrator/train_sink.py | 19 ++++++------ src/prime_rl/orchestrator/types.py | 16 +++++++++- tests/unit/orchestrator/test_advantage.py | 14 ++++++--- 11 files changed, 93 insertions(+), 71 deletions(-) diff --git a/docs/algorithms.md b/docs/algorithms.md index de5b3594ed..1392c3538c 100644 --- a/docs/algorithms.md +++ b/docs/algorithms.md @@ -166,12 +166,12 @@ At runtime, each env's resolved config builds two objects: a `Sampler` (`prime_r | `opsd` | `OPSDAlgorithm` | `score_rollout`: demo-conditioned prefill under the live policy | | `sft` | `SFTDistillAlgorithm` | `score_group`: group-norm credit (feeds filters) | -Each class owns its hooks outright — reading one top to bottom reads the algorithm, and everything on the class is an override point. The two hooks are one scope-and-timing ladder — the wider scope is unlocked by a later barrier, so the two axes coincide. Each is handed the `Rollout` directly — the env's typed trace (`reward`, `nodes`, `num_turns`, ...) with `samples` attached, plus `assign_advantages` to write credit: +Each class owns its hooks outright — reading one top to bottom reads the algorithm, and everything on the class is an override point. The two hooks are one scope-and-timing ladder — the wider scope is unlocked by a later barrier, so the two axes coincide. Each is handed the env's own data directly — a `Rollout` (the typed trace: `reward`, `nodes`, `num_turns`, ... with `samples` attached) on arrival, the group's `Episode`s at group time — plus `assign_advantages` to write credit: - `async score_rollout(rollout)` — one rollout, **on arrival** (as it's tokenized, before its group is complete): rollout-local credit (`rollout.assign_advantages(...)`, scalar broadcast or per-token), observation ce weights, **or** model I/O — query a reference pool (e.g. `self.teacher_pool`, connected in `setup()` via `self.connect(...)`, or the live `self.policy_pool` for opsd) and attach per-token results (e.g. teacher logprobs) with bounded concurrency. No siblings. `echo` weights observation tokens here, identifying env-provided observation nodes by their non-sampled status and source step role attribution, applying the optional user filter, and writing the `ce_weights` stream. Model I/O runs *before* the pre-batch filters, so it pays compute on rollouts that may then be filtered out. -- `score_group(group)` — the cohort, **before filtering** (filters read the streams), synchronous: group-relative credit (GRPO/MaxRL baselines). `group` is a list of `Rollout`. +- `score_group(group)` — the cohort, **before filtering** (filters read the streams), synchronous: group-relative credit (GRPO/MaxRL baselines). `group` is a list of `TrainEpisode`, so an algorithm can compare within an episode as well as across them (`hierarchical_grpo` does); `group_rollouts(group)` flattens it for the ones that only compare across. -The pipeline drives the hooks through two non-virtual methods it never looks inside: `algorithm.finalize_rollout(rollout)` per arrival (rollout-local scoring + reference I/O) and `algorithm.finalize_group(rollouts)` per group (scoring + wire stamping; after this the records are frozen — groups die at stamping). Sample construction (interleaving) is pure pipeline — observation-token provenance is available through structural attribution (`node.sampled`, `node.is_content`) for any algorithm that trains on env-provided tokens. +The pipeline drives the hooks through two non-virtual methods it never looks inside: `algorithm.finalize_rollout(rollout)` per arrival (rollout-local scoring + reference I/O) and `algorithm.finalize_group(episodes)` per group (scoring + wire stamping; after this the records are frozen — groups die at stamping). Sample construction (interleaving) is pure pipeline — observation-token provenance is available through structural attribution (`node.sampled`, `node.is_content`) for any algorithm that trains on env-provided tokens. Class-level declarations state what the algorithm needs: which loss component its action tokens feed (`action_loss_type`). Every class is constructed with its algorithm config plus the one host-owned resource it can't rebuild — the live policy pool (`self.policy_pool`). Everything else an algorithm needs it builds from its own config in `setup()`: `opd` connects its frozen `teacher`; `opsd` builds the renderer for its demonstration hint (tokenizer is always the live policy's — self-distillation has no separate model). The pipeline only ever calls the two `finalize_*` methods — writing your own algorithm is subclassing `Algorithm` and overriding the hooks its signal needs (see [Authoring an Algorithm](#authoring-an-algorithm)). Shared math (efficiency shaping, prefill alignment) lives as plain functions in `prime_rl.orchestrator.algo.advantage`. @@ -427,19 +427,21 @@ There is no config hook that points at user code — a new credit-assignment sch import torch from prime_rl.orchestrator.algo.base import Algorithm +from prime_rl.orchestrator.types import group_rollouts class MyAlgorithm(Algorithm): async def score_group(self, group): - rewards = torch.tensor([rollout.reward for rollout in group], dtype=torch.float32) + rollouts = group_rollouts(group) + rewards = torch.tensor([rollout.reward for rollout in rollouts], dtype=torch.float32) advantages = ... # one value per rollout - for rollout, advantage in zip(group, advantages.tolist(), strict=True): + for rollout, advantage in zip(rollouts, advantages.tolist(), strict=True): rollout.assign_advantages(advantage) ``` -Add a typed `MyAlgoConfig` to `prime_rl.configs.algorithm` and its discriminated union, then register `"my_algo": MyAlgorithm` in `ALGORITHM_CLASSES`. Pick the hook by *when* your signal is ready: `score_rollout` for per-arrival credit or credit that needs a model call (it's `async`), `score_group` for group-relative credit (GRPO/MaxRL). `assign_advantages` takes a scalar (broadcast over the rollout's trainable tokens — the common case) or a full-length per-token list aligned to the concatenated sample token_ids (process rewards, step-level credit; `0.0` off-mask). +Add a typed `MyAlgoConfig` to `prime_rl.configs.algorithm` and its discriminated union, then register `"my_algo": MyAlgorithm` in `ALGORITHM_CLASSES`. Pick the hook by *when* your signal is ready: `score_rollout` for per-arrival credit or credit that needs a model call (it's `async`), `score_group` for group-relative credit (GRPO/MaxRL). `assign_advantages` takes a scalar and broadcasts it over the rollout's trainable tokens, writing it onto the graph's nodes so branches sharing a node cannot disagree about its credit. -Each per-token list must match the rollout's completion-token count exactly — validated loudly when the view writes it. Advantage-based filters and metrics derive from the streams (the zero-advantage filter checks for all-zero streams; logged distributions use per-rollout means). Signals that depend on the live policy's weights (like OPD's reverse KL) cannot be precomputed here; those are reference-scoring algorithms, evaluated in the trainer. +Advantage-based filters and metrics derive from the streams (the zero-advantage filter checks for all-zero streams; logged distributions use per-rollout means). Signals that depend on the live policy's weights (like OPD's reverse KL) cannot be precomputed here; those are reference-scoring algorithms, evaluated in the trainer. ### Reference Scoring diff --git a/src/prime_rl/orchestrator/algo/__init__.py b/src/prime_rl/orchestrator/algo/__init__.py index 24b6be1fe9..5e5f657696 100644 --- a/src/prime_rl/orchestrator/algo/__init__.py +++ b/src/prime_rl/orchestrator/algo/__init__.py @@ -17,10 +17,9 @@ ``finalize_rollout`` / ``finalize_group`` methods the pipeline drives. Advantages are per-token everywhere they are stored or shipped — there is no scalar advantage in the pipeline. An algorithm assigns credit in its scoring - hook via ``TrainRollout.assign_advantages``: a scalar that is *broadcast* over the - rollout's completion tokens (uniform credit, the common case), or an explicit - full-length-N per-token list aligned to the concatenated sample token_ids - (0.0 off-mask). + hook via ``TrainRollout.assign_advantages``, which broadcasts one value over + the rollout's trainable tokens and writes it onto the graph's nodes, where + the tokens themselves live. - ``routing`` — wire-field stamping: per-token component weight streams (rl / ce / ref_kl) and the per-token advantage stream. """ diff --git a/src/prime_rl/orchestrator/algo/base.py b/src/prime_rl/orchestrator/algo/base.py index 11c63ce724..410ff4b858 100644 --- a/src/prime_rl/orchestrator/algo/base.py +++ b/src/prime_rl/orchestrator/algo/base.py @@ -19,7 +19,7 @@ I/O against another model — an inference pool the algorithm connected in ``setup()`` (a frozen teacher) or the live policy (opsd's self-distillation), queried with bounded concurrency. No siblings. -- ``score_group(group)`` — the cohort, on group completion, *before* filtering +- ``score_group(group)`` — the cohort of episodes, on group completion, *before* filtering (filters read the streams): group-relative credit (GRPO/MaxRL baselines). How rollouts are *produced* is not the algorithm's concern: that is the env's @@ -44,12 +44,13 @@ from prime_rl.configs.algorithm import ActionLossType, AlgoConfig, FrozenModelConfig from prime_rl.orchestrator.algo.routing import stamp_advantages, stamp_loss_routing +from prime_rl.orchestrator.types import group_rollouts from prime_rl.utils.logger import get_logger if TYPE_CHECKING: from renderers import RendererConfig - from prime_rl.orchestrator.types import TrainRollout + from prime_rl.orchestrator.types import TrainEpisode, TrainRollout from prime_rl.utils.client import InferencePool @@ -91,9 +92,10 @@ class Algorithm: (``action_loss_type``); - lifecycle — :meth:`setup` connects client pools to the frozen models the algorithm declares, resolving each reference via :meth:`connect`; - - the two scoring hooks, each ``async`` and given the :class:`TrainRollout` - directly — read the trace, write credit via - :meth:`TrainRollout.assign_advantages`. They are + - the two scoring hooks, each ``async`` and given the env's own data + directly — a :class:`TrainRollout` on arrival, the group's + :class:`TrainEpisode`\ s at group time — so a hook reads the trace and + writes credit via :meth:`TrainRollout.assign_advantages`. They are async so either stage may do I/O — e.g. a process-reward model or a teacher at arrival, or a judge at group time whose signal a pre-batch filter then reads; a hook that only does advantage math simply never @@ -103,9 +105,9 @@ class Algorithm: observation ce weights, or per-token results from a model the algorithm connected in :meth:`setup` (e.g. teacher reference logprobs). Default: nothing. - - :meth:`score_group` — the cohort, *before* filtering (filters read the - streams): group-relative credit. Default: nothing — rollouts keep - ``advantages=None``, so advantage-based filters skip them. + - :meth:`score_group` — the cohort of episodes, *before* filtering + (filters read the streams): group-relative credit. Default: nothing — + rollouts keep ``advantages=None``, so advantage-based filters skip them. Model I/O lives in :meth:`score_rollout`: it runs at arrival, *before* the pre-batch filters, so it pays compute on rollouts that may then be filtered @@ -144,9 +146,11 @@ async def score_rollout(self, rollout: TrainRollout) -> None: connected in :meth:`setup`, or the live policy (opsd). No siblings, no group stats.""" - async def score_group(self, group: list[TrainRollout]) -> None: + async def score_group(self, group: list[TrainEpisode]) -> None: """Group phase, the finalized cohort, before filtering: write - group-relative credit.""" + group-relative credit. The cohort arrives as episodes, so an algorithm + can compare within one episode as well as across them; ``group_rollouts`` + flattens it for the algorithms that only compare across.""" async def finalize_rollout(self, rollout: TrainRollout) -> None: """Arrival phase (non-virtual): rollout-local scoring as each rollout is @@ -154,12 +158,12 @@ async def finalize_rollout(self, rollout: TrainRollout) -> None: if rollout.samples: await self.score_rollout(rollout) - async def finalize_group(self, rollouts: list[TrainRollout]) -> None: + async def finalize_group(self, episodes: list[TrainEpisode]) -> None: """Group phase (non-virtual): group-relative scoring, then stamp each sample's wire fields (the advantage stream + loss routing). After this the records are frozen — groups die at stamping.""" - await self.score_group(rollouts) - for rollout in rollouts: + await self.score_group(episodes) + for rollout in group_rollouts(episodes): stamp_advantages(rollout) for sample in rollout.samples: stamp_loss_routing(sample, self.action_loss_type) diff --git a/src/prime_rl/orchestrator/algo/grpo.py b/src/prime_rl/orchestrator/algo/grpo.py index d538fda9bd..c002858bb2 100644 --- a/src/prime_rl/orchestrator/algo/grpo.py +++ b/src/prime_rl/orchestrator/algo/grpo.py @@ -6,9 +6,10 @@ from prime_rl.configs.algorithm import GRPOAlgoConfig from prime_rl.orchestrator.algo.base import Algorithm +from prime_rl.orchestrator.types import group_rollouts if TYPE_CHECKING: - from prime_rl.orchestrator.types import TrainRollout + from prime_rl.orchestrator.types import TrainEpisode from prime_rl.utils.client import InferencePool @@ -21,15 +22,16 @@ def __init__(self, config: GRPOAlgoConfig, policy_pool: InferencePool): super().__init__(config, policy_pool) self.length_penalty = config.length_penalty - async def score_group(self, group: list[TrainRollout]) -> None: - rewards = torch.tensor([rollout.reward for rollout in group], dtype=torch.float32) + async def score_group(self, group: list[TrainEpisode]) -> None: + rollouts = group_rollouts(group) + rewards = torch.tensor([rollout.reward for rollout in rollouts], dtype=torch.float32) length_penalty = self.length_penalty if length_penalty is None: advantages = rewards - rewards.mean() else: - output = torch.tensor([rollout.num_output_tokens for rollout in group], dtype=rewards.dtype) - total = torch.tensor([rollout.num_total_tokens for rollout in group], dtype=rewards.dtype) - turns = torch.tensor([rollout.num_turns for rollout in group], dtype=rewards.dtype) + output = torch.tensor([rollout.num_output_tokens for rollout in rollouts], dtype=rewards.dtype) + total = torch.tensor([rollout.num_total_tokens for rollout in rollouts], dtype=rewards.dtype) + turns = torch.tensor([rollout.num_turns for rollout in rollouts], dtype=rewards.dtype) input = total - output penalty_frac = ( length_penalty.num_output_tokens_weight * (output / output.max().clamp(min=1)) @@ -39,5 +41,5 @@ async def score_group(self, group: list[TrainRollout]) -> None: penalty = rewards.mean() * penalty_frac shaped_rewards = rewards - penalty advantages = shaped_rewards - shaped_rewards.mean() - for rollout, advantage in zip(group, advantages.tolist(), strict=True): + for rollout, advantage in zip(rollouts, advantages.tolist(), strict=True): rollout.assign_advantages(advantage) diff --git a/src/prime_rl/orchestrator/algo/hierarchical_grpo.py b/src/prime_rl/orchestrator/algo/hierarchical_grpo.py index a1b4912ca7..f97bdcae5f 100644 --- a/src/prime_rl/orchestrator/algo/hierarchical_grpo.py +++ b/src/prime_rl/orchestrator/algo/hierarchical_grpo.py @@ -7,7 +7,7 @@ from prime_rl.orchestrator.algo.base import Algorithm if TYPE_CHECKING: - from prime_rl.orchestrator.types import TrainRollout + from prime_rl.orchestrator.types import TrainEpisode, TrainRollout from prime_rl.utils.client import InferencePool @@ -21,19 +21,19 @@ class HierarchicalGRPOAlgorithm(Algorithm): interchangeable. ``episode_agents`` lists the roles, normally ``solver``, that are compared - within one episode. Other roles are compared across the full rollout group. - A comparison group with one trace produces zero advantage.""" + within one episode. Other roles are compared across the whole group. A + comparison group with one trace produces zero advantage.""" def __init__(self, config: HierarchicalGRPOAlgoConfig, policy_pool: InferencePool): super().__init__(config, policy_pool) self.episode_agents = set(config.episode_agents) - async def score_group(self, group: list[TrainRollout]) -> None: + async def score_group(self, group: list[TrainEpisode]) -> None: peers: dict[tuple[str, str | None], list[TrainRollout]] = defaultdict(list) - for rollout in group: - episode_scoped = rollout.agent.name in self.episode_agents - key = (rollout.agent.name, rollout.episode_id if episode_scoped else None) - peers[key].append(rollout) + for episode in group: + for rollout in episode.rollouts: + episode_scoped = rollout.agent.name in self.episode_agents + peers[(rollout.agent.name, episode.id if episode_scoped else None)].append(rollout) for members in peers.values(): baseline = sum(rollout.reward for rollout in members) / len(members) for rollout in members: diff --git a/src/prime_rl/orchestrator/algo/max_rl.py b/src/prime_rl/orchestrator/algo/max_rl.py index 58c233e8d4..d56a4bdd25 100644 --- a/src/prime_rl/orchestrator/algo/max_rl.py +++ b/src/prime_rl/orchestrator/algo/max_rl.py @@ -5,9 +5,10 @@ import torch from prime_rl.orchestrator.algo.base import Algorithm +from prime_rl.orchestrator.types import group_rollouts if TYPE_CHECKING: - from prime_rl.orchestrator.types import TrainRollout + from prime_rl.orchestrator.types import TrainEpisode class MaxRLAlgorithm(Algorithm): @@ -23,9 +24,10 @@ class MaxRLAlgorithm(Algorithm): <= 0 carries no signal and gets zero advantages (the zero-advantage filter drops it, matching the paper's no-success convention).""" - async def score_group(self, group: list[TrainRollout]) -> None: - rewards = torch.tensor([rollout.reward for rollout in group], dtype=torch.float32) + async def score_group(self, group: list[TrainEpisode]) -> None: + rollouts = group_rollouts(group) + rewards = torch.tensor([rollout.reward for rollout in rollouts], dtype=torch.float32) 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): + for rollout, advantage in zip(rollouts, advantages.tolist(), strict=True): rollout.assign_advantages(advantage) diff --git a/src/prime_rl/orchestrator/algo/rae.py b/src/prime_rl/orchestrator/algo/rae.py index db4240f30e..18b2000d98 100644 --- a/src/prime_rl/orchestrator/algo/rae.py +++ b/src/prime_rl/orchestrator/algo/rae.py @@ -5,9 +5,10 @@ from prime_rl.configs.algorithm import RAEAlgoConfig from prime_rl.orchestrator.algo.base import Algorithm +from prime_rl.orchestrator.types import group_rollouts if TYPE_CHECKING: - from prime_rl.orchestrator.types import TrainRollout + from prime_rl.orchestrator.types import TrainEpisode from prime_rl.utils.client import InferencePool @@ -34,8 +35,8 @@ def __init__(self, config: RAEAlgoConfig, policy_pool: InferencePool): self.decay = config.decay self.baselines: dict[str, float] = defaultdict(float) - async def score_group(self, group: list[TrainRollout]) -> None: - for rollout in group: + async def score_group(self, group: list[TrainEpisode]) -> None: + for rollout in group_rollouts(group): baseline = self.baselines[rollout.agent.name] rollout.assign_advantages(rollout.reward - baseline) self.baselines[rollout.agent.name] = self.decay * baseline + (1.0 - self.decay) * rollout.reward diff --git a/src/prime_rl/orchestrator/metrics.py b/src/prime_rl/orchestrator/metrics.py index a2d58ef7b9..cc0bd25392 100644 --- a/src/prime_rl/orchestrator/metrics.py +++ b/src/prime_rl/orchestrator/metrics.py @@ -388,13 +388,6 @@ def to_wandb(self, *, prefix: str, subset: Subset) -> dict[str, float]: return out -def keep_traces(episode: Episode, keep: Callable[[Rollout], bool]) -> Episode | None: - """The episode narrowed to the traces that pass ``keep``, or ``None`` if none do. A subset view - stays a list of episodes so the episode-level aggregates keep describing what survived.""" - traces = [t for t in episode.traces if keep(t)] - return episode.model_copy(update={"traces": traces}) if traces else None - - class TrainRollouts: """The train episodes of one window (everything that came back, errored + filtered + untrainable included). ``effective`` is the clean trainable subset — the same episodes, each @@ -419,9 +412,7 @@ def __iter__(self) -> Iterator[TrainRollout]: @property def effective(self) -> TrainRollouts: - kept = ( - keep_traces(e, lambda r: not r.has_error and not r.is_filtered and r.agent.trainable) for e in self.episodes - ) + kept = (e.narrow(lambda r: not r.has_error and not r.is_filtered and r.agent.trainable) for e in self.episodes) return TrainRollouts([e for e in kept if e is not None]) def by_env(self) -> dict[str, TrainRollouts]: @@ -471,7 +462,7 @@ def group_size(self) -> int: @property def effective(self) -> EvalRollouts: - kept = (keep_traces(e, lambda r: not r.has_error and r.agent.trainable) for e in self.episodes) + kept = (e.narrow(lambda r: not r.has_error and r.agent.trainable) for e in self.episodes) return EvalRollouts([e for e in kept if e is not None], group_size=self.group_size) @property diff --git a/src/prime_rl/orchestrator/train_sink.py b/src/prime_rl/orchestrator/train_sink.py index 5cc2e13805..39cd68a041 100644 --- a/src/prime_rl/orchestrator/train_sink.py +++ b/src/prime_rl/orchestrator/train_sink.py @@ -4,9 +4,9 @@ dispatcher producing more rollouts), then the env algorithm's ``finalize_rollout`` (rollout-local scoring + any reference I/O). Errored and untrainable rollouts skip this. -2. ``process_group`` — filters errored rollouts, hands the trainable - survivors to the env algorithm's ``finalize_group`` (advantages + - per-sample wire stamping), runs the pre-batch filter pass. +2. ``process_group`` — filters errored rollouts, hands the episodes narrowed + to their trainable survivors to the env algorithm's ``finalize_group`` + (advantages + per-sample wire stamping), runs the pre-batch filter pass. 3. ``process_batch`` — applies post-batch filter annotations and assembles the trainer-bound ``TrainingSample`` list. Returns a ``TrainBatch``. @@ -27,7 +27,7 @@ from prime_rl.orchestrator.filters import RolloutFilter, apply_filters from prime_rl.orchestrator.metrics import TrainRollouts from prime_rl.orchestrator.trajectories import trace_to_samples -from prime_rl.orchestrator.types import TrainBatch, TrainEpisode, TrainRollout +from prime_rl.orchestrator.types import TrainBatch, TrainEpisode, TrainRollout, group_rollouts from prime_rl.transport import TrainingSample from prime_rl.utils.logger import get_logger @@ -189,8 +189,7 @@ async def process_group(self, group_id: uuid.UUID) -> None: for episode in episodes: self.pending_rollouts.append(episode) task_idx = group[0].task.data.idx if group else -1 - survivors = [r for r in group if not r.has_error] - num_errored = len(group) - len(survivors) + num_errored = sum(r.has_error for r in group) # Group-scoring envs: any failure makes survivors' rewards unsafe # (computed relative to the missing ones) @@ -201,8 +200,10 @@ async def process_group(self, group_id: uuid.UUID) -> None: f"rollouts={len(group)} (errored={num_errored}) | dropped: group-scored partial" ) return - # Untrainable traces carry no samples and must not skew the group baseline. - survivors = [r for r in survivors if r.agent.trainable] + # Untrainable traces carry no samples and must not skew the group baseline. The cohort + # stays episodes so the algorithm can still see which attempts shared one. + cohort = [n for e in episodes if (n := e.narrow(lambda r: not r.has_error and r.agent.trainable))] + survivors = group_rollouts(cohort) if not survivors: get_logger().debug( f"Finished group | env={env_name} task_idx={task_idx} | " @@ -213,7 +214,7 @@ async def process_group(self, group_id: uuid.UUID) -> None: # Advantages + per-sample wire stamping (advantage stream, loss # routing) are the algorithm's job (finalize_group); the sink only # owns the grouping mechanics. - await env.algorithm.finalize_group(survivors) + await env.algorithm.finalize_group(cohort) # The env has a single sampling temperature; fan it out per token # (context tokens are masked out, so their temperature is don't-care). diff --git a/src/prime_rl/orchestrator/types.py b/src/prime_rl/orchestrator/types.py index 8fd6a100e3..2f399bb5b2 100644 --- a/src/prime_rl/orchestrator/types.py +++ b/src/prime_rl/orchestrator/types.py @@ -3,8 +3,9 @@ from __future__ import annotations import uuid +from collections.abc import Callable, Iterable from dataclasses import dataclass -from typing import TYPE_CHECKING, Any, ClassVar, Generic, Literal, Protocol, cast +from typing import TYPE_CHECKING, Any, ClassVar, Generic, Literal, Protocol, Self, cast import verifiers.v1 as vf from pydantic import ConfigDict, Field @@ -150,6 +151,13 @@ def rollouts(self) -> list[Rollout]: """The episode's traces, typed as the rollouts prime-rl works with.""" return cast(list[Rollout], self.traces) + def narrow(self, keep: Callable[[Rollout], bool]) -> Self | None: + """This episode with only the traces that pass ``keep``, or ``None`` if none do. A subset + stays a list of episodes rather than a flat trace list, so the episode-level aggregates + keep describing what survived. The kept traces are the same objects, not copies.""" + traces = [t for t in self.traces if keep(t)] + return self.model_copy(update={"traces": traces}) if traces else None + def to_record(self) -> dict[str, Any]: """JSON record without raw tensors — the episode form of ``Trace.to_record``, and the unit ``traces.jsonl`` stores: one episode per line, matching what verifiers writes and what its @@ -188,6 +196,12 @@ class EvalEpisode(Episode): step: int = Field(default=0, exclude=True) +def group_rollouts(episodes: Iterable[TrainEpisode]) -> list[TrainRollout]: + """Every trace of a group, flat — the view an algorithm comparing across the whole cohort + wants, where the episode an attempt came from does not matter.""" + return [r for e in episodes for r in e.rollouts] + + @dataclass class InflightEpisode: """One episode in flight, and the facts of the dispatch that will be stamped onto it when it diff --git a/tests/unit/orchestrator/test_advantage.py b/tests/unit/orchestrator/test_advantage.py index a2bcb1adfb..adc307ef33 100644 --- a/tests/unit/orchestrator/test_advantage.py +++ b/tests/unit/orchestrator/test_advantage.py @@ -11,7 +11,7 @@ from prime_rl.orchestrator.algo.grpo import GRPOAlgorithm from prime_rl.orchestrator.algo.max_rl import MaxRLAlgorithm from prime_rl.orchestrator.trajectories import iter_trainable_branches, trace_to_samples -from prime_rl.orchestrator.types import TrainRollout +from prime_rl.orchestrator.types import TrainEpisode, TrainRollout def _build_rollout( @@ -138,6 +138,12 @@ def _make_group(rewards, completion_lengths=None, num_turns=None) -> list[TrainR return rollouts +def _as_episodes(group: list[TrainRollout]) -> list[TrainEpisode]: + """One episode per rollout — the shape a single-agent env produces, and what the + algorithms are handed.""" + return [TrainEpisode.model_construct(id=f"e{i}", traces=[rollout]) for i, rollout in enumerate(group)] + + def _scalar(rollout: TrainRollout) -> float: """The per-rollout advantage scalar an algorithm assigned — broadcast over the rollout's trainable tokens, so every assigned position holds it.""" @@ -148,14 +154,14 @@ def _scalar(rollout: TrainRollout) -> float: def _grpo(group: list[TrainRollout], length_penalty=None) -> list[float]: """Drive ``GRPOAlgorithm.score_group`` and read back each per-rollout scalar.""" algo = GRPOAlgorithm(GRPOAlgoConfig(length_penalty=length_penalty), policy_pool=None) - asyncio.run(algo.score_group(group)) + asyncio.run(algo.score_group(_as_episodes(group))) return [_scalar(rollout) for rollout in group] def _max_rl(group: list[TrainRollout]) -> list[float]: """Drive ``MaxRLAlgorithm.score_group`` and read back each per-rollout scalar.""" algo = MaxRLAlgorithm(MaxRLAlgoConfig(), policy_pool=None) - asyncio.run(algo.score_group(group)) + asyncio.run(algo.score_group(_as_episodes(group))) return [_scalar(rollout) for rollout in group] @@ -217,7 +223,7 @@ def test_linear_context_term_penalizes_more_context(): _build_rollout(1.0, sampled_lengths=[10], obs_lengths=[]), _build_rollout(1.0, sampled_lengths=[10], obs_lengths=[100]), ] - asyncio.run(GRPOAlgorithm(GRPOAlgoConfig(length_penalty=cfg), policy_pool=None).score_group(group)) + asyncio.run(GRPOAlgorithm(GRPOAlgoConfig(length_penalty=cfg), policy_pool=None).score_group(_as_episodes(group))) advs = [_scalar(rollout) for rollout in group] assert advs[0] > advs[1] assert sum(advs) == pytest.approx(0.0, abs=1e-6) From 74f73a411ae64574cad0990ff28a46a3b238a910 Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Tue, 4 Aug 2026 18:37:26 +0000 Subject: [PATCH 35/58] docs: the scoring hooks take episodes Co-Authored-By: Claude Fable 5 --- docs/algorithms.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/algorithms.md b/docs/algorithms.md index 1392c3538c..a2bcd1e7db 100644 --- a/docs/algorithms.md +++ b/docs/algorithms.md @@ -168,7 +168,7 @@ At runtime, each env's resolved config builds two objects: a `Sampler` (`prime_r Each class owns its hooks outright — reading one top to bottom reads the algorithm, and everything on the class is an override point. The two hooks are one scope-and-timing ladder — the wider scope is unlocked by a later barrier, so the two axes coincide. Each is handed the env's own data directly — a `Rollout` (the typed trace: `reward`, `nodes`, `num_turns`, ... with `samples` attached) on arrival, the group's `Episode`s at group time — plus `assign_advantages` to write credit: -- `async score_rollout(rollout)` — one rollout, **on arrival** (as it's tokenized, before its group is complete): rollout-local credit (`rollout.assign_advantages(...)`, scalar broadcast or per-token), observation ce weights, **or** model I/O — query a reference pool (e.g. `self.teacher_pool`, connected in `setup()` via `self.connect(...)`, or the live `self.policy_pool` for opsd) and attach per-token results (e.g. teacher logprobs) with bounded concurrency. No siblings. `echo` weights observation tokens here, identifying env-provided observation nodes by their non-sampled status and source step role attribution, applying the optional user filter, and writing the `ce_weights` stream. Model I/O runs *before* the pre-batch filters, so it pays compute on rollouts that may then be filtered out. +- `async score_rollout(rollout)` — one rollout, **on arrival** (as it's tokenized, before its group is complete): rollout-local credit (`rollout.assign_advantages(...)`), observation ce weights, **or** model I/O — query a reference pool (e.g. `self.teacher_pool`, connected in `setup()` via `self.connect(...)`, or the live `self.policy_pool` for opsd) and attach per-token results (e.g. teacher logprobs) with bounded concurrency. No siblings. `echo` weights observation tokens here, identifying env-provided observation nodes by their non-sampled status and source step role attribution, applying the optional user filter, and writing the `ce_weights` stream. Model I/O runs *before* the pre-batch filters, so it pays compute on rollouts that may then be filtered out. - `score_group(group)` — the cohort, **before filtering** (filters read the streams), synchronous: group-relative credit (GRPO/MaxRL baselines). `group` is a list of `TrainEpisode`, so an algorithm can compare within an episode as well as across them (`hierarchical_grpo` does); `group_rollouts(group)` flattens it for the ones that only compare across. The pipeline drives the hooks through two non-virtual methods it never looks inside: `algorithm.finalize_rollout(rollout)` per arrival (rollout-local scoring + reference I/O) and `algorithm.finalize_group(episodes)` per group (scoring + wire stamping; after this the records are frozen — groups die at stamping). Sample construction (interleaving) is pure pipeline — observation-token provenance is available through structural attribution (`node.sampled`, `node.is_content`) for any algorithm that trains on env-provided tokens. @@ -420,7 +420,7 @@ Both of `kuhn-poker-v1`'s agents late-bind to the run's own model — shared-pol ### Authoring an Algorithm -There is no config hook that points at user code — a new credit-assignment scheme is a new named algorithm in the repo. Subclass `Algorithm`, assign credit in the scoring hook whose timing fits your signal, and register the class. The hook receives the group's `Rollout`s (each the env's typed `verifiers.Trace` — turns, tool calls, metadata in `info` — with `samples` attached) and writes credit via `assign_advantages`: +There is no config hook that points at user code — a new credit-assignment scheme is a new named algorithm in the repo. Subclass `Algorithm`, assign credit in the scoring hook whose timing fits your signal, and register the class. The hook receives the group's `TrainEpisode`s, each holding the env's typed `verifiers.Trace`s — turns, tool calls, metadata in `info` — with `samples` attached, and writes credit via `assign_advantages`: ```python # src/prime_rl/orchestrator/algo/my_algo.py From 8e4af214498c000e11ec9644c9018dcb74bff262 Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Tue, 4 Aug 2026 18:39:19 +0000 Subject: [PATCH 36/58] fix(orchestrator): a forked node's credit is context in the later branch MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A sampled node shared by several branches is trainable only in the first one; the branch view still spreads its credit everywhere it appears, so zero the positions the sample does not train on — the layout the trainer has always been given. Co-Authored-By: Claude Fable 5 --- src/prime_rl/orchestrator/algo/routing.py | 12 ++++--- tests/unit/orchestrator/test_advantage.py | 38 +++++++++++++++++++++++ 2 files changed, 46 insertions(+), 4 deletions(-) diff --git a/src/prime_rl/orchestrator/algo/routing.py b/src/prime_rl/orchestrator/algo/routing.py index 3d3192ba63..aaac58f132 100644 --- a/src/prime_rl/orchestrator/algo/routing.py +++ b/src/prime_rl/orchestrator/algo/routing.py @@ -53,8 +53,12 @@ def stamp_loss_routing(sample: TrainingSample, action_loss_type: ActionLossType) def stamp_advantages(rollout: TrainRollout) -> None: - """Copy each trainable branch's per-token credit onto the sample built from it. The branch - spreads its nodes' values across its own tokens, so the two align by construction; a rollout - that was never scored (opd/opsd) leaves ``None`` and ships no advantage stream.""" + """Copy each trainable branch's per-token credit onto the sample built from it, zeroed where + the sample does not train. The branch spreads its nodes' values across its own tokens, so the + two align by construction, but a node shared with an earlier branch is credited there and is + only context here. A rollout that was never scored (opd/opsd) ships no advantage stream.""" for sample, (branch, _) in zip(rollout.samples, iter_trainable_branches(rollout), strict=True): - sample.advantages = branch.advantages + advantages = branch.advantages + if advantages is None: + continue + sample.advantages = [a if trains else 0.0 for a, trains in zip(advantages, sample.mask, strict=True)] diff --git a/tests/unit/orchestrator/test_advantage.py b/tests/unit/orchestrator/test_advantage.py index adc307ef33..3ce54f56bd 100644 --- a/tests/unit/orchestrator/test_advantage.py +++ b/tests/unit/orchestrator/test_advantage.py @@ -10,6 +10,7 @@ ) from prime_rl.orchestrator.algo.grpo import GRPOAlgorithm from prime_rl.orchestrator.algo.max_rl import MaxRLAlgorithm +from prime_rl.orchestrator.algo.routing import stamp_advantages from prime_rl.orchestrator.trajectories import iter_trainable_branches, trace_to_samples from prime_rl.orchestrator.types import TrainEpisode, TrainRollout @@ -272,3 +273,40 @@ def test_unassigned_credit_is_not_zero_credit(): scored = _build_rollout(0.0, sampled_lengths=[2]) scored.assign_advantages(0.0) assert scored.advantages == [0.0, 0.0] and not scored.is_trainable # scored, but no gradient + + +def test_stamp_advantages_zeros_a_shared_node_in_the_later_branch(): + """A forked node is trained in the first branch containing it; in the later branch it is + context, so its credit must not ride along on that branch's sample.""" + root = vf.MessageNode( + message=vf.AssistantMessage(role="assistant", content=""), + sampled=True, + token_ids=[1, 2], + mask=[True, True], + logprobs=[-0.1, -0.2], + ) + leaves = [ + vf.MessageNode( + parent=0, + message=vf.AssistantMessage(role="assistant", content=""), + sampled=True, + token_ids=[3], + mask=[True], + logprobs=[-0.3], + ) + for _ in range(2) + ] + rollout = TrainRollout[vf.TaskData]( + task=vf.TraceTask(type="Task", data=vf.TaskData(idx=0, prompt=None)), + agent=vf.AgentInfo(config=vf.AgentConfig()), + nodes=[root, *leaves], + rewards={"reward": vf.Reward(score=0.0)}, + ) + rollout.env_name = "test" + rollout.samples = trace_to_samples(rollout, env_name="test") + rollout.assign_advantages(0.5) + stamp_advantages(rollout) + + first, second = rollout.samples + assert first.advantages == [0.5, 0.5, 0.5] + assert second.advantages == [0.0, 0.0, 0.5] From 7162b2524299564f8f35a1a24e98db9a43ad18e1 Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Tue, 4 Aug 2026 18:49:03 +0000 Subject: [PATCH 37/58] refactor(orchestrator)!: one Episode, discriminated by vf's run record The train/eval split was a type distinction the producer already knows: the dispatcher is what decides which path an episode is on. It now writes that into vf's own RunInfo when the episode lands, so TrainEpisode, EvalEpisode and the KIND ClassVar all collapse into one Episode and the main loop reads run.type instead of asking isinstance. An eval episode's step rides in run.step; a train one's is filled by the loop that knows which batch window is collecting. Also re-pins deps/verifiers to vf#2245 as merged on main. Co-Authored-By: Claude Fable 5 --- deps/verifiers | 2 +- src/prime_rl/orchestrator/algo/base.py | 8 +- src/prime_rl/orchestrator/algo/grpo.py | 4 +- .../orchestrator/algo/hierarchical_grpo.py | 4 +- src/prime_rl/orchestrator/algo/max_rl.py | 4 +- src/prime_rl/orchestrator/algo/rae.py | 4 +- src/prime_rl/orchestrator/dispatcher.py | 6 +- src/prime_rl/orchestrator/eval_sink.py | 8 +- src/prime_rl/orchestrator/metrics.py | 10 +-- src/prime_rl/orchestrator/orchestrator.py | 28 +++---- src/prime_rl/orchestrator/train_sink.py | 6 +- src/prime_rl/orchestrator/types.py | 78 +++++++++---------- tests/unit/orchestrator/test_advantage.py | 6 +- tests/unit/orchestrator/test_metrics.py | 36 ++++----- 14 files changed, 93 insertions(+), 111 deletions(-) diff --git a/deps/verifiers b/deps/verifiers index 7eba3e65d2..48591d3735 160000 --- a/deps/verifiers +++ b/deps/verifiers @@ -1 +1 @@ -Subproject commit 7eba3e65d2e5df2249519b387f2e7cc6658abf7b +Subproject commit 48591d3735f3bb42dd93c7878adaa6ad6c5ff6be diff --git a/src/prime_rl/orchestrator/algo/base.py b/src/prime_rl/orchestrator/algo/base.py index 410ff4b858..a7962c0327 100644 --- a/src/prime_rl/orchestrator/algo/base.py +++ b/src/prime_rl/orchestrator/algo/base.py @@ -50,7 +50,7 @@ if TYPE_CHECKING: from renderers import RendererConfig - from prime_rl.orchestrator.types import TrainEpisode, TrainRollout + from prime_rl.orchestrator.types import Episode, TrainRollout from prime_rl.utils.client import InferencePool @@ -94,7 +94,7 @@ class Algorithm: the algorithm declares, resolving each reference via :meth:`connect`; - the two scoring hooks, each ``async`` and given the env's own data directly — a :class:`TrainRollout` on arrival, the group's - :class:`TrainEpisode`\ s at group time — so a hook reads the trace and + :class:`Episode`\ s at group time — so a hook reads the trace and writes credit via :meth:`TrainRollout.assign_advantages`. They are async so either stage may do I/O — e.g. a process-reward model or a teacher at arrival, or a judge at group time whose signal a pre-batch @@ -146,7 +146,7 @@ async def score_rollout(self, rollout: TrainRollout) -> None: connected in :meth:`setup`, or the live policy (opsd). No siblings, no group stats.""" - async def score_group(self, group: list[TrainEpisode]) -> None: + async def score_group(self, group: list[Episode]) -> None: """Group phase, the finalized cohort, before filtering: write group-relative credit. The cohort arrives as episodes, so an algorithm can compare within one episode as well as across them; ``group_rollouts`` @@ -158,7 +158,7 @@ async def finalize_rollout(self, rollout: TrainRollout) -> None: if rollout.samples: await self.score_rollout(rollout) - async def finalize_group(self, episodes: list[TrainEpisode]) -> None: + async def finalize_group(self, episodes: list[Episode]) -> None: """Group phase (non-virtual): group-relative scoring, then stamp each sample's wire fields (the advantage stream + loss routing). After this the records are frozen — groups die at stamping.""" diff --git a/src/prime_rl/orchestrator/algo/grpo.py b/src/prime_rl/orchestrator/algo/grpo.py index c002858bb2..1ca2024320 100644 --- a/src/prime_rl/orchestrator/algo/grpo.py +++ b/src/prime_rl/orchestrator/algo/grpo.py @@ -9,7 +9,7 @@ from prime_rl.orchestrator.types import group_rollouts if TYPE_CHECKING: - from prime_rl.orchestrator.types import TrainEpisode + from prime_rl.orchestrator.types import Episode from prime_rl.utils.client import InferencePool @@ -22,7 +22,7 @@ def __init__(self, config: GRPOAlgoConfig, policy_pool: InferencePool): super().__init__(config, policy_pool) self.length_penalty = config.length_penalty - async def score_group(self, group: list[TrainEpisode]) -> None: + async def score_group(self, group: list[Episode]) -> None: rollouts = group_rollouts(group) rewards = torch.tensor([rollout.reward for rollout in rollouts], dtype=torch.float32) length_penalty = self.length_penalty diff --git a/src/prime_rl/orchestrator/algo/hierarchical_grpo.py b/src/prime_rl/orchestrator/algo/hierarchical_grpo.py index f97bdcae5f..6495115c9e 100644 --- a/src/prime_rl/orchestrator/algo/hierarchical_grpo.py +++ b/src/prime_rl/orchestrator/algo/hierarchical_grpo.py @@ -7,7 +7,7 @@ from prime_rl.orchestrator.algo.base import Algorithm if TYPE_CHECKING: - from prime_rl.orchestrator.types import TrainEpisode, TrainRollout + from prime_rl.orchestrator.types import Episode, TrainRollout from prime_rl.utils.client import InferencePool @@ -28,7 +28,7 @@ def __init__(self, config: HierarchicalGRPOAlgoConfig, policy_pool: InferencePoo super().__init__(config, policy_pool) self.episode_agents = set(config.episode_agents) - async def score_group(self, group: list[TrainEpisode]) -> None: + async def score_group(self, group: list[Episode]) -> None: peers: dict[tuple[str, str | None], list[TrainRollout]] = defaultdict(list) for episode in group: for rollout in episode.rollouts: diff --git a/src/prime_rl/orchestrator/algo/max_rl.py b/src/prime_rl/orchestrator/algo/max_rl.py index d56a4bdd25..f013b24df2 100644 --- a/src/prime_rl/orchestrator/algo/max_rl.py +++ b/src/prime_rl/orchestrator/algo/max_rl.py @@ -8,7 +8,7 @@ from prime_rl.orchestrator.types import group_rollouts if TYPE_CHECKING: - from prime_rl.orchestrator.types import TrainEpisode + from prime_rl.orchestrator.types import Episode class MaxRLAlgorithm(Algorithm): @@ -24,7 +24,7 @@ class MaxRLAlgorithm(Algorithm): <= 0 carries no signal and gets zero advantages (the zero-advantage filter drops it, matching the paper's no-success convention).""" - async def score_group(self, group: list[TrainEpisode]) -> None: + async def score_group(self, group: list[Episode]) -> None: rollouts = group_rollouts(group) rewards = torch.tensor([rollout.reward for rollout in rollouts], dtype=torch.float32) mean = rewards.mean() diff --git a/src/prime_rl/orchestrator/algo/rae.py b/src/prime_rl/orchestrator/algo/rae.py index 18b2000d98..33abfdc1e5 100644 --- a/src/prime_rl/orchestrator/algo/rae.py +++ b/src/prime_rl/orchestrator/algo/rae.py @@ -8,7 +8,7 @@ from prime_rl.orchestrator.types import group_rollouts if TYPE_CHECKING: - from prime_rl.orchestrator.types import TrainEpisode + from prime_rl.orchestrator.types import Episode from prime_rl.utils.client import InferencePool @@ -35,7 +35,7 @@ def __init__(self, config: RAEAlgoConfig, policy_pool: InferencePool): self.decay = config.decay self.baselines: dict[str, float] = defaultdict(float) - async def score_group(self, group: list[TrainEpisode]) -> None: + async def score_group(self, group: list[Episode]) -> None: for rollout in group_rollouts(group): baseline = self.baselines[rollout.agent.name] rollout.assign_advantages(rollout.reward - baseline) diff --git a/src/prime_rl/orchestrator/dispatcher.py b/src/prime_rl/orchestrator/dispatcher.py index 7bb8e6e8b9..3cf6cd60f0 100644 --- a/src/prime_rl/orchestrator/dispatcher.py +++ b/src/prime_rl/orchestrator/dispatcher.py @@ -131,11 +131,13 @@ def __init__( eval_source: EvalSource | None, policy_pool: InferencePool, policy: Policy, + run_id: str, max_inflight_episodes: int, tasks_per_minute: float | None, max_off_policy_steps: int, ) -> None: self.policy = policy + self.run_id = run_id self.train_envs = train_envs self.eval_envs = eval_envs # Train rollouts go to the env sampler's pool; eval always @@ -563,7 +565,9 @@ async def emit_episode( for rollout in episode.traces: rollout.env_name = meta.env_name rollout.group_id = meta.group_id - await self.out_q.put(meta.stamp(episode, policy_version=policy_version, eval_step=eval_step)) + await self.out_q.put( + meta.stamp(episode, run_id=self.run_id, policy_version=policy_version, eval_step=eval_step) + ) async def emit_failed_episodes(self, meta: InflightEpisode, group: GroupState | None, error: vf.Error) -> None: """Emit one traceless episode per rollout the task owed, so the sink still counts its way to diff --git a/src/prime_rl/orchestrator/eval_sink.py b/src/prime_rl/orchestrator/eval_sink.py index 1a892a1f05..9c09e4daa5 100644 --- a/src/prime_rl/orchestrator/eval_sink.py +++ b/src/prime_rl/orchestrator/eval_sink.py @@ -19,7 +19,7 @@ from prime_rl.orchestrator.envs import EvalEnvs from prime_rl.orchestrator.metrics import EvalRollouts -from prime_rl.orchestrator.types import EvalBatch, EvalEpisode, Rollout +from prime_rl.orchestrator.types import Episode, EvalBatch, Rollout from prime_rl.utils.logger import get_logger @@ -28,13 +28,13 @@ class EvalSink: def __init__(self, *, eval_envs: EvalEnvs) -> None: self.eval_envs = eval_envs - self.pending_groups: dict[uuid.UUID, list[EvalEpisode]] = defaultdict(list) + self.pending_groups: dict[uuid.UUID, list[Episode]] = defaultdict(list) # Episodes arrived per group / per batch bucket — the finalization counts. self.pending_group_episodes: dict[uuid.UUID, int] = defaultdict(int) - self.pending_batches: dict[tuple[str, int], list[EvalEpisode]] = defaultdict(list) + self.pending_batches: dict[tuple[str, int], list[Episode]] = defaultdict(list) self.pending_batch_episodes: dict[tuple[str, int], int] = defaultdict(int) - def add(self, episode: EvalEpisode) -> EvalBatch | None: + def add(self, episode: Episode) -> EvalBatch | None: """Process one episode arrival; finalize the group on the ``group_size``-th episode and the per-env epoch on the ``num_examples × group_size``-th. A failed episode brings no rollouts but still counts toward both.""" diff --git a/src/prime_rl/orchestrator/metrics.py b/src/prime_rl/orchestrator/metrics.py index cc0bd25392..d5a0cf0e08 100644 --- a/src/prime_rl/orchestrator/metrics.py +++ b/src/prime_rl/orchestrator/metrics.py @@ -26,7 +26,7 @@ from prime_rl.orchestrator.utils import compute_pass_metrics if TYPE_CHECKING: - from prime_rl.orchestrator.types import EvalEpisode, Rollout, TrainEpisode, TrainRollout + from prime_rl.orchestrator.types import Episode, Rollout, TrainRollout Subset = Literal["all", "effective"] @@ -394,10 +394,10 @@ class TrainRollouts: narrowed to its surviving traces; ``metrics`` builds ``TrainMetrics`` over them. Sized and iterated by *rollout*, since that is the unit almost every consumer wants.""" - def __init__(self, episodes: list[TrainEpisode] | None = None) -> None: + def __init__(self, episodes: list[Episode] | None = None) -> None: self.episodes = episodes if episodes is not None else [] - def append(self, episode: TrainEpisode) -> None: + def append(self, episode: Episode) -> None: self.episodes.append(episode) @property @@ -416,7 +416,7 @@ def effective(self) -> TrainRollouts: return TrainRollouts([e for e in kept if e is not None]) def by_env(self) -> dict[str, TrainRollouts]: - grouped: dict[str, list[TrainEpisode]] = {} + grouped: dict[str, list[Episode]] = {} for episode in self.episodes: grouped.setdefault(episode.env_name, []).append(episode) return {env: TrainRollouts(episodes) for env, episodes in grouped.items()} @@ -432,7 +432,7 @@ class EvalRollouts: ``group_size`` (rollouts per example, the ``avg@k`` k) is derived from the full epoch and carried onto ``effective`` so both subsets share one stable key; ``metrics`` builds ``EvalMetrics``.""" - def __init__(self, episodes: list[EvalEpisode] | None = None, group_size: int | None = None) -> None: + def __init__(self, episodes: list[Episode] | None = None, group_size: int | None = None) -> None: self.episodes = episodes if episodes is not None else [] self._group_size = group_size diff --git a/src/prime_rl/orchestrator/orchestrator.py b/src/prime_rl/orchestrator/orchestrator.py index d3527b5f86..1f8377c8db 100644 --- a/src/prime_rl/orchestrator/orchestrator.py +++ b/src/prime_rl/orchestrator/orchestrator.py @@ -28,7 +28,6 @@ from typing import TYPE_CHECKING import tomli_w -import verifiers.v1 as vf from modelexpress import p2p_pb2 from modelexpress.client import MxClient @@ -59,7 +58,6 @@ from prime_rl.orchestrator.types import ( Episode, EvalBatch, - EvalEpisode, Policy, Progress, TrainBatch, @@ -407,6 +405,7 @@ async def setup(self) -> None: eval_source=self.eval_source, policy_pool=self.policy_inference, policy=self.policy, + run_id=self.run_id, max_inflight_episodes=config.max_inflight_episodes, tasks_per_minute=config.tasks_per_minute, max_off_policy_steps=config.max_off_policy_steps, @@ -528,26 +527,21 @@ async def main_loop(self) -> None: # Every completed episode — errored, filtered, or never batched — lands in the # ``all`` trace file the moment it arrives, so it survives crashes and drains. One # episode per line, so an episode that produced no traces still records why. - # Train episodes belong to the batch window currently collecting (``progress.step``), - # eval ones to the step whose eval triggered them. - is_eval = isinstance(episode, EvalEpisode) - step = episode.step if isinstance(episode, EvalEpisode) else self.progress.step - run: vf.RunInfo = ( - vf.EvalRunInfo(id=self.run_id, step=step) if is_eval else vf.TrainRunInfo(id=self.run_id, step=step) - ) - episode.record_run( - run, - env_name=episode.env_name, - group_id=str(episode.group_id), - policy_version=episode.policy_version, - ) + # An eval episode already knows its step (the eval that triggered it); a train one + # belongs to the batch window collecting right now, which only this loop knows. + run = episode.run + assert run is not None, "the dispatcher records the run when the episode lands" + if run.type == "train": + run.step = self.progress.step + step = run.step + assert step is not None await asyncio.to_thread( save_episodes, [episode.to_record()], - get_trace_path(self.config.output_dir, step, episode.KIND, "all"), + get_trace_path(self.config.output_dir, step, run.type, "all"), ) - if isinstance(episode, EvalEpisode): + if run.type == "eval": assert self.eval_sink is not None # eval rollouts only emitted when eval is configured eval_batch = self.eval_sink.add(episode) if eval_batch is not None: diff --git a/src/prime_rl/orchestrator/train_sink.py b/src/prime_rl/orchestrator/train_sink.py index 39cd68a041..8aee8820f2 100644 --- a/src/prime_rl/orchestrator/train_sink.py +++ b/src/prime_rl/orchestrator/train_sink.py @@ -27,7 +27,7 @@ from prime_rl.orchestrator.filters import RolloutFilter, apply_filters from prime_rl.orchestrator.metrics import TrainRollouts from prime_rl.orchestrator.trajectories import trace_to_samples -from prime_rl.orchestrator.types import TrainBatch, TrainEpisode, TrainRollout, group_rollouts +from prime_rl.orchestrator.types import Episode, TrainBatch, TrainRollout, group_rollouts from prime_rl.transport import TrainingSample from prime_rl.utils.logger import get_logger @@ -79,7 +79,7 @@ def __init__( # Keyed by the dispatcher's group UUID. ``(env_name, task_idx)`` # isn't unique — the same task can be re-sampled while an # earlier group is still in flight - self.pending_groups: dict[uuid.UUID, list[TrainEpisode]] = defaultdict(list) + self.pending_groups: dict[uuid.UUID, list[Episode]] = defaultdict(list) # Episodes arrived per group — the finalization count (an episode may # add several traces to ``pending_groups`` but counts once here). self.pending_group_episodes: dict[uuid.UUID, int] = defaultdict(int) @@ -124,7 +124,7 @@ def pending_batch_by_env(self) -> dict[str, int]: counts[r.env_name] += 1 return dict(counts) - async def add(self, episode: TrainEpisode) -> TrainBatch | None: + async def add(self, episode: Episode) -> TrainBatch | None: """Process one episode arrival; finalize the group on the ``group_size``-th episode; return a ``TrainBatch`` if the finalization pushed (or left) the batch over its threshold. Arrivals into diff --git a/src/prime_rl/orchestrator/types.py b/src/prime_rl/orchestrator/types.py index 2f399bb5b2..4c34dea188 100644 --- a/src/prime_rl/orchestrator/types.py +++ b/src/prime_rl/orchestrator/types.py @@ -5,7 +5,7 @@ import uuid from collections.abc import Callable, Iterable from dataclasses import dataclass -from typing import TYPE_CHECKING, Any, ClassVar, Generic, Literal, Protocol, Self, cast +from typing import TYPE_CHECKING, Any, Generic, Literal, Protocol, Self, cast import verifiers.v1 as vf from pydantic import ConfigDict, Field @@ -132,24 +132,25 @@ class Episode(vf.WireEpisode): raised before reaching the env — are minted the same way. So ``is_empty`` is simply "no traces", and ``last_error`` says why in one vocabulary for every cause. - Train and eval are the two subclasses rather than a discriminator field, so each carries only - what its path means: an eval episode has a step it belongs to, a train episode has the policy - it was generated from.""" + Which path it is on is ``run.type``, vf's own discriminator, stamped by the dispatcher when the + episode lands — so there is one episode class and no prime-rl-side kind.""" model_config = ConfigDict(arbitrary_types_allowed=True) # traces are ``Rollout``s - KIND: ClassVar[RolloutKind] - """Which path this episode is on, for the run record and the dispatcher's counters.""" - env_name: str = Field(default="", exclude=True) + """The env as prime-rl names it (the config key), which is not vf's ``env.id``.""" group_id: uuid.UUID = Field(default_factory=uuid.uuid4, exclude=True) policy_version: int = Field(default=0, exclude=True) """The policy that generated it — the thing being trained on one path, measured on the other.""" + off_policy_steps: int = Field(default=0, exclude=True) + """How stale it was by the time it shipped. Always 0 on the eval path, which never trains.""" @property - def rollouts(self) -> list[Rollout]: - """The episode's traces, typed as the rollouts prime-rl works with.""" - return cast(list[Rollout], self.traces) + def rollouts(self) -> list[TrainRollout]: + """The episode's traces, typed as the rollouts prime-rl works with. Every trace is built as + a ``TrainRollout`` (``Env.run`` is shared), so the training fields are always reachable — + on the eval path they simply stay empty.""" + return cast(list[TrainRollout], self.traces) def narrow(self, keep: Callable[[Rollout], bool]) -> Self | None: """This episode with only the traces that pass ``keep``, or ``None`` if none do. A subset @@ -173,30 +174,7 @@ def is_empty(self) -> bool: return not self.traces -class TrainEpisode(Episode): - """An episode collected for training, which alone can go stale relative to the live policy.""" - - KIND: ClassVar[RolloutKind] = "train" - - off_policy_steps: int = Field(default=0, exclude=True) - """How stale it was by the time it shipped — meaningless for eval, which never trains on it.""" - - @property - def rollouts(self) -> list[TrainRollout]: - return cast(list[TrainRollout], self.traces) - - -class EvalEpisode(Episode): - """An episode collected for one eval epoch. ``step`` is the training step whose eval triggered - it — always known, unlike on the train path, so it is not optional here. It is the source for - the ``run.step`` each of its traces gets stamped with on arrival.""" - - KIND: ClassVar[RolloutKind] = "eval" - - step: int = Field(default=0, exclude=True) - - -def group_rollouts(episodes: Iterable[TrainEpisode]) -> list[TrainRollout]: +def group_rollouts(episodes: Iterable[Episode]) -> list[TrainRollout]: """Every trace of a group, flat — the view an algorithm comparing across the whole cohort wants, where the episode an attempt came from does not matter.""" return [r for e in episodes for r in e.rollouts] @@ -218,20 +196,34 @@ class InflightEpisode: off_policy_steps: int = 0 eval_step: int | None = None - def stamp(self, episode: vf.WireEpisode, *, policy_version: int, eval_step: int | None) -> Episode: + def stamp(self, episode: vf.WireEpisode, *, run_id: str, policy_version: int, eval_step: int | None) -> Episode: """Mint the landed episode: the env's own, carrying the dispatch it came from. The group's values win over this dispatch's when it is still alive, so they are passed in rather than - read off ``self``.""" - common = { + read off ``self``. + + The run record is written here because this is where a dispatch's facts become an episode's, + and its ``type`` is what tells the rest of the orchestrator which path the episode is on. A + train episode's step is not known yet — it belongs to whichever batch window is collecting + when it lands, so the main loop fills it in.""" + landed = Episode.model_construct( **dict(episode), - "env_name": self.env_name, - "group_id": self.group_id, - "policy_version": policy_version, - } + env_name=self.env_name, + group_id=self.group_id, + policy_version=policy_version, + off_policy_steps=self.off_policy_steps, + ) if self.kind == "eval": assert eval_step is not None, "eval episode missing its step" - return EvalEpisode.model_construct(**common, step=eval_step) - return TrainEpisode.model_construct(**common, off_policy_steps=self.off_policy_steps) + run: vf.RunInfo = vf.EvalRunInfo(id=run_id, step=eval_step) + else: + run = vf.TrainRunInfo(id=run_id) + landed.record_run( + run, + env_name=self.env_name, + group_id=str(self.group_id), + policy_version=policy_version, + ) + return landed @dataclass diff --git a/tests/unit/orchestrator/test_advantage.py b/tests/unit/orchestrator/test_advantage.py index 3ce54f56bd..95505a28f1 100644 --- a/tests/unit/orchestrator/test_advantage.py +++ b/tests/unit/orchestrator/test_advantage.py @@ -12,7 +12,7 @@ from prime_rl.orchestrator.algo.max_rl import MaxRLAlgorithm from prime_rl.orchestrator.algo.routing import stamp_advantages from prime_rl.orchestrator.trajectories import iter_trainable_branches, trace_to_samples -from prime_rl.orchestrator.types import TrainEpisode, TrainRollout +from prime_rl.orchestrator.types import Episode, TrainRollout def _build_rollout( @@ -139,10 +139,10 @@ def _make_group(rewards, completion_lengths=None, num_turns=None) -> list[TrainR return rollouts -def _as_episodes(group: list[TrainRollout]) -> list[TrainEpisode]: +def _as_episodes(group: list[TrainRollout]) -> list[Episode]: """One episode per rollout — the shape a single-agent env produces, and what the algorithms are handed.""" - return [TrainEpisode.model_construct(id=f"e{i}", traces=[rollout]) for i, rollout in enumerate(group)] + return [Episode.model_construct(id=f"e{i}", traces=[rollout]) for i, rollout in enumerate(group)] def _scalar(rollout: TrainRollout) -> float: diff --git a/tests/unit/orchestrator/test_metrics.py b/tests/unit/orchestrator/test_metrics.py index 457e205cbd..f79abaecd0 100644 --- a/tests/unit/orchestrator/test_metrics.py +++ b/tests/unit/orchestrator/test_metrics.py @@ -8,7 +8,7 @@ import verifiers.v1 as vf from prime_rl.orchestrator.metrics import EvalRollouts, Stat, TrainRollouts -from prime_rl.orchestrator.types import EvalEpisode, InflightEpisode, Rollout, TrainEpisode, TrainRollout +from prime_rl.orchestrator.types import Episode, InflightEpisode, Rollout, TrainRollout from prime_rl.orchestrator.utils import compute_pass_metrics _ids = count() @@ -82,13 +82,13 @@ def mk( ) -def ep(*rollouts, env_name: str = "env", errors=(), cls=TrainEpisode): +def ep(*rollouts, env_name: str = "env", errors=(), cls=Episode): """One episode over these traces. ``model_construct`` skips validation so the duck-typed stand-ins above can stand in for real ones.""" return cls.model_construct(id=f"e{next(_ids)}", traces=list(rollouts), env_name=env_name, errors=list(errors)) -def solo(rollouts, cls=TrainEpisode): +def solo(rollouts, cls=Episode): """Each rollout as its own single-trace episode — the single-agent shape.""" return [ep(r, cls=cls) for r in rollouts] @@ -274,12 +274,12 @@ def test_train_only_metrics_absent_from_eval(): assert out["train/agg/all/agent/is_filtered/mean"] == 0.5 assert out["train/agg/all/agent/filters/gibberish/mean"] == 0.5 assert "train/agg/all/is_trainable/mean" not in out # pipeline verdicts are per-trace - eval_out = EvalRollouts(solo(rollouts, cls=EvalEpisode)).metrics.to_wandb(prefix="eval/x", subset="all") + eval_out = EvalRollouts(solo(rollouts)).metrics.to_wandb(prefix="eval/x", subset="all") assert not any("is_trainable" in k or "is_filtered" in k or "/filters/" in k for k in eval_out) def test_eval_avg_at_k_and_pass_k(): - binary = EvalRollouts(solo([mk(reward=1.0, group_id="g0"), mk(reward=0.0, group_id="g0")], cls=EvalEpisode)) + binary = EvalRollouts(solo([mk(reward=1.0, group_id="g0"), mk(reward=0.0, group_id="g0")])) eff = binary.effective.metrics.to_wandb(prefix="eval/x", subset="effective") assert eff["eval/x/effective/agent/avg@2"] == 0.5 # mean reward under avg@ (k from the groups) assert "eval/x/effective/avg@2" not in eff # scores are per-agent, never pooled @@ -287,7 +287,7 @@ def test_eval_avg_at_k_and_pass_k(): all_out = binary.metrics.to_wandb(prefix="eval/x", subset="all") assert all_out["eval/x/all/agent/avg@2"] == 0.5 assert not any("pass@" in k or "pass^" in k for k in all_out) # pass@k effective-only - non_binary = EvalRollouts(solo([mk(reward=0.5, group_id="g0"), mk(reward=1.0, group_id="g0")], cls=EvalEpisode)) + non_binary = EvalRollouts(solo([mk(reward=0.5, group_id="g0"), mk(reward=1.0, group_id="g0")])) assert not any("pass@" in k for k in non_binary.effective.metrics.to_wandb(prefix="eval/x", subset="effective")) @@ -312,16 +312,6 @@ def test_traceless_episode_keeps_its_reason(): assert len(pool.effective.episodes) == 1 # it survives nothing, so the subset drops it -def test_train_and_eval_episodes_carry_only_their_own_facts(): - """The path is the type, not a discriminator field: an eval episode always knows its step, and - only a train episode can go stale — neither can be asked for the other's fact.""" - assert (TrainEpisode.KIND, EvalEpisode.KIND) == ("train", "eval") - train, evaluation = TrainEpisode.model_construct(), EvalEpisode.model_construct() - assert train.off_policy_steps == 0 and "off_policy_steps" not in EvalEpisode.model_fields - assert evaluation.step == 0 and "step" not in TrainEpisode.model_fields - assert train.policy_version == evaluation.policy_version == 0 # both measure a policy - - def test_training_state_is_train_only(): """Only a train rollout carries trainer-bound state; an eval trace has no field for it. Credit is not among them — it lives on the graph's nodes, which every trace has.""" @@ -333,19 +323,21 @@ def test_training_state_is_train_only(): def test_inflight_episode_stamps_what_lands(): """The dispatch and the landed episode are one pair: `stamp` is the only place the facts of a - dispatch become facts of an episode, and it mints the class the path calls for.""" + dispatch become facts of an episode, and the run it writes is what tells the rest of the + orchestrator which path the episode is on.""" wire = vf.WireEpisode.model_construct(id="e", traces=[]) inflight = InflightEpisode( kind="train", env_name="rt", group_id=uuid4(), policy_version=3, episodes_owed=1, off_policy_steps=2 ) - train = inflight.stamp(wire, policy_version=7, eval_step=None) - assert isinstance(train, TrainEpisode) and train.KIND == "train" + train = inflight.stamp(wire, run_id="r", policy_version=7, eval_step=None) + assert train.run is not None and train.run.type == "train" + assert train.run.step is None # the batch window it lands in is not known yet assert (train.env_name, train.policy_version, train.off_policy_steps) == ("rt", 7, 2) # group's version wins - evaluation = replace(inflight, kind="eval").stamp(wire, policy_version=7, eval_step=12) - assert isinstance(evaluation, EvalEpisode) and evaluation.step == 12 + evaluation = replace(inflight, kind="eval").stamp(wire, run_id="r", policy_version=7, eval_step=12) + assert evaluation.run is not None and (evaluation.run.type, evaluation.run.step) == ("eval", 12) with pytest.raises(AssertionError): # an eval episode without its step is not representable - replace(inflight, kind="eval").stamp(wire, policy_version=7, eval_step=None) + replace(inflight, kind="eval").stamp(wire, run_id="r", policy_version=7, eval_step=None) def test_empty_is_not_the_same_as_failed(): From c136ce0612f7331ecbd7dc743768e80242e51522 Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Tue, 4 Aug 2026 18:55:48 +0000 Subject: [PATCH 38/58] fix(orchestrator): the eval sink reads its epoch off the run Co-Authored-By: Claude Fable 5 --- src/prime_rl/orchestrator/eval_sink.py | 15 ++++++++++++--- 1 file changed, 12 insertions(+), 3 deletions(-) diff --git a/src/prime_rl/orchestrator/eval_sink.py b/src/prime_rl/orchestrator/eval_sink.py index 9c09e4daa5..564594f485 100644 --- a/src/prime_rl/orchestrator/eval_sink.py +++ b/src/prime_rl/orchestrator/eval_sink.py @@ -17,12 +17,21 @@ import uuid from collections import defaultdict +import verifiers.v1 as vf + from prime_rl.orchestrator.envs import EvalEnvs from prime_rl.orchestrator.metrics import EvalRollouts from prime_rl.orchestrator.types import Episode, EvalBatch, Rollout from prime_rl.utils.logger import get_logger +def eval_step_of(episode: Episode) -> int: + """The eval epoch an episode belongs to, off the run the dispatcher recorded when it landed. + Only the eval path has one, which is why this lives here and not on ``Episode``.""" + assert isinstance(episode.run, vf.EvalRunInfo) and episode.run.step is not None + return episode.run.step + + class EvalSink: """Constructed only when eval is configured.""" @@ -42,7 +51,7 @@ def add(self, episode: Episode) -> EvalBatch | None: group_id = episode.group_id for rollout in episode.rollouts: self.process_rollout(rollout) - bkey = (env_name, episode.step) + bkey = (env_name, eval_step_of(episode)) self.pending_groups[group_id].append(episode) self.pending_group_episodes[group_id] += 1 if self.pending_group_episodes[group_id] >= self.group_size_for(env_name): @@ -73,7 +82,7 @@ def batch_progress(self) -> list[tuple[str, int, int, int, int]]: env_name = episodes[0].env_name if self.eval_envs.get(env_name).requires_group_scoring: continue - bkey = (env_name, episodes[0].step) + bkey = (env_name, eval_step_of(episodes[0])) buffered[bkey] = buffered.get(bkey, 0) + self.pending_group_episodes.get(group_id, 0) return [ ( @@ -105,7 +114,7 @@ def process_group(self, group_id: uuid.UUID) -> None: # Read the group's facts off an episode, not a trace: every episode in it may have # produced none (a whole group cancelled off-policy). env_name = finished[0].env_name - eval_step = finished[0].step + eval_step = eval_step_of(finished[0]) group = [t for e in finished for t in e.rollouts] task_idx = group[0].task.data.idx if group else -1 bucket = self.pending_batches[(env_name, eval_step)] From 46446e3a9f060ce919f8a3a0ae54a724a3c9dc7c Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Tue, 4 Aug 2026 19:05:54 +0000 Subject: [PATCH 39/58] refactor(orchestrator)!: per-agent metrics are built from episodes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit TraceMetrics took a loose trace list, so a trace had to carry the example it answered for solve_rates and pass@k to bucket by. It now takes the episodes narrowed to one agent (Episode.narrow, keyed by the agent names vf.Episode.by_agent reports), which is where group_id already lives — so the field comes off the trace entirely. Co-Authored-By: Claude Fable 5 --- src/prime_rl/orchestrator/dispatcher.py | 1 - src/prime_rl/orchestrator/metrics.py | 42 ++++++++++++++--------- src/prime_rl/orchestrator/orchestrator.py | 2 +- src/prime_rl/orchestrator/types.py | 7 ++-- tests/unit/orchestrator/test_filters.py | 2 -- tests/unit/orchestrator/test_metrics.py | 31 ++++++++++------- 6 files changed, 47 insertions(+), 38 deletions(-) diff --git a/src/prime_rl/orchestrator/dispatcher.py b/src/prime_rl/orchestrator/dispatcher.py index 3cf6cd60f0..25f9f15598 100644 --- a/src/prime_rl/orchestrator/dispatcher.py +++ b/src/prime_rl/orchestrator/dispatcher.py @@ -564,7 +564,6 @@ async def emit_episode( for rollout in episode.traces: rollout.env_name = meta.env_name - rollout.group_id = meta.group_id await self.out_q.put( meta.stamp(episode, run_id=self.run_id, policy_version=policy_version, eval_step=eval_step) ) diff --git a/src/prime_rl/orchestrator/metrics.py b/src/prime_rl/orchestrator/metrics.py index d5a0cf0e08..85bbc60d40 100644 --- a/src/prime_rl/orchestrator/metrics.py +++ b/src/prime_rl/orchestrator/metrics.py @@ -172,7 +172,14 @@ class TraceMetrics(StatGroup): """Trace-level metrics for one agent, every one of them flat over that agent's traces: one sample is one trace, so a fan-out (n same-agent traces in one episode, e.g. n solvers) simply contributes n samples. Weighing whole episodes against each other is the episode level's job; - inside a seat the trace is the unit, which is also what the advantage computation samples.""" + inside a seat the trace is the unit, which is also what the advantage computation samples. + + Built from the episodes narrowed to that agent, not a loose trace list — the episode is the + atomic unit, so anything episode-scoped (which example a trace answered) is still in reach.""" + + def __init__(self, episodes: list[Episode]) -> None: + self.episodes = episodes + super().__init__([t for e in episodes for t in e.rollouts]) DISTRIBUTIONS = ("reward", "num_total_tokens", "num_input_tokens", "num_output_tokens", "num_turns", "num_branches") RATES = ("is_truncated", "is_completed") @@ -227,8 +234,8 @@ def solve_rates(self) -> dict[str, float]: (every trace scored 1.0), and ``solved_some`` (the mixed remainder — the GRPO-signal groups).""" groups: dict = {} - for r in self.rollouts: - groups.setdefault(r.group_id, []).append(r) + for e in self.episodes: + groups.setdefault(e.group_id, []).extend(e.rollouts) n_groups = len(groups) solved_none = sum(1 for g in groups.values() if sum(r.reward for r in g) == 0) solved_all = sum(1 for g in groups.values() if all(r.reward == 1.0 for r in g)) @@ -271,12 +278,13 @@ def __init__(self, episodes: list[Episode]) -> None: self.rollouts: list[Rollout] = [t for e in episodes for t in e.traces] def by_agent(self) -> dict[str, TraceMetrics]: - """Per-agent metric views, merging each episode's own ``vf.Episode.by_agent`` grouping.""" - per_agent: dict[str, list[Rollout]] = {} - for episode in self.episodes: - for name, traces in episode.by_agent.items(): - per_agent.setdefault(name, []).extend(traces) - return {name: TraceMetrics(rollouts) for name, rollouts in sorted(per_agent.items())} + """Per-agent metric views: the episodes narrowed to one agent's traces, keyed by the agent + names each episode's own ``vf.Episode.by_agent`` reports.""" + names = sorted({name for e in self.episodes for name in e.by_agent}) + return { + name: TraceMetrics([n for e in self.episodes if (n := e.narrow(lambda r: r.agent.name == name))]) + for name in names + } # Episode-level count metrics, one value per episode — ``vf.Episode``'s own aggregates. @property @@ -350,14 +358,14 @@ def to_wandb(self, *, prefix: str, subset: Subset) -> dict[str, float]: return out -def pass_at_k(rollouts: list[Rollout]) -> dict[str, float]: +def pass_at_k(episodes: list[Episode]) -> dict[str, float]: """pass@k / pass^k averaged over examples; ``{}`` for non-binary rewards.""" - rewards = [r.reward for r in rollouts] + rewards = [r.reward for e in episodes for r in e.rollouts] if not set(rewards).issubset({0.0, 1.0}): return {} by_example: dict = {} - for r in rollouts: - by_example.setdefault(r.group_id, []).append(r.reward) + for e in episodes: + by_example.setdefault(e.group_id, []).extend(r.reward for r in e.rollouts) per_example = [compute_pass_metrics(rs) for rs in by_example.values()] keys = sorted({k for d in per_example for k in d}) return {k: sum(d[k] for d in per_example if k in d) / sum(1 for d in per_example if k in d) for k in keys} @@ -384,7 +392,7 @@ def to_wandb(self, *, prefix: str, subset: Subset) -> dict[str, float]: p = f"{prefix}/{subset}/{agent}" out[f"{p}/avg@{self.group_size}"] = traces.stats()["reward"].mean() if subset == "effective": - out |= {f"{p}/{k}": v for k, v in pass_at_k(traces.rollouts).items()} + out |= {f"{p}/{k}": v for k, v in pass_at_k(traces.episodes).items()} return out @@ -455,9 +463,9 @@ def group_size(self) -> int: if self._group_size is not None: return self._group_size counts: dict = {} - for r in self.rollouts: - if r.agent.trainable: - counts[r.group_id] = counts.get(r.group_id, 0) + 1 + for e in self.episodes: + trainable = sum(1 for r in e.rollouts if r.agent.trainable) + counts[e.group_id] = counts.get(e.group_id, 0) + trainable return max(counts.values(), default=0) @property diff --git a/src/prime_rl/orchestrator/orchestrator.py b/src/prime_rl/orchestrator/orchestrator.py index 1f8377c8db..af6400ab58 100644 --- a/src/prime_rl/orchestrator/orchestrator.py +++ b/src/prime_rl/orchestrator/orchestrator.py @@ -660,7 +660,7 @@ async def finalize_train_batch(self, batch: TrainBatch) -> None: num_input = sum(r.num_input_tokens for r in effective) num_output = sum(r.num_output_tokens for r in effective) num_rollouts = len(batch.rollouts) - num_unique_examples = len({r.group_id for r in batch.rollouts}) + num_unique_examples = len({e.group_id for e in batch.rollouts.episodes}) metrics |= { "progress/tokens": num_tokens, "progress/input_tokens": num_input, diff --git a/src/prime_rl/orchestrator/types.py b/src/prime_rl/orchestrator/types.py index 4c34dea188..bdbae307eb 100644 --- a/src/prime_rl/orchestrator/types.py +++ b/src/prime_rl/orchestrator/types.py @@ -61,9 +61,9 @@ class GroupState: class Rollout(vf.Trace[DataT], Generic[DataT]): """A completed rollout: the env's typed ``vf.Trace`` *is* the rollout, carrying only the links - prime-rl needs to place a loose trace back among its peers — its episode, its comparison group, - its env. Everything about the dispatch itself lives on the ``Episode``, which is the thing that - was dispatched. All added fields are ``exclude=True``, so dumping a Rollout yields a plain + a consumer that works in loose traces — the sample monitors — needs to place one back among its + peers. Anything episode-scoped is read off the ``Episode``, which is the atomic unit everything + else passes around. All added fields are ``exclude=True``, so dumping a Rollout yields a plain trace on the wire; ``vf.Trace.record_run`` mirrors them into ``info`` on arrival so the on-disk records stay fully placeable. @@ -71,7 +71,6 @@ class Rollout(vf.Trace[DataT], Generic[DataT]): (``rollout.reward``, ``rollout.nodes``, ``rollout.num_turns``).""" env_name: str = Field(default="", exclude=True) - group_id: uuid.UUID = Field(default_factory=uuid.uuid4, exclude=True) # Links the traces of one episode; stamped into ``info`` on arrival so # saved records keep their grouping. episode_id: str = Field(default="", exclude=True) diff --git a/tests/unit/orchestrator/test_filters.py b/tests/unit/orchestrator/test_filters.py index c344de9010..4d35c842b6 100644 --- a/tests/unit/orchestrator/test_filters.py +++ b/tests/unit/orchestrator/test_filters.py @@ -1,5 +1,4 @@ import math -import uuid import verifiers.v1 as vf @@ -64,7 +63,6 @@ def _make_rollout( rewards={"reward": vf.Reward(score=reward)}, ) rollout.env_name = "test" - rollout.group_id = uuid.uuid4() return rollout diff --git a/tests/unit/orchestrator/test_metrics.py b/tests/unit/orchestrator/test_metrics.py index f79abaecd0..ee0b2f406a 100644 --- a/tests/unit/orchestrator/test_metrics.py +++ b/tests/unit/orchestrator/test_metrics.py @@ -32,7 +32,6 @@ def mk( metrics: dict | None = None, rewards: dict | None = None, env_name: str = "env", - group_id: str = "g0", trainable: bool = True, is_trainable: bool = True, is_filtered: bool = False, @@ -64,7 +63,6 @@ def mk( stop_condition=stop_condition, metrics=metrics or {}, env_name=env_name, - group_id=group_id, agent=SimpleNamespace(trainable=trainable, name=agent_name), is_trainable=is_trainable, is_filtered=is_filtered, @@ -82,19 +80,23 @@ def mk( ) -def ep(*rollouts, env_name: str = "env", errors=(), cls=Episode): +def ep(*rollouts, env_name: str = "env", errors=(), group_id="g0", cls=Episode): """One episode over these traces. ``model_construct`` skips validation so the duck-typed stand-ins above can stand in for real ones.""" - return cls.model_construct(id=f"e{next(_ids)}", traces=list(rollouts), env_name=env_name, errors=list(errors)) + return cls.model_construct( + id=f"e{next(_ids)}", traces=list(rollouts), env_name=env_name, errors=list(errors), group_id=group_id + ) -def solo(rollouts, cls=Episode): - """Each rollout as its own single-trace episode — the single-agent shape.""" - return [ep(r, cls=cls) for r in rollouts] +def solo(rollouts, group_ids=None, cls=Episode): + """Each rollout as its own single-trace episode — the single-agent shape. ``group_ids`` gives + the example each answers; by default they all answer the same one.""" + ids = group_ids or ["g0"] * len(rollouts) + return [ep(r, group_id=g, cls=cls) for r, g in zip(rollouts, ids, strict=True)] -def train_wandb(rollouts, subset: str = "all") -> dict: - return TrainRollouts(solo(rollouts)).metrics.to_wandb(prefix="train/agg", subset=subset) +def train_wandb(rollouts, subset: str = "all", group_ids=None) -> dict: + return TrainRollouts(solo(rollouts, group_ids)).metrics.to_wandb(prefix="train/agg", subset=subset) def test_stat(): @@ -204,7 +206,9 @@ def test_boolean_rates_and_error_breakdown_all_only(): def test_solve_rates(): groups = {"A": [1.0, 1.0], "B": [0.0, 0.0], "C": [1.0, 0.0], "D": [1.0, 0.0]} # all / none / some / some - out = train_wandb([mk(reward=r, group_id=g) for g, rs in groups.items() for r in rs]) + out = train_wandb( + [mk(reward=r) for _, rs in groups.items() for r in rs], group_ids=[g for g, rs in groups.items() for _ in rs] + ) rates = ( out["train/agg/all/agent/solved_all"], out["train/agg/all/agent/solved_none"], @@ -279,7 +283,7 @@ def test_train_only_metrics_absent_from_eval(): def test_eval_avg_at_k_and_pass_k(): - binary = EvalRollouts(solo([mk(reward=1.0, group_id="g0"), mk(reward=0.0, group_id="g0")])) + binary = EvalRollouts(solo([mk(reward=1.0), mk(reward=0.0)])) eff = binary.effective.metrics.to_wandb(prefix="eval/x", subset="effective") assert eff["eval/x/effective/agent/avg@2"] == 0.5 # mean reward under avg@ (k from the groups) assert "eval/x/effective/avg@2" not in eff # scores are per-agent, never pooled @@ -287,7 +291,7 @@ def test_eval_avg_at_k_and_pass_k(): all_out = binary.metrics.to_wandb(prefix="eval/x", subset="all") assert all_out["eval/x/all/agent/avg@2"] == 0.5 assert not any("pass@" in k or "pass^" in k for k in all_out) # pass@k effective-only - non_binary = EvalRollouts(solo([mk(reward=0.5, group_id="g0"), mk(reward=1.0, group_id="g0")])) + non_binary = EvalRollouts(solo([mk(reward=0.5), mk(reward=1.0)])) assert not any("pass@" in k for k in non_binary.effective.metrics.to_wandb(prefix="eval/x", subset="effective")) @@ -318,7 +322,8 @@ def test_training_state_is_train_only(): assert {"samples", "is_filtered", "filter_results"} <= set(TrainRollout.model_fields) assert not {"samples", "is_filtered", "filter_results"} & set(Rollout.model_fields) assert "advantages" not in TrainRollout.model_fields # derived from the nodes - assert {"env_name", "group_id", "episode_id"} <= set(Rollout.model_fields) # links stay shared + assert {"env_name", "episode_id"} <= set(Rollout.model_fields) # links stay shared + assert "group_id" not in Rollout.model_fields # the example a trace answered is the episode's def test_inflight_episode_stamps_what_lands(): From c1f33d3cf8109fccdb13709e99325a465bd39741 Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Tue, 4 Aug 2026 20:12:19 +0000 Subject: [PATCH 40/58] feat(orchestrator)!: split degeneracy detection from the drop policy MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Gibberish and repetition were filters with an enforce flag, which meant their monitoring numbers were censored: apply_filters stopped at the first hit, so a rollout flagged as gibberish was never measured for repetition. They are now detectors — pure per-trace measurements that all run on every trace and report per agent. Zero-advantage stops being a plugin. It is not a property of the generation (it is only knowable after the group scores, and it is what is_trainable already means), so it becomes the drop policy's built-in default. Rollouts that were never scored stay: opd/opsd assign no credit and train through reference KL. Co-Authored-By: Claude Fable 5 --- configs/ci/nightly-fft/wiki-search.toml | 4 - configs/debug/algo/echo.toml | 10 +- docs/algorithms.md | 67 ++- examples/advanced/glm-5.2/swe.toml | 5 +- examples/advanced/intellect-3.1/rl.toml | 4 - examples/basic/wiki-search/rl.toml | 4 - .../src/prime_rl/configs/orchestrator.py | 70 +-- src/prime_rl/orchestrator/algo/base.py | 6 +- src/prime_rl/orchestrator/detectors.py | 129 ++++++ src/prime_rl/orchestrator/filters.py | 172 -------- src/prime_rl/orchestrator/metrics.py | 8 +- src/prime_rl/orchestrator/orchestrator.py | 25 +- src/prime_rl/orchestrator/train_sink.py | 78 ++-- src/prime_rl/orchestrator/types.py | 5 +- tests/unit/orchestrator/test_detectors.py | 269 ++++++++++++ tests/unit/orchestrator/test_filters.py | 406 ------------------ tests/unit/orchestrator/test_metrics.py | 16 +- 17 files changed, 536 insertions(+), 742 deletions(-) create mode 100644 src/prime_rl/orchestrator/detectors.py delete mode 100644 src/prime_rl/orchestrator/filters.py create mode 100644 tests/unit/orchestrator/test_detectors.py delete mode 100644 tests/unit/orchestrator/test_filters.py diff --git a/configs/ci/nightly-fft/wiki-search.toml b/configs/ci/nightly-fft/wiki-search.toml index 9814e6351b..62b40a2767 100644 --- a/configs/ci/nightly-fft/wiki-search.toml +++ b/configs/ci/nightly-fft/wiki-search.toml @@ -19,10 +19,6 @@ batch_size = 512 group_size = 16 oversampling_factor = 2.0 -[[orchestrator.pre_batch_filters]] -type = "zero_advantage" -enforce = true - [[orchestrator.train.source]] name = "wiki-search" diff --git a/configs/debug/algo/echo.toml b/configs/debug/algo/echo.toml index c05beb865e..1e66648c8d 100644 --- a/configs/debug/algo/echo.toml +++ b/configs/debug/algo/echo.toml @@ -16,6 +16,10 @@ name = "debug-echo" batch_size = 32 group_size = 4 +# ECHO learns from observation tokens even when the GRPO advantage collapses +# to zero — keep zero-advantage rollouts in the batch. +drop_zero_advantage = false + # alphabet-sort's feedback arrives as user messages, so train the user role # instead of echo's tool default. [orchestrator.algo] @@ -44,12 +48,6 @@ type = "subprocess" [orchestrator.train.sampling] max_completion_tokens = 512 -# ECHO learns from observation tokens even when the GRPO advantage collapses -# to zero — keep zero-advantage rollouts in the batch. -[[orchestrator.post_batch_filters]] -type = "zero_advantage" -enforce = false - # Fine-tune inherits the PrimeIntellect Qwen3 template byte-for-byte. [orchestrator.renderer] name = "prime-qwen3" diff --git a/docs/algorithms.md b/docs/algorithms.md index a2bcd1e7db..64d38699ad 100644 --- a/docs/algorithms.md +++ b/docs/algorithms.md @@ -1,6 +1,6 @@ # Algorithms -This page covers the math and the configurable algorithmic components: the algorithm abstraction and its algorithms, how off-policy training works, the loss components and advantage functions, how to plug in your own, the filters applied between rollout and training, and how multi-turn rollouts get merged into training samples. +This page covers the math and the configurable algorithmic components: the algorithm abstraction and its algorithms, how off-policy training works, the loss components and advantage functions, how to plug in your own, the degeneracy detectors and drop policy applied between rollout and training, and how multi-turn rollouts get merged into training samples. ## Table of Contents @@ -21,7 +21,7 @@ This page covers the math and the configurable algorithmic components: the algor - [Self-Play Advantage (RAE)](#self-play-advantage-rae) - [Authoring an Algorithm](#authoring-an-algorithm) - [Reference Scoring](#reference-scoring) -- [Filters](#filters) +- [Detectors and the drop policy](#detectors-and-the-drop-policy) - [Multi-Turn Trajectories](#multi-turn-trajectories) - [Extension Property](#extension-property) - [Best-Effort Interleaving](#best-effort-interleaving) @@ -164,12 +164,12 @@ At runtime, each env's resolved config builds two objects: a `Sampler` (`prime_r | `hierarchical_grpo` | `HierarchicalGRPOAlgorithm` | `score_group`: GRPO baseline per episode for solvers, per group for the proposer | | `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) | +| `sft` | `SFTDistillAlgorithm` | `score_group`: group-norm credit (feeds the drop policy) | Each class owns its hooks outright — reading one top to bottom reads the algorithm, and everything on the class is an override point. The two hooks are one scope-and-timing ladder — the wider scope is unlocked by a later barrier, so the two axes coincide. Each is handed the env's own data directly — a `Rollout` (the typed trace: `reward`, `nodes`, `num_turns`, ... with `samples` attached) on arrival, the group's `Episode`s at group time — plus `assign_advantages` to write credit: -- `async score_rollout(rollout)` — one rollout, **on arrival** (as it's tokenized, before its group is complete): rollout-local credit (`rollout.assign_advantages(...)`), observation ce weights, **or** model I/O — query a reference pool (e.g. `self.teacher_pool`, connected in `setup()` via `self.connect(...)`, or the live `self.policy_pool` for opsd) and attach per-token results (e.g. teacher logprobs) with bounded concurrency. No siblings. `echo` weights observation tokens here, identifying env-provided observation nodes by their non-sampled status and source step role attribution, applying the optional user filter, and writing the `ce_weights` stream. Model I/O runs *before* the pre-batch filters, so it pays compute on rollouts that may then be filtered out. -- `score_group(group)` — the cohort, **before filtering** (filters read the streams), synchronous: group-relative credit (GRPO/MaxRL baselines). `group` is a list of `TrainEpisode`, so an algorithm can compare within an episode as well as across them (`hierarchical_grpo` does); `group_rollouts(group)` flattens it for the ones that only compare across. +- `async score_rollout(rollout)` — one rollout, **on arrival** (as it's tokenized, before its group is complete): rollout-local credit (`rollout.assign_advantages(...)`), observation ce weights, **or** model I/O — query a reference pool (e.g. `self.teacher_pool`, connected in `setup()` via `self.connect(...)`, or the live `self.policy_pool` for opsd) and attach per-token results (e.g. teacher logprobs) with bounded concurrency. No siblings. `echo` weights observation tokens here, identifying env-provided observation nodes by their non-sampled status and source step role attribution, applying the optional user filter, and writing the `ce_weights` stream. Model I/O runs *before* the drop policy, so it pays compute on rollouts that may then be dropped. +- `score_group(group)` — the cohort, **before the drop policy** (which reads the streams), synchronous: group-relative credit (GRPO/MaxRL baselines). `group` is a list of `TrainEpisode`, so an algorithm can compare within an episode as well as across them (`hierarchical_grpo` does); `group_rollouts(group)` flattens it for the ones that only compare across. The pipeline drives the hooks through two non-virtual methods it never looks inside: `algorithm.finalize_rollout(rollout)` per arrival (rollout-local scoring + reference I/O) and `algorithm.finalize_group(episodes)` per group (scoring + wire stamping; after this the records are frozen — groups die at stamping). Sample construction (interleaving) is pure pipeline — observation-token provenance is available through structural attribution (`node.sampled`, `node.is_content`) for any algorithm that trains on env-provided tokens. @@ -305,8 +305,8 @@ The per-token training signal is set by `algo.type` and the [algorithm](#the-alg | `rae` | `rl` | Reward minus a per-agent EMA baseline (SPIRAL's role-conditioned advantage estimation) — for multi-agent self-play envs. | | `hierarchical_grpo` | `rl` | GRPO for proposer-solver envs: solvers are compared within one proposed problem, while proposers are compared across proposals. | | `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 (`teacher`, 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. | +| `opd` | `ref_kl` | On-policy distillation: per-token reverse KL to a reference model (`teacher`, an inline frozen hosted model), evaluated in the trainer from shipped reference logprobs. No credit — rollouts keep `advantages = None` (so zero-advantage dropping never applies) 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` (so zero-advantage dropping never applies) and ship no advantage stream. | | `sft` | `ce` | Cross-entropy on the sampled tokens. Assigns no advantage — trains on every sampled token. | ### Default Advantage @@ -441,7 +441,7 @@ class MyAlgorithm(Algorithm): Add a typed `MyAlgoConfig` to `prime_rl.configs.algorithm` and its discriminated union, then register `"my_algo": MyAlgorithm` in `ALGORITHM_CLASSES`. Pick the hook by *when* your signal is ready: `score_rollout` for per-arrival credit or credit that needs a model call (it's `async`), `score_group` for group-relative credit (GRPO/MaxRL). `assign_advantages` takes a scalar and broadcasts it over the rollout's trainable tokens, writing it onto the graph's nodes so branches sharing a node cannot disagree about its credit. -Advantage-based filters and metrics derive from the streams (the zero-advantage filter checks for all-zero streams; logged distributions use per-rollout means). Signals that depend on the live policy's weights (like OPD's reverse KL) cannot be precomputed here; those are reference-scoring algorithms, evaluated in the trainer. +The drop policy and metrics derive from the streams (zero-advantage dropping checks for an all-zero stream; logged distributions use per-rollout means). Signals that depend on the live policy's weights (like OPD's reverse KL) cannot be precomputed here; those are reference-scoring algorithms, evaluated in the trainer. ### Reference Scoring @@ -456,30 +456,51 @@ type = "opsd" demo_key = "demonstration" ``` -Scoring runs at arrival, *before* the pre-batch filters, so a rollout that is later filtered still cost its reference compute — accepted for the simpler one-rollout-at-a-time shape (advantage-based filters never fire for opd/opsd anyway, since neither assigns an advantage). +Scoring runs at arrival, *before* the drop policy, so a rollout that is later dropped still cost its reference compute — accepted for the simpler one-rollout-at-a-time shape (zero-advantage dropping never applies to opd/opsd anyway, since neither assigns an advantage). -## Filters +## Detectors and the drop policy -Filters drop rollouts between scoring and training. Built-ins (composable): +Two separate things: **measuring** what a rollout looks like, and **deciding** whether to train on it. -| Filter | Effect | -|---|---| -| `gibberish` | Drops rollouts whose mean log-prob fall below a threshold — usually a sign of degenerate output. | -| `repetition` | Drops rollouts with high n-gram repetition. | -| `zero_advantage` | Drops rollouts whose advantage is zero, so the trainer doesn't waste tokens on them. | +Detectors measure. Each one asks a single question of a trace's tokens, runs on every trace as soon +as it is tokenized, and reports its rate per agent alongside reward and truncation — whether or not +anything acts on it. Every detector measures every trace, so one detection never hides another. -The default `[orchestrator]` config registers all three in both filter slots: `post_batch_filters` enforce by default (flagged rollouts are recorded but not shipped to the trainer), while `pre_batch_filters` run in monitor mode (`enforce = false`); flip `enforce = true` there to drop matching rollouts before they consume a slot in the batch. Setting a slot replaces its defaults wholesale: +| Detector | Measures | +|---|---| +| `gibberish` | rare tokens (high BPE id) generated at high entropy — degenerate output | +| `repetition` | a long stretch of very-high-confidence tokens — a repetition loop | ```toml -[[orchestrator.post_batch_filters]] -type = "zero_advantage" +[orchestrator.detectors] +drop = ["gibberish"] # measured always; this says which detections also drop -[[orchestrator.post_batch_filters]] -type = "repetition" -threshold = 0.4 +[orchestrator.detectors.repetition] +window = 2000 ``` -Filtered rollouts still appear in W&B distributions, just not in the trainer batch — useful for spotting whether filtering is doing its job. +Set a detector to `false` to stop measuring it (`[orchestrator.detectors] repetition = false`). + +The drop policy decides. It runs once, when a finalized group's credit is assigned and the rollouts +would enter the batch buffer, and it has two inputs: the detections the run listed in +`detectors.drop`, and **zero credit**, which drops on its own: + +```toml +[orchestrator] +drop_zero_advantage = true # the default +``` + +A scored rollout whose every token is worth nothing — a GRPO group where all rollouts earned the +same reward — produces no gradient, so training on it is a wasted forward pass. It is not a plugin +because it is not a property of the generation: it is only knowable after the group is scored, and +it is what `TrainRollout.is_trainable` already means. + +A rollout that was **never** scored is not zero-credit. `opd` / `opsd` assign no advantages at all +and train through reference KL, so they are never dropped by this rule. + +Dropped rollouts still appear in the metrics window and the `all` trace file — they just don't ship. +`{scope}/{subset}//detected//mean` is each detector's rate, and +`{scope}/{subset}//is_filtered/mean` is the share the policy dropped. ## Multi-Turn Trajectories diff --git a/examples/advanced/glm-5.2/swe.toml b/examples/advanced/glm-5.2/swe.toml index c397f717f6..a9133b475b 100644 --- a/examples/advanced/glm-5.2/swe.toml +++ b/examples/advanced/glm-5.2/swe.toml @@ -84,9 +84,8 @@ id = "bash" type = "prime" labels = ["glm5-pd-disag", "swe-bench-verified"] -[[orchestrator.post_batch_filters]] -type = "gibberish" -enforce = true +[orchestrator.detectors] +drop = ["gibberish"] [inference] enable_expert_parallel = true diff --git a/examples/advanced/intellect-3.1/rl.toml b/examples/advanced/intellect-3.1/rl.toml index 174639662f..a22f5ab7fd 100644 --- a/examples/advanced/intellect-3.1/rl.toml +++ b/examples/advanced/intellect-3.1/rl.toml @@ -119,10 +119,6 @@ id = "null" [orchestrator.train.source.env.agent.runtime] type = "subprocess" -[[orchestrator.pre_batch_filters]] -type = "zero_advantage" -enforce = true - [orchestrator.eval] interval = 25 diff --git a/examples/basic/wiki-search/rl.toml b/examples/basic/wiki-search/rl.toml index 8f9d37945a..068f0868ba 100644 --- a/examples/basic/wiki-search/rl.toml +++ b/examples/basic/wiki-search/rl.toml @@ -40,10 +40,6 @@ name = "qwen3-4b-wiki-search" [orchestrator.train.sampling] max_completion_tokens = 512 -[[orchestrator.pre_batch_filters]] -type = "zero_advantage" -enforce = true - [[orchestrator.train.source]] name = "wiki-search" diff --git a/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py b/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py index 3434786721..0b625999e0 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py @@ -357,11 +357,8 @@ class CheckpointConfig(BaseConfig): # Flags rare tokens generated at high entropy (Section 5.2, https://arxiv.org/abs/2510.02387). -class GibberishFilterConfig(BaseConfig): - type: Literal["gibberish"] = "gibberish" - - enforce: bool = False - """When True, skip detected rollouts entirely so they are not sent to the trainer. When False, only track detection metrics.""" +class GibberishDetectorConfig(BaseConfig): + """Flags rare tokens generated at high entropy (Section 5.2, https://arxiv.org/abs/2510.02387).""" token_id_threshold: int = 100_000 """Token IDs above this are candidates for gibberish. BPE tokens are sorted by merge order.""" @@ -370,14 +367,9 @@ class GibberishFilterConfig(BaseConfig): """Offset from uniform-distribution logprob. Threshold = ``-log(vocab_size) - logprob_offset``.""" -# Flags rollouts stuck in a repetition loop: emits high-confidence tokens for an extended stretch. -# Flagged when `window` consecutive tokens are each sampled with probability above `prob_threshold`. -# (Section 3.2, https://arxiv.org/abs/2506.13585) -class RepetitionFilterConfig(BaseConfig): - type: Literal["repetition"] = "repetition" - - enforce: bool = False - """When True, skip detected rollouts entirely so they are not sent to the trainer. When False, only track detection metrics.""" +class RepetitionDetectorConfig(BaseConfig): + """Flags rollouts stuck in a repetition loop: high-confidence tokens for an extended stretch + (Section 3.2, https://arxiv.org/abs/2506.13585).""" window: int = Field(3_000, ge=1) """Consecutive high-probability steps required to flag the rollout.""" @@ -386,18 +378,15 @@ class RepetitionFilterConfig(BaseConfig): """Tokens sampled with probability above this are considered repetitive. Consecutive such tokens count toward the window.""" -# Flags rollouts with zero advantage. -class ZeroAdvantageFilterConfig(BaseConfig): - type: Literal["zero_advantage"] = "zero_advantage" - - enforce: bool = True - """When True, skip detected rollouts entirely so they are not sent to the trainer. When False, only track detection metrics.""" +class DetectorsConfig(BaseConfig): + """Degeneracy detectors. Each one measures every trace and reports its rate per agent; set + ``drop`` to also keep what it flags out of the training batch.""" + gibberish: GibberishDetectorConfig | None = GibberishDetectorConfig() + repetition: RepetitionDetectorConfig | None = RepetitionDetectorConfig() -FilterConfig: TypeAlias = Annotated[ - GibberishFilterConfig | RepetitionFilterConfig | ZeroAdvantageFilterConfig, - Field(discriminator="type"), -] + drop: list[Literal["gibberish", "repetition"]] = [] + """Which detections keep a rollout out of the batch. Measuring is always on; dropping is not.""" class FileSystemWeightBroadcastConfig(BaseConfig): @@ -473,23 +462,14 @@ class OrchestratorConfig(BaseConfig): eval: EvalConfig | None = None """Evaluation configuration.""" - pre_batch_filters: list[FilterConfig] = [ - GibberishFilterConfig(enforce=False), - RepetitionFilterConfig(enforce=False), - ZeroAdvantageFilterConfig(enforce=False), - ] - """Filters applied *before* a rollout enters the training batch buffer. - All three filter types are registered in monitor mode by default; flip ``enforce=true`` per type - to drop matching rollouts before they consume a slot in the batch (e.g. a zero-advantage group - never makes it into a training batch).""" - - post_batch_filters: list[FilterConfig] = [ - GibberishFilterConfig(), - RepetitionFilterConfig(), - ZeroAdvantageFilterConfig(), - ] - """Filters applied *after* a batch has been assembled. Each filter annotates each rollout; - rollouts flagged by an enforcing filter are still recorded but not shipped to the trainer.""" + detectors: DetectorsConfig = DetectorsConfig() + """Degeneracy detectors, measured on every trace and reported per agent. ``detectors.drop`` + says which detections also keep a rollout out of the training batch.""" + + drop_zero_advantage: bool = True + """Keep scored rollouts whose credit is all zero out of the batch — a GRPO group that earned a + uniform reward carries no gradient, so training on it is wasted compute. Rollouts that were + never scored (opd/opsd train through reference KL, not credit) are unaffected.""" log: LogConfig = LogConfig() @@ -582,16 +562,6 @@ def auto_setup_prime_monitor_run_name(self): self.prime_monitor.run_name = self.wandb.name return self - @model_validator(mode="after") - def validate_unique_filter_types(self): - for slot_name in ("pre_batch_filters", "post_batch_filters"): - types = [f.type for f in getattr(self, slot_name)] - if len(types) != len(set(types)): - raise ValueError( - f"Duplicate filter types in {slot_name}: {types}. Each filter type may only appear once per slot." - ) - return self - @model_validator(mode="after") def inherit_env_algorithms(self): """Envs without their own algorithm inherit the top-level one. diff --git a/src/prime_rl/orchestrator/algo/base.py b/src/prime_rl/orchestrator/algo/base.py index a7962c0327..3d578f8be6 100644 --- a/src/prime_rl/orchestrator/algo/base.py +++ b/src/prime_rl/orchestrator/algo/base.py @@ -107,11 +107,11 @@ class Algorithm: nothing. - :meth:`score_group` — the cohort of episodes, *before* filtering (filters read the streams): group-relative credit. Default: nothing — - rollouts keep ``advantages=None``, so advantage-based filters skip them. + rollouts keep ``advantages=None``, which the drop policy reads as unscored, + not as zero credit. Model I/O lives in :meth:`score_rollout`: it runs at arrival, *before* the - pre-batch filters, so it pays compute on rollouts that may then be filtered - out — accepted for the simpler one-rollout-at-a-time shape. + drop policy, so it pays compute on rollouts that may then be dropped — accepted for the simpler one-rollout-at-a-time shape. Constructed with the algorithm config it interprets plus the live policy pool (``self.policy_pool`` — always available, never closed by the diff --git a/src/prime_rl/orchestrator/detectors.py b/src/prime_rl/orchestrator/detectors.py new file mode 100644 index 0000000000..4a0f56d797 --- /dev/null +++ b/src/prime_rl/orchestrator/detectors.py @@ -0,0 +1,129 @@ +"""Degeneracy detectors: per-trace measurements of pathological generation. + +A detector answers one question about one trace's tokens — is it gibberish, is it stuck in a +repetition loop — and nothing else. Every configured detector runs on every trace as soon as it is +tokenized, so its rate is a trace metric like reward or truncation, reported per agent whether or +not anything acts on it. + +What to *do* about a detection is a separate decision, made once when the batch is assembled +(``prime_rl.orchestrator.train_sink``). Keeping the two apart is what lets every detector measure +every trace: a policy that stops at the first hit would leave the rest unmeasured. +""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import TYPE_CHECKING, Protocol + +from prime_rl.configs.orchestrator import DetectorsConfig +from prime_rl.utils.logger import get_logger + +if TYPE_CHECKING: + from prime_rl.orchestrator.types import TrainRollout + + +class Detector(Protocol): + name: str + + def detect(self, rollout: TrainRollout) -> bool: ... + + +@dataclass +class GibberishDetector: + """Rare tokens generated at high entropy. + + A token counts when both: + - id(token) > token_id_threshold (rare BPE token) + - logprob(token) < -log(vocab_size) - logprob_offset (high entropy) + + References: + Section 5.2, https://arxiv.org/abs/2510.02387 + """ + + name: str + token_id_threshold: int + logprob_threshold: float + + def detect(self, rollout: TrainRollout) -> bool: + for branch in rollout.branches: + # branch.{token_ids,logprobs,sampled_mask} are flat and mutually aligned; the raw + # node arrays are not (node.logprobs covers only the sampled suffix, not the + # generation-prompt scaffold that token_ids/mask also span). + for token_id, logprob, sampled in zip(branch.token_ids, branch.logprobs, branch.sampled_mask): + if sampled and token_id > self.token_id_threshold and logprob < self.logprob_threshold: + return True + return False + + +@dataclass +class RepetitionDetector: + """A repetition loop: a long stretch of very-high-confidence tokens. + + Counts consecutive tokens with logprob > log(prob_threshold); a streak reaching ``window`` + is a detection. + + References: + Section 3.2, https://arxiv.org/abs/2506.13585 + """ + + name: str + window: int + logprob_threshold: float + + def detect(self, rollout: TrainRollout) -> bool: + for branch in rollout.branches: + # Aligned branch streams (see GibberishDetector), and reset the streak per branch: + # flat rollout.nodes interleaves distinct root->leaf paths (compaction/subagents), + # so a per-node walk would run a streak across a branch boundary. + consecutive = 0 + for logprob, sampled in zip(branch.logprobs, branch.sampled_mask): + if not sampled: + continue + consecutive = consecutive + 1 if logprob > self.logprob_threshold else 0 + if consecutive >= self.window: + return True + return False + + +def setup_detectors(config: DetectorsConfig, vocab_size: int) -> list[Detector]: + detectors: list[Detector] = [] + if config.gibberish is not None: + detectors.append( + GibberishDetector( + name="gibberish", + token_id_threshold=config.gibberish.token_id_threshold, + logprob_threshold=-math.log(vocab_size) - config.gibberish.logprob_offset, + ) + ) + if config.repetition is not None: + detectors.append( + RepetitionDetector( + name="repetition", + window=config.repetition.window, + logprob_threshold=math.log(config.repetition.prob_threshold), + ) + ) + if detectors: + get_logger().info(f"Measuring {len(detectors)} degeneracy detector(s): {', '.join(d.name for d in detectors)}") + if config.drop: + get_logger().info(f"Dropping detected rollouts: {', '.join(sorted(config.drop))}") + return detectors + + +def detect(detectors: list[Detector], rollout: TrainRollout) -> None: + """Measure every detector on one trace, writing the verdicts to ``rollout.detections``.""" + rollout.detections = {d.name: d.detect(rollout) for d in detectors} + + +def drop_reasons(rollout: TrainRollout, *, drop_detections: list[str], drop_zero_advantage: bool) -> list[str]: + """Why this rollout should not be trained on, if anything. + + A detection only drops when the run asked it to — measuring is always on, acting is opt-in. + Zero credit drops on its own: a scored rollout whose every token is worth nothing produces no + gradient, so the forward pass is wasted. A rollout that was never scored is *not* zero-credit — + opd/opsd train through reference KL and assign no advantages at all, so they must survive.""" + reasons = [name for name in drop_detections if rollout.detections.get(name)] + if drop_zero_advantage and rollout.advantages is not None and not rollout.is_trainable: + reasons.append("zero_advantage") + return reasons diff --git a/src/prime_rl/orchestrator/filters.py b/src/prime_rl/orchestrator/filters.py deleted file mode 100644 index 01f241c1c0..0000000000 --- a/src/prime_rl/orchestrator/filters.py +++ /dev/null @@ -1,172 +0,0 @@ -"""Orchestrator-side rollout filters for detecting degenerate generations. - -Filters run after rollouts complete, inspecting token IDs and logprobs to -detect gibberish or repetition. Detection metrics are always tracked. -When enforce=True, detected rollouts are skipped entirely during training and -are not sent to the trainer. Reward is kept as-is for baseline calculation. -""" - -from __future__ import annotations - -import math -from dataclasses import dataclass -from typing import TYPE_CHECKING, Protocol - -from prime_rl.configs.orchestrator import FilterConfig -from prime_rl.utils.logger import get_logger - -if TYPE_CHECKING: - from prime_rl.orchestrator.types import TrainRollout - - -@dataclass -class FilterResult: - detected: bool - - -class RolloutFilter(Protocol): - name: str - enforce: bool - - def check(self, rollout: TrainRollout) -> FilterResult: ... - - -@dataclass -class GibberishFilter: - """Flags rollouts containing rare tokens generated at high entropy. - - A token is flagged when both: - - id(token) > token_id_threshold (rare BPE token) - - logprob(token) < -log(vocab_size) - logprob_offset (high entropy) - - References: - Section 5.2, https://arxiv.org/abs/2510.02387 - """ - - name: str - token_id_threshold: int - logprob_threshold: float - enforce: bool = False - - def check(self, rollout: TrainRollout) -> FilterResult: - for branch in rollout.branches: - # branch.{token_ids,logprobs,sampled_mask} are flat and mutually aligned; the raw - # node arrays are not (node.logprobs covers only the sampled suffix, not the - # generation-prompt scaffold that token_ids/mask also span). - for token_id, logprob, sampled in zip(branch.token_ids, branch.logprobs, branch.sampled_mask): - if not sampled: - continue - if token_id > self.token_id_threshold and logprob < self.logprob_threshold: - return FilterResult(detected=True) - return FilterResult(detected=False) - - -@dataclass -class RepetitionFilter: - """Flags rollouts with pathological repetition loops. - - Counts consecutive tokens where logprob > log(prob_threshold), indicating - the model is generating with very high confidence. When the streak reaches - the window size, the rollout is flagged. - - References: - Section 3.2, https://arxiv.org/abs/2506.13585 - """ - - name: str - window: int - logprob_threshold: float - enforce: bool = False - - def check(self, rollout: TrainRollout) -> FilterResult: - for branch in rollout.branches: - # Aligned branch streams (see GibberishFilter), and reset the streak per branch: - # flat rollout.nodes interleaves distinct root->leaf paths (compaction/subagents), - # so a per-node walk would run a streak across a branch boundary. - consecutive = 0 - for logprob, sampled in zip(branch.logprobs, branch.sampled_mask): - if not sampled: - continue - if logprob > self.logprob_threshold: - consecutive += 1 - else: - consecutive = 0 - if consecutive >= self.window: - return FilterResult(detected=True) - return FilterResult(detected=False) - - -@dataclass -class ZeroAdvantageFilter: - """Flags rollouts whose advantage stream is all zero (e.g. all rollouts in - a GRPO group earned the same reward, so the centered advantage collapses).""" - - name: str - enforce: bool = True - - def check(self, rollout: TrainRollout) -> FilterResult: - if rollout.advantages is not None and all(a == 0.0 for a in rollout.advantages): - return FilterResult(detected=True) - return FilterResult(detected=False) - - -def setup_filter(config: FilterConfig, vocab_size: int) -> RolloutFilter: - """Create a RolloutFilter from a filter config.""" - if config.type == "gibberish": - return GibberishFilter( - name="gibberish", - token_id_threshold=config.token_id_threshold, - logprob_threshold=-math.log(vocab_size) - config.logprob_offset, - enforce=config.enforce, - ) - elif config.type == "repetition": - return RepetitionFilter( - name="repetition", - window=config.window, - logprob_threshold=math.log(config.prob_threshold), - enforce=config.enforce, - ) - elif config.type == "zero_advantage": - return ZeroAdvantageFilter( - name="zero_advantage", - enforce=config.enforce, - ) - raise ValueError(f"Unknown filter type: {config.type}") - - -def setup_filters(configs: list[FilterConfig], vocab_size: int, *, kind: str) -> list[RolloutFilter]: - """Create RolloutFilters from a list of filter configs.""" - filters = [setup_filter(config, vocab_size) for config in configs] - if filters: - get_logger().info(f"Configured {len(filters)} {kind} rollout filter(s):") - for config, filt in zip(configs, filters): - mode = "Enforcing" if filt.enforce else "Monitoring" - params = ", ".join(f"{k}={v}" for k, v in config.model_dump().items()) - get_logger().info(f" {mode} {filt.name} filter ({params})") - return filters - - -def apply_filters(filters: list[RolloutFilter], rollouts: list[TrainRollout]) -> None: - """Flag ``TrainRollout``\\ s in place with per-filter detection + drop decision. - - Each rollout's ``filter_results`` dict records per-filter detection bools; - ``is_filtered`` is True iff an enforcing filter detected it. First matching - filter wins per rollout (no double-counting). Reward and trajectory tokens - are left untouched so the rollout can still contribute to baseline - calculations and metric aggregation. - """ - for rollout in rollouts: - rollout.filter_results = {f.name: False for f in filters} - rollout.is_filtered = False - - if not filters: - return - - for rollout in rollouts: - for filt in filters: - result = filt.check(rollout) - if result.detected: - rollout.filter_results[filt.name] = True - if filt.enforce: - rollout.is_filtered = True - break diff --git a/src/prime_rl/orchestrator/metrics.py b/src/prime_rl/orchestrator/metrics.py index 85bbc60d40..6d456c629c 100644 --- a/src/prime_rl/orchestrator/metrics.py +++ b/src/prime_rl/orchestrator/metrics.py @@ -343,16 +343,16 @@ def reward(self) -> Stat: def to_wandb(self, *, prefix: str, subset: Subset) -> dict[str, float]: out = super().to_wandb(prefix=prefix, subset=subset) - # The pipeline verdicts are per-trace (an untrainable seat is 0.0 throughout, and filters - # only ever run on trainable survivors), so they read per agent like the rest. + # Detections are measured on every trace; the drop verdict is only ever reached for + # trainable survivors (an untrainable seat is 0.0 throughout). Both read per agent. for agent, traces in self.by_agent().items(): p = f"{prefix}/{subset}/{agent}" rollouts = traces.rollouts out[f"{p}/is_trainable/mean"] = sum(float(r.is_trainable) for r in rollouts) / len(rollouts) out[f"{p}/is_filtered/mean"] = sum(float(r.is_filtered) for r in rollouts) / len(rollouts) - names = sorted({name for r in rollouts for name in r.filter_results}) + names = sorted({name for r in rollouts for name in r.detections}) out |= { - f"{p}/filters/{name}/mean": sum(1 for r in rollouts if r.filter_results.get(name)) / len(rollouts) + f"{p}/detected/{name}/mean": sum(1 for r in rollouts if r.detections.get(name)) / len(rollouts) for name in names } return out diff --git a/src/prime_rl/orchestrator/orchestrator.py b/src/prime_rl/orchestrator/orchestrator.py index af6400ab58..cc225ee5ce 100644 --- a/src/prime_rl/orchestrator/orchestrator.py +++ b/src/prime_rl/orchestrator/orchestrator.py @@ -5,7 +5,7 @@ - ``RolloutDispatcher`` schedules rollouts; emits ``Episode`` (train/eval discriminated by ``kind``) on its queue — the env's own episode, or a traceless one saying why there is none. -- ``TrainSink`` ingests train rollouts (tokenize → advantages → filters) +- ``TrainSink`` ingests train rollouts (tokenize → detect → advantages → drop policy) and returns a ``TrainBatch`` when the threshold is met. - ``EvalSink`` ingests eval rollouts and returns an ``EvalBatch`` (the full returned cohort) on epoch completion. @@ -42,11 +42,11 @@ import prime_rl._compat # noqa: F401 — patch ring_flash_attn compat before transitive imports from prime_rl.configs.orchestrator import OrchestratorConfig from prime_rl.orchestrator.ckpt import setup_ckpt_manager +from prime_rl.orchestrator.detectors import setup_detectors from prime_rl.orchestrator.dispatcher import DispatcherMetrics, DispatcherMode, RolloutDispatcher from prime_rl.orchestrator.envs import EvalEnvs, TrainEnvs from prime_rl.orchestrator.eval_sink import EvalSink from prime_rl.orchestrator.eval_source import EvalSource -from prime_rl.orchestrator.filters import setup_filters from prime_rl.orchestrator.inference_metrics import InferenceMetricsCollector from prime_rl.orchestrator.patches import ( monkey_patch_chat_completion_logprobs, @@ -95,9 +95,8 @@ # shutdown wedges (env-server ZMQ recv, vLLM admin aclose, etc) SHUTDOWN_TIMEOUT_S = 300 -# Abort after this many consecutive train batches drop all rollouts to -# post-batch filters — usually a misconfigured filter or homogeneous-reward -# dataset; fail loudly instead of spinning +# Abort after this many consecutive train batches drop every rollout — usually an over-eager +# drop policy or a homogeneous-reward dataset; fail loudly instead of spinning MAX_CONSECUTIVE_EMPTY_BATCHES = 10 # Maximum batches the orchestrator may run ahead of the trainer. The @@ -254,8 +253,7 @@ async def setup(self) -> None: self.usage_reporter = UsageReporter() # Filters apply to train rollouts only - pre_filters = setup_filters(config.pre_batch_filters, vocab_size=self.tokenizer.vocab_size, kind="pre-batch") - post_filters = setup_filters(config.post_batch_filters, vocab_size=self.tokenizer.vocab_size, kind="post-batch") + detectors = setup_detectors(config.detectors, vocab_size=self.tokenizer.vocab_size) get_logger().info("Loading training environments") self.train_envs = TrainEnvs( @@ -417,8 +415,9 @@ async def setup(self) -> None: mm_token_type_ids_mapping=self.mm_token_type_ids_mapping, batch_size=config.batch_size, token_batch_size=config.token_batch_size, - pre_filters=pre_filters, - post_filters=post_filters, + detectors=detectors, + drop_detections=config.detectors.drop, + drop_zero_advantage=config.drop_zero_advantage, ) self.eval_sink = EvalSink(eval_envs=self.eval_envs) if self.eval_envs is not None else None self.watcher = WeightWatcher( @@ -589,7 +588,7 @@ async def finalize_train_batch(self, batch: TrainBatch) -> None: if self.consecutive_empty_batches >= MAX_CONSECUTIVE_EMPTY_BATCHES: raise RuntimeError( f"{self.consecutive_empty_batches} consecutive empty train batches — " - "check filter config (pre_batch_filters / post_batch_filters) or task difficulty." + "check the drop policy (detectors.drop / drop_zero_advantage) or task difficulty." ) return self.consecutive_empty_batches = 0 @@ -678,11 +677,9 @@ async def finalize_train_batch(self, batch: TrainBatch) -> None: for env_name, env_pool in batch.rollouts.by_env().items(): metrics[f"batch/{env_name}"] = len(env_pool) / len(batch.rollouts) if self.train_sink.pre_filter_seen > 0: - metrics["pre_filters/all/dropped_rate"] = ( - self.train_sink.pre_filter_dropped / self.train_sink.pre_filter_seen - ) + metrics["dropped/all/rate"] = self.train_sink.pre_filter_dropped / self.train_sink.pre_filter_seen for name, count in self.train_sink.pre_filter_dropped_by_name.items(): - metrics[f"pre_filters/all/{name}/rate"] = count / self.train_sink.pre_filter_seen + metrics[f"dropped/all/{name}/rate"] = count / self.train_sink.pre_filter_seen self.monitor.log(metrics, step=step) self.wait_for_policy_time = 0.0 self.monitor.log_samples(effective.rollouts, step=step) diff --git a/src/prime_rl/orchestrator/train_sink.py b/src/prime_rl/orchestrator/train_sink.py index 8aee8820f2..c81ec982ed 100644 --- a/src/prime_rl/orchestrator/train_sink.py +++ b/src/prime_rl/orchestrator/train_sink.py @@ -1,14 +1,15 @@ """TrainSink: three-level rollout sink for the training side. 1. ``process_rollout`` — eager per-rollout tokenization (overlaps with - dispatcher producing more rollouts), then the env algorithm's - ``finalize_rollout`` (rollout-local scoring + any reference I/O). Errored - and untrainable rollouts skip this. + dispatcher producing more rollouts), the degeneracy detectors, then the env + algorithm's ``finalize_rollout`` (rollout-local scoring + any reference + I/O). Errored and untrainable rollouts skip this. 2. ``process_group`` — filters errored rollouts, hands the episodes narrowed to their trainable survivors to the env algorithm's ``finalize_group`` - (advantages + per-sample wire stamping), runs the pre-batch filter pass. -3. ``process_batch`` — applies post-batch filter annotations and assembles - the trainer-bound ``TrainingSample`` list. Returns a ``TrainBatch``. + (advantages + per-sample wire stamping), then applies the drop policy — + the one place a rollout is kept out of training. +3. ``process_batch`` — pops a cohort and flattens it into the trainer-bound + ``TrainingSample`` list. Returns a ``TrainBatch``. ``add()`` takes one ``Episode`` and returns ``TrainBatch | None``; group accounting counts episodes, never loose traces. @@ -23,8 +24,8 @@ from collections import defaultdict from prime_rl.configs.orchestrator import OrchestratorConfig +from prime_rl.orchestrator.detectors import Detector, detect, drop_reasons from prime_rl.orchestrator.envs import TrainEnvs -from prime_rl.orchestrator.filters import RolloutFilter, apply_filters from prime_rl.orchestrator.metrics import TrainRollouts from prime_rl.orchestrator.trajectories import trace_to_samples from prime_rl.orchestrator.types import Episode, TrainBatch, TrainRollout, group_rollouts @@ -57,8 +58,9 @@ def __init__( mm_token_type_ids_mapping: dict[int, int] | None, batch_size: int | None, token_batch_size: int | None, - pre_filters: list[RolloutFilter], - post_filters: list[RolloutFilter], + detectors: list[Detector], + drop_detections: list[str], + drop_zero_advantage: bool, ) -> None: assert (batch_size is None) != (token_batch_size is None), ( "Exactly one of batch_size / token_batch_size must be set" @@ -69,8 +71,9 @@ def __init__( self.mm_token_type_ids_mapping = mm_token_type_ids_mapping self.batch_size = batch_size self.token_batch_size = token_batch_size - self.pre_filters = pre_filters - self.post_filters = post_filters + self.detectors = detectors + self.drop_detections = drop_detections + self.drop_zero_advantage = drop_zero_advantage # Observation window for the next shipped batch: rollouts of groups # finalized since the last ship (errored + filtered + survivors). @@ -165,6 +168,7 @@ async def process_rollout(self, rollout: TrainRollout) -> None: mm_token_type_ids_mapping=self.mm_token_type_ids_mapping, ) rollout.samples = samples or [] + detect(self.detectors, rollout) # Arrival phase: rollout-local scoring (raw reward, echo observation # weighting, opd/opsd reference logprobs) runs as soon as the rollout is # tokenized — before its group is complete. @@ -173,7 +177,7 @@ async def process_rollout(self, rollout: TrainRollout) -> None: async def process_group(self, group_id: uuid.UUID) -> None: """Finalize one GRPO group: drop errored rollouts (the whole group when ``requires_group_scoring`` and any failed), assign advantages, - run pre-batch filters, append survivors to ``pending_batch``.""" + apply the drop policy, append what survives to ``pending_batch``.""" episodes = self.pending_groups.pop(group_id, []) self.pending_group_episodes.pop(group_id, None) if not episodes: @@ -223,44 +227,40 @@ async def process_group(self, group_id: uuid.UUID) -> None: for sample in r.samples: sample.temperatures = [temperature] * len(sample.token_ids) - if self.pre_filters: - apply_filters(self.pre_filters, survivors) - filtered_by_name: dict[str, int] = {} - num_filtered = 0 + # Credit is assigned, so the drop decision can be made now: every detection was already + # measured at tokenization, and zero credit is only knowable after the group scored. + dropped_by_reason: dict[str, int] = {} + num_dropped = 0 for r in survivors: self.pre_filter_seen += 1 - if r.is_filtered: + reasons = drop_reasons( + r, drop_detections=self.drop_detections, drop_zero_advantage=self.drop_zero_advantage + ) + r.is_filtered = bool(reasons) + if reasons: self.pre_filter_dropped += 1 - num_filtered += 1 - for name, hit in r.filter_results.items(): - if hit: - self.pre_filter_dropped_by_name[name] = self.pre_filter_dropped_by_name.get(name, 0) + 1 - filtered_by_name[name] = filtered_by_name.get(name, 0) + 1 + num_dropped += 1 + for reason in reasons: + self.pre_filter_dropped_by_name[reason] = self.pre_filter_dropped_by_name.get(reason, 0) + 1 + dropped_by_reason[reason] = dropped_by_reason.get(reason, 0) + 1 continue - # Reset annotations so the post-batch filter pass starts clean - r.filter_results = {} - r.is_filtered = False self.pending_batch.append(r) if self.token_batch_size is not None: self.pending_tokens += payload_tokens(r) - # Per-group summary. One line per finalized group; per-filter - # detection breakdown lives at debug level in ``apply_filters`` rewards = [r.reward for r in survivors] avg_reward = sum(rewards) / len(rewards) if rewards else 0.0 - filter_str = ", ".join(f"{n}={c}" for n, c in filtered_by_name.items()) if filtered_by_name else "—" + drop_str = ", ".join(f"{n}={c}" for n, c in dropped_by_reason.items()) if dropped_by_reason else "—" get_logger().debug( f"Finished group | env={env_name} task_idx={task_idx} | " - f"rollouts={len(group)} (errored={num_errored}, filtered={num_filtered}) | " - f"reward={avg_reward:.4f} | filters: {filter_str}" + f"rollouts={len(group)} (errored={num_errored}, dropped={num_dropped}) | " + f"reward={avg_reward:.4f} | dropped: {drop_str}" ) def process_batch(self) -> TrainBatch: - """Pop a cohort off ``pending_batch`` (by rollout count when - ``batch_size`` is set, by token count when ``token_batch_size`` is - set), apply post-batch filter annotations, and assemble the - trainer-bound ``TrainingSample`` list. Overflow stays for the next - batch.""" + """Pop a cohort off ``pending_batch`` (by rollout count when ``batch_size`` is set, by + token count when ``token_batch_size`` is set) and flatten it into the trainer-bound + ``TrainingSample`` list. Overflow stays for the next batch.""" if self.batch_size is not None: cohort = self.pending_batch[: self.batch_size] self.pending_batch = self.pending_batch[self.batch_size :] @@ -277,12 +277,10 @@ def process_batch(self) -> TrainBatch: self.pending_batch = self.pending_batch[cut:] self.pending_tokens -= running - if self.post_filters: - apply_filters(self.post_filters, cohort) - # Samples are pre-built by ``process_rollout``; ``process_group`` already stamped the - # advantage stream and loss routing on each sample. Filtered rollouts don't ship. - samples: list[TrainingSample] = [sample for r in cohort if not r.is_filtered for sample in r.samples] + # advantage stream and loss routing on each sample, and decided what ships. Past this line + # the batch is samples — the episode it came from has served its purpose. + samples: list[TrainingSample] = [sample for r in cohort for sample in r.samples] # ``rollouts`` is the observation window — every rollout of every group finalized since the # last ship (errored + filtered + survivors) — while ``samples`` is the shipped cohort's diff --git a/src/prime_rl/orchestrator/types.py b/src/prime_rl/orchestrator/types.py index bdbae307eb..978f7c22a7 100644 --- a/src/prime_rl/orchestrator/types.py +++ b/src/prime_rl/orchestrator/types.py @@ -84,8 +84,11 @@ class TrainRollout(Rollout[DataT], Generic[DataT]): model_config = ConfigDict(arbitrary_types_allowed=True) # ``samples`` holds msgspec structs samples: list[TrainingSample] = Field(default_factory=list, exclude=True) + detections: dict[str, bool] = Field(default_factory=dict, exclude=True) + """What each degeneracy detector measured on this trace — a measurement, not a verdict.""" is_filtered: bool = Field(default=False, exclude=True) - filter_results: dict[str, bool] = Field(default_factory=dict, exclude=True) + """The sink's verdict: this rollout is not trained on. Kept for the metrics window, which + reports what came back as well as what shipped.""" def assign_advantages(self, value: float) -> None: """Write ``value`` as the credit for every trainable token, node by node. Credit lives on diff --git a/tests/unit/orchestrator/test_detectors.py b/tests/unit/orchestrator/test_detectors.py new file mode 100644 index 0000000000..84e3730b51 --- /dev/null +++ b/tests/unit/orchestrator/test_detectors.py @@ -0,0 +1,269 @@ +import math + +import verifiers.v1 as vf + +from prime_rl.configs.orchestrator import DetectorsConfig, GibberishDetectorConfig, RepetitionDetectorConfig +from prime_rl.orchestrator.detectors import ( + GibberishDetector, + RepetitionDetector, + detect, + drop_reasons, + setup_detectors, +) +from prime_rl.orchestrator.types import TrainRollout + + +def _assistant_node(token_ids: list[int], logprobs: list[float]) -> vf.MessageNode: + """An assistant node whose tokens are all model-sampled (the detectors read each node's + masked-True tokens + logprobs).""" + return vf.MessageNode( + message=vf.AssistantMessage(content="x"), + token_ids=token_ids, + mask=[True] * len(token_ids), + logprobs=logprobs, + ) + + +def _scaffold_assistant_node( + completion_ids: list[int], completion_logprobs: list[float], *, scaffold: int = 2 +) -> vf.MessageNode: + """A realistic v1 assistant node: a leading generation-prompt scaffold (mask=False, not + model-sampled) then the sampled completion. ``logprobs`` cover only the completion suffix + (vLLM returns logprobs for generated tokens only) — the exact layout where per-node + ``zip(token_ids, logprobs, mask)`` mispairs and the branch streams normalize.""" + return vf.MessageNode( + message=vf.AssistantMessage(content="x"), + token_ids=[1] * scaffold + completion_ids, + mask=[False] * scaffold + [True] * len(completion_ids), + logprobs=completion_logprobs, + ) + + +def _make_rollout( + completion_ids: list[int], + completion_logprobs: list[float], + *, + reward: float = 1.0, + multi_step: bool = False, +) -> TrainRollout: + """Build a ``TrainRollout`` (a message-graph trace) carrying the completion tokens — enough for + the detectors to inspect each node's sampled tokens / logprobs.""" + if multi_step: + mid = len(completion_ids) // 2 + nodes = [ + _assistant_node(completion_ids[:mid], completion_logprobs[:mid]), + _assistant_node(completion_ids[mid:], completion_logprobs[mid:]), + ] + else: + nodes = [_assistant_node(completion_ids, completion_logprobs)] + rollout = TrainRollout[vf.TaskData]( + task=vf.TraceTask(type="Task", data=vf.TaskData(idx=0, prompt="")), + agent=vf.AgentInfo(config=vf.AgentConfig()), + nodes=nodes, + rewards={"reward": vf.Reward(score=reward)}, + ) + rollout.env_name = "test" + return rollout + + +def _make_gibberish_detector(vocab_size=128_000, token_id_threshold=100_000, logprob_offset=2.0): + return GibberishDetector( + name="gibberish", + token_id_threshold=token_id_threshold, + logprob_threshold=-math.log(vocab_size) - logprob_offset, + ) + + +def _make_repetition_detector(window=5, prob_threshold=0.99): + return RepetitionDetector(name="repetition", window=window, logprob_threshold=math.log(prob_threshold)) + + +# --- GibberishDetector --- + + +def test_gibberish_detects_rare_low_prob_token(): + gibberish = _make_gibberish_detector() + + detected = gibberish.detect( + _make_rollout( + completion_ids=[50, 120_000, 80], + completion_logprobs=[-1.0, gibberish.logprob_threshold - 1.0, -0.5], + ) + ) + assert detected is True + + +def test_gibberish_ignores_normal_tokens(): + gibberish = _make_gibberish_detector() + + detected = gibberish.detect( + _make_rollout( + completion_ids=[10, 200, 5000], + completion_logprobs=[-1.0, -2.0, -3.0], + ) + ) + assert detected is False + + +def test_gibberish_ignores_high_prob_rare_token(): + gibberish = _make_gibberish_detector() + + detected = gibberish.detect( + _make_rollout( + completion_ids=[120_000], + completion_logprobs=[-0.5], + ) + ) + assert detected is False + + +def test_gibberish_works_across_trajectory_steps(): + gibberish = _make_gibberish_detector() + + detected = gibberish.detect( + _make_rollout( + completion_ids=[50, 60, 120_000, 80], + completion_logprobs=[-1.0, -0.5, gibberish.logprob_threshold - 1.0, -0.5], + multi_step=True, + ) + ) + assert detected is True + + +def test_gibberish_aligns_logprobs_under_generation_prompt_scaffold(): + """Regression: the assistant node carries a generation-prompt scaffold (mask=False) and + suffix-only logprobs, and the gibberish token is the LAST completion token. The old + per-node ``zip(token_ids, logprobs, mask)`` truncated at len(logprobs) and never examined + it; reading the aligned branch streams detects it.""" + gibberish = _make_gibberish_detector() + + rollout = TrainRollout[vf.TaskData]( + task=vf.TraceTask(type="Task", data=vf.TaskData(idx=0, prompt="")), + agent=vf.AgentInfo(config=vf.AgentConfig()), + nodes=[_scaffold_assistant_node([50, 80, 120_000], [-1.0, -0.5, gibberish.logprob_threshold - 1.0])], + rewards={"reward": vf.Reward(score=1.0)}, + ) + + detected = gibberish.detect(rollout) + assert detected is True + + +# --- RepetitionDetector --- + + +def test_repetition_triggers_after_window(): + repetition = _make_repetition_detector(window=5) + + detected = repetition.detect( + _make_rollout( + completion_ids=list(range(5)), + completion_logprobs=[-0.001] * 5, + ) + ) + assert detected is True + + +def test_repetition_no_trigger_below_window(): + repetition = _make_repetition_detector(window=5) + + detected = repetition.detect( + _make_rollout( + completion_ids=list(range(4)), + completion_logprobs=[-0.001] * 4, + ) + ) + assert detected is False + + +def test_repetition_resets_on_low_prob(): + repetition = _make_repetition_detector(window=5) + + logprobs = [-0.001] * 3 + [-2.0] + [-0.001] * 3 + detected = repetition.detect( + _make_rollout( + completion_ids=list(range(7)), + completion_logprobs=logprobs, + ) + ) + assert detected is False + + +def test_repetition_varied_probs_no_trigger(): + repetition = _make_repetition_detector(window=3) + + detected = repetition.detect( + _make_rollout( + completion_ids=list(range(6)), + completion_logprobs=[-0.001, -3.0, -0.001, -3.0, -0.001, -3.0], + ) + ) + assert detected is False + + +# --- setup + drop policy --- + + +def test_setup_detectors_builds_both_and_derives_thresholds(): + detectors = setup_detectors(DetectorsConfig(), vocab_size=128_000) + assert [d.name for d in detectors] == ["gibberish", "repetition"] + gibberish, repetition = detectors + assert gibberish.token_id_threshold == 100_000 + assert gibberish.logprob_threshold == -math.log(128_000) - 2.0 + assert repetition.logprob_threshold == math.log(0.99) + + +def test_setup_detectors_can_turn_one_off(): + config = DetectorsConfig(gibberish=None, repetition=RepetitionDetectorConfig(window=7)) + assert [d.name for d in setup_detectors(config, vocab_size=128_000)] == ["repetition"] + + +def test_setup_detectors_thresholds_follow_config(): + config = DetectorsConfig(gibberish=GibberishDetectorConfig(token_id_threshold=50_000, logprob_offset=1.0)) + gibberish = setup_detectors(config, vocab_size=1_000)[0] + assert (gibberish.token_id_threshold, gibberish.logprob_threshold) == (50_000, -math.log(1_000) - 1.0) + + +def test_detect_measures_every_detector(): + """Every detector measures every trace — one hit does not shadow the others, which is what + makes a detection rate a metric rather than a by-product of the drop order.""" + gibberish = _make_gibberish_detector() + rollout = _make_rollout( + completion_ids=[120_000] * 5, + completion_logprobs=[gibberish.logprob_threshold - 1.0] * 5, + ) + detect([gibberish, _make_repetition_detector(window=5)], rollout) + assert rollout.detections == {"gibberish": True, "repetition": False} + + +def _drop(rollout, drop=(), zero_advantage=True): + return drop_reasons(rollout, drop_detections=list(drop), drop_zero_advantage=zero_advantage) + + +def test_detection_only_drops_when_asked(): + rollout = _make_rollout(completion_ids=[1], completion_logprobs=[-1.0]) + rollout.detections = {"gibberish": True} + rollout.assign_advantages(0.5) + assert _drop(rollout) == [] # measured, but nothing asked for it to drop + assert _drop(rollout, drop=["gibberish"]) == ["gibberish"] + + +def test_zero_credit_drops_by_default(): + rollout = _make_rollout(completion_ids=[1], completion_logprobs=[-1.0]) + rollout.assign_advantages(0.0) + assert _drop(rollout) == ["zero_advantage"] + assert _drop(rollout, zero_advantage=False) == [] + + +def test_unscored_rollout_is_not_zero_credit(): + """opd/opsd assign no credit at all and train through reference KL — dropping them as + zero-advantage would ship an empty batch for every distillation run.""" + rollout = _make_rollout(completion_ids=[1], completion_logprobs=[-1.0]) + assert rollout.advantages is None + assert _drop(rollout) == [] + + +def test_reasons_accumulate(): + rollout = _make_rollout(completion_ids=[1], completion_logprobs=[-1.0]) + rollout.detections = {"gibberish": True, "repetition": True} + rollout.assign_advantages(0.0) + assert _drop(rollout, drop=["gibberish", "repetition"]) == ["gibberish", "repetition", "zero_advantage"] diff --git a/tests/unit/orchestrator/test_filters.py b/tests/unit/orchestrator/test_filters.py deleted file mode 100644 index 4d35c842b6..0000000000 --- a/tests/unit/orchestrator/test_filters.py +++ /dev/null @@ -1,406 +0,0 @@ -import math - -import verifiers.v1 as vf - -from prime_rl.configs.orchestrator import GibberishFilterConfig, RepetitionFilterConfig -from prime_rl.orchestrator.filters import ( - GibberishFilter, - RepetitionFilter, - apply_filters, - setup_filter, - setup_filters, -) -from prime_rl.orchestrator.types import TrainRollout - - -def _assistant_node(token_ids: list[int], logprobs: list[float]) -> vf.MessageNode: - """An assistant node whose tokens are all model-sampled (the filters read each node's - masked-True tokens + logprobs).""" - return vf.MessageNode( - message=vf.AssistantMessage(content="x"), - token_ids=token_ids, - mask=[True] * len(token_ids), - logprobs=logprobs, - ) - - -def _scaffold_assistant_node( - completion_ids: list[int], completion_logprobs: list[float], *, scaffold: int = 2 -) -> vf.MessageNode: - """A realistic v1 assistant node: a leading generation-prompt scaffold (mask=False, not - model-sampled) then the sampled completion. ``logprobs`` cover only the completion suffix - (vLLM returns logprobs for generated tokens only) — the exact layout where per-node - ``zip(token_ids, logprobs, mask)`` mispairs and the branch streams normalize.""" - return vf.MessageNode( - message=vf.AssistantMessage(content="x"), - token_ids=[1] * scaffold + completion_ids, - mask=[False] * scaffold + [True] * len(completion_ids), - logprobs=completion_logprobs, - ) - - -def _make_rollout( - completion_ids: list[int], - completion_logprobs: list[float], - *, - reward: float = 1.0, - multi_step: bool = False, -) -> TrainRollout: - """Build a ``TrainRollout`` (a message-graph trace) carrying the completion tokens — enough for - the filters to inspect each node's sampled tokens / logprobs.""" - if multi_step: - mid = len(completion_ids) // 2 - nodes = [ - _assistant_node(completion_ids[:mid], completion_logprobs[:mid]), - _assistant_node(completion_ids[mid:], completion_logprobs[mid:]), - ] - else: - nodes = [_assistant_node(completion_ids, completion_logprobs)] - rollout = TrainRollout[vf.TaskData]( - task=vf.TraceTask(type="Task", data=vf.TaskData(idx=0, prompt="")), - agent=vf.AgentInfo(config=vf.AgentConfig()), - nodes=nodes, - rewards={"reward": vf.Reward(score=reward)}, - ) - rollout.env_name = "test" - return rollout - - -def _make_gibberish_filter(vocab_size=128_000, token_id_threshold=100_000, logprob_offset=2.0, enforce=False): - logprob_threshold = -math.log(vocab_size) - logprob_offset - return GibberishFilter( - name="gibberish", token_id_threshold=token_id_threshold, logprob_threshold=logprob_threshold, enforce=enforce - ) - - -def _make_repetition_filter(window=5, prob_threshold=0.99, enforce=False): - return RepetitionFilter( - name="repetition", window=window, logprob_threshold=math.log(prob_threshold), enforce=enforce - ) - - -# --- GibberishFilter tests --- - - -def test_gibberish_detects_rare_low_prob_token(): - gibberish_filter = _make_gibberish_filter() - - result = gibberish_filter.check( - _make_rollout( - completion_ids=[50, 120_000, 80], - completion_logprobs=[-1.0, gibberish_filter.logprob_threshold - 1.0, -0.5], - ) - ) - assert result.detected is True - - -def test_gibberish_ignores_normal_tokens(): - gibberish_filter = _make_gibberish_filter() - - result = gibberish_filter.check( - _make_rollout( - completion_ids=[10, 200, 5000], - completion_logprobs=[-1.0, -2.0, -3.0], - ) - ) - assert result.detected is False - - -def test_gibberish_ignores_high_prob_rare_token(): - gibberish_filter = _make_gibberish_filter() - - result = gibberish_filter.check( - _make_rollout( - completion_ids=[120_000], - completion_logprobs=[-0.5], - ) - ) - assert result.detected is False - - -def test_gibberish_works_across_trajectory_steps(): - gibberish_filter = _make_gibberish_filter() - - result = gibberish_filter.check( - _make_rollout( - completion_ids=[50, 60, 120_000, 80], - completion_logprobs=[-1.0, -0.5, gibberish_filter.logprob_threshold - 1.0, -0.5], - multi_step=True, - ) - ) - assert result.detected is True - - -def test_gibberish_aligns_logprobs_under_generation_prompt_scaffold(): - """Regression: the assistant node carries a generation-prompt scaffold (mask=False) and - suffix-only logprobs, and the gibberish token is the LAST completion token. The old - per-node ``zip(token_ids, logprobs, mask)`` truncated at len(logprobs) and never examined - it; reading the aligned branch streams detects it.""" - gibberish_filter = _make_gibberish_filter() - - rollout = TrainRollout[vf.TaskData]( - task=vf.TraceTask(type="Task", data=vf.TaskData(idx=0, prompt="")), - agent=vf.AgentInfo(config=vf.AgentConfig()), - nodes=[_scaffold_assistant_node([50, 80, 120_000], [-1.0, -0.5, gibberish_filter.logprob_threshold - 1.0])], - rewards={"reward": vf.Reward(score=1.0)}, - ) - - result = gibberish_filter.check(rollout) - assert result.detected is True - - -# --- RepetitionFilter tests --- - - -def test_repetition_triggers_after_window(): - repetition_filter = _make_repetition_filter(window=5) - - result = repetition_filter.check( - _make_rollout( - completion_ids=list(range(5)), - completion_logprobs=[-0.001] * 5, - ) - ) - assert result.detected is True - - -def test_repetition_no_trigger_below_window(): - repetition_filter = _make_repetition_filter(window=5) - - result = repetition_filter.check( - _make_rollout( - completion_ids=list(range(4)), - completion_logprobs=[-0.001] * 4, - ) - ) - assert result.detected is False - - -def test_repetition_resets_on_low_prob(): - repetition_filter = _make_repetition_filter(window=5) - - logprobs = [-0.001] * 3 + [-2.0] + [-0.001] * 3 - result = repetition_filter.check( - _make_rollout( - completion_ids=list(range(7)), - completion_logprobs=logprobs, - ) - ) - assert result.detected is False - - -def test_repetition_varied_probs_no_trigger(): - repetition_filter = _make_repetition_filter(window=3) - - result = repetition_filter.check( - _make_rollout( - completion_ids=list(range(6)), - completion_logprobs=[-0.001, -3.0, -0.001, -3.0, -0.001, -3.0], - ) - ) - assert result.detected is False - - -# --- setup_filter / setup_filters tests --- - - -def test_setup_filter_gibberish(): - config = GibberishFilterConfig(token_id_threshold=100_000, logprob_offset=2.0) - gibberish_filter = setup_filter(config, vocab_size=128_000) - assert isinstance(gibberish_filter, GibberishFilter) - assert gibberish_filter.name == "gibberish" - assert gibberish_filter.token_id_threshold == 100_000 - assert abs(gibberish_filter.logprob_threshold - (-math.log(128_000) - 2.0)) < 1e-10 - assert gibberish_filter.enforce is False - - -def test_setup_filter_gibberish_enforce(): - config = GibberishFilterConfig(enforce=True) - gibberish_filter = setup_filter(config, vocab_size=128_000) - assert gibberish_filter.enforce is True - - -def test_setup_filter_repetition(): - config = RepetitionFilterConfig(window=3_000, prob_threshold=0.99) - repetition_filter = setup_filter(config, vocab_size=128_000) - assert isinstance(repetition_filter, RepetitionFilter) - assert repetition_filter.name == "repetition" - assert repetition_filter.window == 3_000 - assert abs(repetition_filter.logprob_threshold - math.log(0.99)) < 1e-10 - assert repetition_filter.enforce is False - - -def test_setup_filter_repetition_enforce(): - config = RepetitionFilterConfig(enforce=True) - repetition_filter = setup_filter(config, vocab_size=128_000) - assert repetition_filter.enforce is True - - -def test_setup_filters_multiple(): - configs = [ - GibberishFilterConfig(), - RepetitionFilterConfig(), - ] - filters = setup_filters(configs, vocab_size=128_000, kind="post-batch") - assert len(filters) == 2 - assert filters[0].name == "gibberish" - assert filters[1].name == "repetition" - - -# --- apply_filters tests (enforce=True) --- - - -def test_apply_filters_enforced_flags_rollout(): - gibberish_filter = _make_gibberish_filter(enforce=True) - - rollout = _make_rollout( - completion_ids=[120_000], - completion_logprobs=[gibberish_filter.logprob_threshold - 1.0], - reward=1.0, - ) - - apply_filters([gibberish_filter], [rollout]) - - assert rollout.reward == 1.0 - assert rollout.nodes[0].token_ids == [120_000] - assert rollout.nodes[0].mask == [True] - assert rollout.stop_condition is None - assert rollout.filter_results == {"gibberish": True} - assert rollout.is_filtered is True - - -def test_apply_filters_preserves_clean_rollouts(): - gibberish_filter = _make_gibberish_filter(enforce=True) - - rollout = _make_rollout( - completion_ids=[50, 60, 70], - completion_logprobs=[-1.0, -2.0, -1.5], - reward=1.0, - ) - - apply_filters([gibberish_filter], [rollout]) - - assert rollout.reward == 1.0 - assert rollout.nodes[0].token_ids == [50, 60, 70] - assert all(rollout.nodes[0].mask) - assert rollout.stop_condition is None - assert rollout.filter_results == {"gibberish": False} - assert rollout.is_filtered is False - - -def test_apply_filters_first_filter_wins(): - gibberish_filter = _make_gibberish_filter(enforce=True) - repetition_filter = _make_repetition_filter(window=2, enforce=True) - - rollout = _make_rollout( - completion_ids=[120_000, 1, 2], - completion_logprobs=[gibberish_filter.logprob_threshold - 1.0, -0.001, -0.001], - reward=1.0, - ) - - apply_filters([gibberish_filter, repetition_filter], [rollout]) - - assert rollout.stop_condition is None - assert rollout.filter_results == {"gibberish": True, "repetition": False} - assert rollout.is_filtered is True - - -def test_apply_filters_empty_list(): - rollout = _make_rollout( - completion_ids=[1, 2, 3], - completion_logprobs=[-1.0, -1.0, -1.0], - ) - apply_filters([], [rollout]) - assert rollout.filter_results == {} - assert rollout.is_filtered is False - assert rollout.reward == 1.0 - - -def test_apply_filters_mixed_batch(): - gibberish_filter = _make_gibberish_filter(enforce=True) - - clean = _make_rollout(completion_ids=[50], completion_logprobs=[-1.0], reward=1.0) - dirty = _make_rollout( - completion_ids=[120_000], completion_logprobs=[gibberish_filter.logprob_threshold - 1.0], reward=1.0 - ) - - apply_filters([gibberish_filter], [clean, dirty]) - - assert clean.reward == 1.0 - assert dirty.reward == 1.0 - assert clean.is_filtered is False - assert dirty.is_filtered is True - - -def test_apply_filters_enforced_preserves_rollout_tokens(): - gibberish_filter = _make_gibberish_filter(enforce=True) - - rollout = _make_rollout( - completion_ids=[10, 120_000, 30], - completion_logprobs=[-1.0, gibberish_filter.logprob_threshold - 1.0, -0.5], - reward=1.0, - ) - - apply_filters([gibberish_filter], [rollout]) - - assert rollout.nodes[0].token_ids == [10, 120_000, 30] - assert rollout.nodes[0].logprobs == [ - -1.0, - gibberish_filter.logprob_threshold - 1.0, - -0.5, - ] - assert rollout.nodes[0].mask == [True, True, True] - assert rollout.is_filtered is True - - -def test_apply_filters_preserves_existing_stop_condition(): - gibberish_filter = _make_gibberish_filter(enforce=True) - - rollout = _make_rollout( - completion_ids=[120_000], - completion_logprobs=[gibberish_filter.logprob_threshold - 1.0], - reward=1.0, - ) - rollout.stop_condition = "generation_truncated" - - apply_filters([gibberish_filter], [rollout]) - - assert rollout.stop_condition == "generation_truncated" - assert rollout.is_filtered is True - - -# --- apply_filters tests (monitor-only, enforce=False) --- - - -def test_apply_filters_monitor_only_tracks_detection(): - gibberish_filter = _make_gibberish_filter(enforce=False) - - rollout = _make_rollout( - completion_ids=[120_000], - completion_logprobs=[gibberish_filter.logprob_threshold - 1.0], - reward=1.0, - ) - - apply_filters([gibberish_filter], [rollout]) - - assert rollout.reward == 1.0 - assert all(rollout.nodes[0].mask) - assert rollout.stop_condition is None - assert rollout.filter_results == {"gibberish": True} - assert rollout.is_filtered is False - - -def test_apply_filters_monitor_only_mixed_batch(): - gibberish_filter = _make_gibberish_filter(enforce=False) - - clean = _make_rollout(completion_ids=[50], completion_logprobs=[-1.0], reward=1.0) - dirty = _make_rollout( - completion_ids=[120_000], completion_logprobs=[gibberish_filter.logprob_threshold - 1.0], reward=1.0 - ) - - apply_filters([gibberish_filter], [clean, dirty]) - - assert clean.reward == 1.0 - assert dirty.reward == 1.0 - assert clean.is_filtered is False - assert dirty.is_filtered is False diff --git a/tests/unit/orchestrator/test_metrics.py b/tests/unit/orchestrator/test_metrics.py index ee0b2f406a..5fd4961b8a 100644 --- a/tests/unit/orchestrator/test_metrics.py +++ b/tests/unit/orchestrator/test_metrics.py @@ -35,7 +35,7 @@ def mk( trainable: bool = True, is_trainable: bool = True, is_filtered: bool = False, - filter_results: dict | None = None, + detections: dict | None = None, setup: float = 0.0, agent: float = 0.0, agent_model: float = 0.0, @@ -66,7 +66,7 @@ def mk( agent=SimpleNamespace(trainable=trainable, name=agent_name), is_trainable=is_trainable, is_filtered=is_filtered, - filter_results=filter_results or {}, + detections=detections or {}, timing=SimpleNamespace( setup=SimpleNamespace(duration=setup), agent=SimpleNamespace( @@ -270,16 +270,16 @@ def test_nested_timing(): def test_train_only_metrics_absent_from_eval(): rollouts = [ - mk(is_trainable=True, is_filtered=True, filter_results={"gibberish": True}), - mk(is_trainable=False, filter_results={"gibberish": False}), + mk(is_trainable=True, is_filtered=True, detections={"gibberish": True}), + mk(is_trainable=False, detections={"gibberish": False}), ] out = train_wandb(rollouts) assert out["train/agg/all/agent/is_trainable/mean"] == 0.5 assert out["train/agg/all/agent/is_filtered/mean"] == 0.5 - assert out["train/agg/all/agent/filters/gibberish/mean"] == 0.5 + assert out["train/agg/all/agent/detected/gibberish/mean"] == 0.5 assert "train/agg/all/is_trainable/mean" not in out # pipeline verdicts are per-trace eval_out = EvalRollouts(solo(rollouts)).metrics.to_wandb(prefix="eval/x", subset="all") - assert not any("is_trainable" in k or "is_filtered" in k or "/filters/" in k for k in eval_out) + assert not any("is_trainable" in k or "is_filtered" in k or "/detected/" in k for k in eval_out) def test_eval_avg_at_k_and_pass_k(): @@ -319,8 +319,8 @@ def test_traceless_episode_keeps_its_reason(): def test_training_state_is_train_only(): """Only a train rollout carries trainer-bound state; an eval trace has no field for it. Credit is not among them — it lives on the graph's nodes, which every trace has.""" - assert {"samples", "is_filtered", "filter_results"} <= set(TrainRollout.model_fields) - assert not {"samples", "is_filtered", "filter_results"} & set(Rollout.model_fields) + assert {"samples", "is_filtered", "detections"} <= set(TrainRollout.model_fields) + assert not {"samples", "is_filtered", "detections"} & set(Rollout.model_fields) assert "advantages" not in TrainRollout.model_fields # derived from the nodes assert {"env_name", "episode_id"} <= set(Rollout.model_fields) # links stay shared assert "group_id" not in Rollout.model_fields # the example a trace answered is the episode's From c8eac1740678b226f6432baa6a46469a60d16b7c Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Tue, 4 Aug 2026 21:28:58 +0000 Subject: [PATCH 41/58] refactor(orchestrator)!: prime-rl stops extending Episode MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Everything the orchestrator adds to a dispatch now has a place on vf's own episode: the env it ran (env.name), the group it was planned in (group), and the run it belongs to — which on the training path carries the policy version and how stale it got, and tells train from online eval by kind. Episode is an alias for vf.WireEpisode, and what were methods are free functions over it. Co-Authored-By: Claude Fable 5 --- deps/verifiers | 2 +- .../orchestrator/algo/hierarchical_grpo.py | 3 +- src/prime_rl/orchestrator/eval_sink.py | 17 ++- src/prime_rl/orchestrator/metrics.py | 23 ++-- src/prime_rl/orchestrator/orchestrator.py | 19 +-- src/prime_rl/orchestrator/train_sink.py | 28 ++-- src/prime_rl/orchestrator/types.py | 127 ++++++++---------- tests/unit/orchestrator/test_metrics.py | 44 ++++-- 8 files changed, 136 insertions(+), 127 deletions(-) diff --git a/deps/verifiers b/deps/verifiers index 48591d3735..70490aa3ef 160000 --- a/deps/verifiers +++ b/deps/verifiers @@ -1 +1 @@ -Subproject commit 48591d3735f3bb42dd93c7878adaa6ad6c5ff6be +Subproject commit 70490aa3efe7080e1a4b138da58eb63b295b1771 diff --git a/src/prime_rl/orchestrator/algo/hierarchical_grpo.py b/src/prime_rl/orchestrator/algo/hierarchical_grpo.py index 6495115c9e..f9a2b6eb54 100644 --- a/src/prime_rl/orchestrator/algo/hierarchical_grpo.py +++ b/src/prime_rl/orchestrator/algo/hierarchical_grpo.py @@ -5,6 +5,7 @@ from prime_rl.configs.algorithm import HierarchicalGRPOAlgoConfig from prime_rl.orchestrator.algo.base import Algorithm +from prime_rl.orchestrator.types import rollouts_of if TYPE_CHECKING: from prime_rl.orchestrator.types import Episode, TrainRollout @@ -31,7 +32,7 @@ def __init__(self, config: HierarchicalGRPOAlgoConfig, policy_pool: InferencePoo async def score_group(self, group: list[Episode]) -> None: peers: dict[tuple[str, str | None], list[TrainRollout]] = defaultdict(list) for episode in group: - for rollout in episode.rollouts: + for rollout in rollouts_of(episode): episode_scoped = rollout.agent.name in self.episode_agents peers[(rollout.agent.name, episode.id if episode_scoped else None)].append(rollout) for members in peers.values(): diff --git a/src/prime_rl/orchestrator/eval_sink.py b/src/prime_rl/orchestrator/eval_sink.py index 564594f485..aea0649983 100644 --- a/src/prime_rl/orchestrator/eval_sink.py +++ b/src/prime_rl/orchestrator/eval_sink.py @@ -14,14 +14,13 @@ from __future__ import annotations -import uuid from collections import defaultdict import verifiers.v1 as vf from prime_rl.orchestrator.envs import EvalEnvs from prime_rl.orchestrator.metrics import EvalRollouts -from prime_rl.orchestrator.types import Episode, EvalBatch, Rollout +from prime_rl.orchestrator.types import Episode, EvalBatch, Rollout, env_name_of, group_id_of, rollouts_of from prime_rl.utils.logger import get_logger @@ -37,9 +36,9 @@ class EvalSink: def __init__(self, *, eval_envs: EvalEnvs) -> None: self.eval_envs = eval_envs - self.pending_groups: dict[uuid.UUID, list[Episode]] = defaultdict(list) + self.pending_groups: dict[str, list[Episode]] = defaultdict(list) # Episodes arrived per group / per batch bucket — the finalization counts. - self.pending_group_episodes: dict[uuid.UUID, int] = defaultdict(int) + self.pending_group_episodes: dict[str, int] = defaultdict(int) self.pending_batches: dict[tuple[str, int], list[Episode]] = defaultdict(list) self.pending_batch_episodes: dict[tuple[str, int], int] = defaultdict(int) @@ -47,9 +46,9 @@ def add(self, episode: Episode) -> EvalBatch | None: """Process one episode arrival; finalize the group on the ``group_size``-th episode and the per-env epoch on the ``num_examples × group_size``-th. A failed episode brings no rollouts but still counts toward both.""" - env_name = episode.env_name - group_id = episode.group_id - for rollout in episode.rollouts: + env_name = env_name_of(episode) + group_id = group_id_of(episode) + for rollout in rollouts_of(episode): self.process_rollout(rollout) bkey = (env_name, eval_step_of(episode)) self.pending_groups[group_id].append(episode) @@ -106,7 +105,7 @@ def process_rollout(self, rollout: Rollout) -> None: # ── level 2: per-group (move into batch bucket) ─────────────────────── - def process_group(self, group_id: uuid.UUID) -> None: + def process_group(self, group_id: str) -> None: finished = self.pending_groups.pop(group_id, []) episodes = self.pending_group_episodes.pop(group_id, 0) if not finished: @@ -115,7 +114,7 @@ def process_group(self, group_id: uuid.UUID) -> None: # produced none (a whole group cancelled off-policy). env_name = finished[0].env_name eval_step = eval_step_of(finished[0]) - group = [t for e in finished for t in e.rollouts] + group = [t for e in finished for t in rollouts_of(e)] task_idx = group[0].task.data.idx if group else -1 bucket = self.pending_batches[(env_name, eval_step)] bucket.extend(finished) diff --git a/src/prime_rl/orchestrator/metrics.py b/src/prime_rl/orchestrator/metrics.py index 6d456c629c..bcbb61865b 100644 --- a/src/prime_rl/orchestrator/metrics.py +++ b/src/prime_rl/orchestrator/metrics.py @@ -23,6 +23,7 @@ from typing import TYPE_CHECKING, Any, Callable, Iterator, Literal +from prime_rl.orchestrator.types import env_name_of, group_id_of, narrow, rollouts_of from prime_rl.orchestrator.utils import compute_pass_metrics if TYPE_CHECKING: @@ -179,7 +180,7 @@ class TraceMetrics(StatGroup): def __init__(self, episodes: list[Episode]) -> None: self.episodes = episodes - super().__init__([t for e in episodes for t in e.rollouts]) + super().__init__([t for e in episodes for t in rollouts_of(e)]) DISTRIBUTIONS = ("reward", "num_total_tokens", "num_input_tokens", "num_output_tokens", "num_turns", "num_branches") RATES = ("is_truncated", "is_completed") @@ -235,7 +236,7 @@ def solve_rates(self) -> dict[str, float]: groups).""" groups: dict = {} for e in self.episodes: - groups.setdefault(e.group_id, []).extend(e.rollouts) + groups.setdefault(group_id_of(e), []).extend(rollouts_of(e)) n_groups = len(groups) solved_none = sum(1 for g in groups.values() if sum(r.reward for r in g) == 0) solved_all = sum(1 for g in groups.values() if all(r.reward == 1.0 for r in g)) @@ -282,7 +283,7 @@ def by_agent(self) -> dict[str, TraceMetrics]: names each episode's own ``vf.Episode.by_agent`` reports.""" names = sorted({name for e in self.episodes for name in e.by_agent}) return { - name: TraceMetrics([n for e in self.episodes if (n := e.narrow(lambda r: r.agent.name == name))]) + name: TraceMetrics([n for e in self.episodes if (n := narrow(e, lambda r: r.agent.name == name))]) for name in names } @@ -360,12 +361,12 @@ def to_wandb(self, *, prefix: str, subset: Subset) -> dict[str, float]: def pass_at_k(episodes: list[Episode]) -> dict[str, float]: """pass@k / pass^k averaged over examples; ``{}`` for non-binary rewards.""" - rewards = [r.reward for e in episodes for r in e.rollouts] + rewards = [r.reward for e in episodes for r in rollouts_of(e)] if not set(rewards).issubset({0.0, 1.0}): return {} by_example: dict = {} for e in episodes: - by_example.setdefault(e.group_id, []).extend(r.reward for r in e.rollouts) + by_example.setdefault(group_id_of(e), []).extend(r.reward for r in rollouts_of(e)) per_example = [compute_pass_metrics(rs) for rs in by_example.values()] keys = sorted({k for d in per_example for k in d}) return {k: sum(d[k] for d in per_example if k in d) / sum(1 for d in per_example if k in d) for k in keys} @@ -410,7 +411,7 @@ def append(self, episode: Episode) -> None: @property def rollouts(self) -> list[TrainRollout]: - return [t for e in self.episodes for t in e.rollouts] + return [t for e in self.episodes for t in rollouts_of(e)] def __len__(self) -> int: return sum(len(e.traces) for e in self.episodes) @@ -420,13 +421,13 @@ def __iter__(self) -> Iterator[TrainRollout]: @property def effective(self) -> TrainRollouts: - kept = (e.narrow(lambda r: not r.has_error and not r.is_filtered and r.agent.trainable) for e in self.episodes) + kept = (narrow(e, lambda r: not r.has_error and not r.is_filtered and r.agent.trainable) for e in self.episodes) return TrainRollouts([e for e in kept if e is not None]) def by_env(self) -> dict[str, TrainRollouts]: grouped: dict[str, list[Episode]] = {} for episode in self.episodes: - grouped.setdefault(episode.env_name, []).append(episode) + grouped.setdefault(env_name_of(episode), []).append(episode) return {env: TrainRollouts(episodes) for env, episodes in grouped.items()} @property @@ -464,13 +465,13 @@ def group_size(self) -> int: return self._group_size counts: dict = {} for e in self.episodes: - trainable = sum(1 for r in e.rollouts if r.agent.trainable) - counts[e.group_id] = counts.get(e.group_id, 0) + trainable + trainable = sum(1 for r in rollouts_of(e) if r.agent.trainable) + counts[group_id_of(e)] = counts.get(group_id_of(e), 0) + trainable return max(counts.values(), default=0) @property def effective(self) -> EvalRollouts: - kept = (e.narrow(lambda r: not r.has_error and r.agent.trainable) for e in self.episodes) + kept = (narrow(e, lambda r: not r.has_error and r.agent.trainable) for e in self.episodes) return EvalRollouts([e for e in kept if e is not None], group_size=self.group_size) @property diff --git a/src/prime_rl/orchestrator/orchestrator.py b/src/prime_rl/orchestrator/orchestrator.py index cc225ee5ce..18d4522227 100644 --- a/src/prime_rl/orchestrator/orchestrator.py +++ b/src/prime_rl/orchestrator/orchestrator.py @@ -61,6 +61,7 @@ Policy, Progress, TrainBatch, + run_of, ) from prime_rl.orchestrator.utils import ( get_weight_dir, @@ -528,19 +529,18 @@ async def main_loop(self) -> None: # episode per line, so an episode that produced no traces still records why. # An eval episode already knows its step (the eval that triggered it); a train one # belongs to the batch window collecting right now, which only this loop knows. - run = episode.run - assert run is not None, "the dispatcher records the run when the episode lands" - if run.type == "train": + run = run_of(episode) + if run.kind == "train": run.step = self.progress.step step = run.step assert step is not None await asyncio.to_thread( save_episodes, [episode.to_record()], - get_trace_path(self.config.output_dir, step, run.type, "all"), + get_trace_path(self.config.output_dir, step, run.kind, "all"), ) - if run.type == "eval": + if run.kind == "eval": assert self.eval_sink is not None # eval rollouts only emitted when eval is configured eval_batch = self.eval_sink.add(episode) if eval_batch is not None: @@ -627,7 +627,8 @@ async def finalize_train_batch(self, batch: TrainBatch) -> None: # sourced rollouts stay 0 (their sampler doesn't follow the policy). for train_episode in batch.rollouts.episodes: if self.train_envs.get(train_episode.env_name).sampler.samples_from_live_policy: - train_episode.off_policy_steps = (step - 1) - train_episode.policy_version + run = run_of(train_episode) + run.off_policy_steps = (step - 1) - (run.policy_version or 0) # The effective (clean, trained-on) subset lands in the per-step ``effective`` trace file # at ship time; the full arrival window already streamed into ``all`` on arrival. @@ -806,7 +807,7 @@ def log_train_batch(self, batch: TrainBatch, *, step: int, step_time: float) -> n_effective = len(effective) n_trainable = sum(1 for r in effective if r.is_trainable) trainable_rate = (n_trainable / n_effective) if n_effective else 0.0 - max_off_policy = max((e.off_policy_steps for e in effective.episodes), default=0) + max_off_policy = max((run_of(e).off_policy_steps for e in effective.episodes), default=0) head = ( f"Step {step} | {format_time(step_time):>7} | Reward {eff.reward.mean():.4f} | " @@ -830,7 +831,7 @@ def log_train_batch(self, batch: TrainBatch, *, step: int, step_time: float) -> lines.append( f"╰─ {env_name:<{name_width}} | Ratio {ratio:.1%} | Reward {env_eff.reward.mean():.4f} | " f"Turns {env_eff.num_turns.mean():.1f} | Branches {env_eff.num_branches.mean():.1f} | " - f"Max Off-Policy {max((e.off_policy_steps for e in env_eff_pool.episodes), default=0)} | " + f"Max Off-Policy {max((run_of(e).off_policy_steps for e in env_eff_pool.episodes), default=0)} | " f"Error {pool.metrics.has_error.mean():.1%} | Truncation {env_eff.is_truncated.mean():.1%}" ) get_logger().success("\n\t\t ".join(lines)) @@ -851,7 +852,7 @@ async def finalize_eval_batch(self, batch: EvalBatch) -> None: save_episodes, records, get_trace_path(self.config.output_dir, batch.step, "eval", "effective") ) self.monitor.log_eval_samples(batch.rollouts, env_name=batch.env_name, step=batch.step) - policy_versions = {e.policy_version for e in batch.rollouts.episodes} + policy_versions = {run_of(e).policy_version for e in batch.rollouts.episodes} policy_version = min(policy_versions) if len(policy_versions) > 1: get_logger().warning( diff --git a/src/prime_rl/orchestrator/train_sink.py b/src/prime_rl/orchestrator/train_sink.py index c81ec982ed..5780580a99 100644 --- a/src/prime_rl/orchestrator/train_sink.py +++ b/src/prime_rl/orchestrator/train_sink.py @@ -20,7 +20,6 @@ from __future__ import annotations import asyncio -import uuid from collections import defaultdict from prime_rl.configs.orchestrator import OrchestratorConfig @@ -28,7 +27,16 @@ from prime_rl.orchestrator.envs import TrainEnvs from prime_rl.orchestrator.metrics import TrainRollouts from prime_rl.orchestrator.trajectories import trace_to_samples -from prime_rl.orchestrator.types import Episode, TrainBatch, TrainRollout, group_rollouts +from prime_rl.orchestrator.types import ( + Episode, + TrainBatch, + TrainRollout, + env_name_of, + group_id_of, + group_rollouts, + narrow, + rollouts_of, +) from prime_rl.transport import TrainingSample from prime_rl.utils.logger import get_logger @@ -82,10 +90,10 @@ def __init__( # Keyed by the dispatcher's group UUID. ``(env_name, task_idx)`` # isn't unique — the same task can be re-sampled while an # earlier group is still in flight - self.pending_groups: dict[uuid.UUID, list[Episode]] = defaultdict(list) + self.pending_groups: dict[str, list[Episode]] = defaultdict(list) # Episodes arrived per group — the finalization count (an episode may # add several traces to ``pending_groups`` but counts once here). - self.pending_group_episodes: dict[uuid.UUID, int] = defaultdict(int) + self.pending_group_episodes: dict[str, int] = defaultdict(int) self.pending_batch: list[TrainRollout] = [] # Running payload-token total of ``pending_batch`` (token-batched # runs), kept in sync on append/pop so the readiness check never @@ -133,9 +141,9 @@ async def add(self, episode: Episode) -> TrainBatch | None: pushed (or left) the batch over its threshold. Arrivals into still-incomplete groups never ship a batch. A failed episode brings no rollouts, but still counts toward the group so finalization triggers.""" - group_id = episode.group_id - env_name = episode.env_name - for rollout in episode.rollouts: + group_id = group_id_of(episode) + env_name = env_name_of(episode) + for rollout in rollouts_of(episode): await self.process_rollout(rollout) self.pending_groups[group_id].append(episode) self.pending_group_episodes[group_id] += 1 @@ -174,7 +182,7 @@ async def process_rollout(self, rollout: TrainRollout) -> None: # tokenized — before its group is complete. await self.train_envs.get(rollout.env_name).algorithm.finalize_rollout(rollout) - async def process_group(self, group_id: uuid.UUID) -> None: + async def process_group(self, group_id: str) -> None: """Finalize one GRPO group: drop errored rollouts (the whole group when ``requires_group_scoring`` and any failed), assign advantages, apply the drop policy, append what survives to ``pending_batch``.""" @@ -185,7 +193,7 @@ async def process_group(self, group_id: uuid.UUID) -> None: # Read the group's facts off an episode, not a trace: every episode in it may have # produced none (a whole group cancelled off-policy). env_name = episodes[0].env_name - group = [t for e in episodes for t in e.rollouts] + group = [t for e in episodes for t in rollouts_of(e)] # Window membership follows group finalization, not arrival: a rollout # only becomes observable (metrics / persistence) once its whole group # is finalized, so a batch's window never claims rollouts of a group @@ -206,7 +214,7 @@ async def process_group(self, group_id: uuid.UUID) -> None: return # Untrainable traces carry no samples and must not skew the group baseline. The cohort # stays episodes so the algorithm can still see which attempts shared one. - cohort = [n for e in episodes if (n := e.narrow(lambda r: not r.has_error and r.agent.trainable))] + cohort = [n for e in episodes if (n := narrow(e, lambda r: not r.has_error and r.agent.trainable))] survivors = group_rollouts(cohort) if not survivors: get_logger().debug( diff --git a/src/prime_rl/orchestrator/types.py b/src/prime_rl/orchestrator/types.py index 978f7c22a7..85227d11d0 100644 --- a/src/prime_rl/orchestrator/types.py +++ b/src/prime_rl/orchestrator/types.py @@ -5,12 +5,11 @@ import uuid from collections.abc import Callable, Iterable from dataclasses import dataclass -from typing import TYPE_CHECKING, Any, Generic, Literal, Protocol, Self, cast +from typing import TYPE_CHECKING, Generic, Literal, Protocol, cast import verifiers.v1 as vf from pydantic import ConfigDict, Field from verifiers.v1.task import DataT -from verifiers.v1.trace import EXCLUDE_FIELDS from prime_rl.transport import TrainingSample @@ -123,63 +122,52 @@ def is_trainable(self) -> bool: return bool(advantages) and any(a != 0.0 for a in advantages) -class Episode(vf.WireEpisode): - """The env's own ``vf.Episode`` extended with the facts of the dispatch it came from — the only - thing prime-rl genuinely adds, so the episode itself travels rather than a wrapper around it. - Those fields are ``exclude=True``, so dumping an Episode yields a plain wire episode. +Episode = vf.WireEpisode +"""The env's own episode, unextended: everything prime-rl needs to say about a dispatch has a +place on it already — the env it ran (``env.name``), the group it was planned in (``group``), and +the run it belongs to (``run``, which on the training path carries the policy version and how +stale it got).""" - An episode that produced no traces is not a special case and needs no stand-in rollout: vf - already records why on ``errors`` (its ``run_episode`` puts the exception there and returns the - episode with ``ok`` false), and prime-rl's own outcomes — an off-policy cancel, a task that - raised before reaching the env — are minted the same way. So ``is_empty`` is simply "no - traces", and ``last_error`` says why in one vocabulary for every cause. - Which path it is on is ``run.type``, vf's own discriminator, stamped by the dispatcher when the - episode lands — so there is one episode class and no prime-rl-side kind.""" +def narrow(episode: Episode, keep: Callable[[Rollout], bool]) -> Episode | None: + """The episode with only the traces that pass ``keep``, or ``None`` if none do. A subset stays + a list of episodes rather than a flat trace list, so the episode-level aggregates keep + describing what survived. The kept traces are the same objects, not copies.""" + traces = [t for t in episode.traces if keep(t)] + return episode.model_copy(update={"traces": traces}) if traces else None - model_config = ConfigDict(arbitrary_types_allowed=True) # traces are ``Rollout``s - env_name: str = Field(default="", exclude=True) - """The env as prime-rl names it (the config key), which is not vf's ``env.id``.""" - group_id: uuid.UUID = Field(default_factory=uuid.uuid4, exclude=True) - policy_version: int = Field(default=0, exclude=True) - """The policy that generated it — the thing being trained on one path, measured on the other.""" - off_policy_steps: int = Field(default=0, exclude=True) - """How stale it was by the time it shipped. Always 0 on the eval path, which never trains.""" +def rollouts_of(episode: Episode) -> list[TrainRollout]: + """An episode's traces, typed as the rollouts prime-rl works with. Every trace is built as a + ``TrainRollout`` (``Env.run`` is shared), so the training fields are always reachable — on the + eval path they simply stay empty.""" + return cast(list[TrainRollout], episode.traces) - @property - def rollouts(self) -> list[TrainRollout]: - """The episode's traces, typed as the rollouts prime-rl works with. Every trace is built as - a ``TrainRollout`` (``Env.run`` is shared), so the training fields are always reachable — - on the eval path they simply stay empty.""" - return cast(list[TrainRollout], self.traces) - - def narrow(self, keep: Callable[[Rollout], bool]) -> Self | None: - """This episode with only the traces that pass ``keep``, or ``None`` if none do. A subset - stays a list of episodes rather than a flat trace list, so the episode-level aggregates - keep describing what survived. The kept traces are the same objects, not copies.""" - traces = [t for t in self.traces if keep(t)] - return self.model_copy(update={"traces": traces}) if traces else None - - def to_record(self) -> dict[str, Any]: - """JSON record without raw tensors — the episode form of ``Trace.to_record``, and the unit - ``traces.jsonl`` stores: one episode per line, matching what verifiers writes and what its - ``read_episodes`` expects.""" - return self.model_dump(mode="json", exclude={"traces": {"__all__": EXCLUDE_FIELDS}}) - @property - def is_empty(self) -> bool: - """Whether nothing came back at all — ``last_error`` then carries the reason. Not the - same as failing: an episode can error and still have traces (vf keeps the completed subset - and marks its clean siblings failed), and that failure is accounted for through those - traces. vf's ``ok`` is the success sentinel; this is only "there is nothing here".""" - return not self.traces +def group_id_of(episode: Episode) -> str: + """The group an episode was planned in. The dispatcher plans every episode into one, so this + is always set by the time anything downstream asks.""" + assert episode.group is not None, "the dispatcher plans every episode into a group" + return episode.group.id + + +def env_name_of(episode: Episode) -> str: + """The env as prime-rl names it — its config key, which is not vf's ``env.id``.""" + return episode.env.name or "" + + +def run_of(episode: Episode) -> vf.TrainRunInfo: + """An episode's run record. Every episode the orchestrator produces belongs to the training + run — the ones it trains on and the ones it evaluates along the way — so this is always a + ``TrainRunInfo``, and ``kind`` says which of the two it is.""" + assert isinstance(episode.run, vf.TrainRunInfo), "the dispatcher records the run on arrival" + return episode.run def group_rollouts(episodes: Iterable[Episode]) -> list[TrainRollout]: - """Every trace of a group, flat — the view an algorithm comparing across the whole cohort + """Every trace of a group, flat — the view an algorithm comparing across the whole group wants, where the episode an attempt came from does not matter.""" - return [r for e in episodes for r in e.rollouts] + return [r for e in episodes for r in rollouts_of(e)] @dataclass @@ -198,34 +186,27 @@ class InflightEpisode: off_policy_steps: int = 0 eval_step: int | None = None - def stamp(self, episode: vf.WireEpisode, *, run_id: str, policy_version: int, eval_step: int | None) -> Episode: - """Mint the landed episode: the env's own, carrying the dispatch it came from. The group's - values win over this dispatch's when it is still alive, so they are passed in rather than - read off ``self``. - - The run record is written here because this is where a dispatch's facts become an episode's, - and its ``type`` is what tells the rest of the orchestrator which path the episode is on. A - train episode's step is not known yet — it belongs to whichever batch window is collecting - when it lands, so the main loop fills it in.""" - landed = Episode.model_construct( - **dict(episode), - env_name=self.env_name, - group_id=self.group_id, - policy_version=policy_version, - off_policy_steps=self.off_policy_steps, - ) + def stamp(self, episode: Episode, *, run_id: str, policy_version: int, eval_step: int | None) -> Episode: # noqa: E501 + """Write the dispatch's facts onto the landed episode, in the places the episode already + has for them. The group's values win over this dispatch's when it is still alive, so they + are passed in rather than read off ``self``. + + The run's ``kind`` says which path the episode is on, which is all the rest of the + orchestrator needs to route it. An eval episode knows its step here; an episode to train on + does not — it belongs to whichever batch window is collecting when it lands, so the main + loop fills that in.""" + episode.env.name = self.env_name + episode.group = vf.GroupInfo(id=str(self.group_id)) if self.kind == "eval": assert eval_step is not None, "eval episode missing its step" - run: vf.RunInfo = vf.EvalRunInfo(id=run_id, step=eval_step) - else: - run = vf.TrainRunInfo(id=run_id) - landed.record_run( - run, - env_name=self.env_name, - group_id=str(self.group_id), + episode.run = vf.TrainRunInfo( + id=run_id, + kind=self.kind, + step=eval_step, policy_version=policy_version, + off_policy_steps=0 if self.kind == "eval" else self.off_policy_steps, ) - return landed + return episode @dataclass diff --git a/tests/unit/orchestrator/test_metrics.py b/tests/unit/orchestrator/test_metrics.py index 5fd4961b8a..11f5ec1f37 100644 --- a/tests/unit/orchestrator/test_metrics.py +++ b/tests/unit/orchestrator/test_metrics.py @@ -8,7 +8,14 @@ import verifiers.v1 as vf from prime_rl.orchestrator.metrics import EvalRollouts, Stat, TrainRollouts -from prime_rl.orchestrator.types import Episode, InflightEpisode, Rollout, TrainRollout +from prime_rl.orchestrator.types import ( + Episode, + InflightEpisode, + Rollout, + TrainRollout, + rollouts_of, + run_of, +) from prime_rl.orchestrator.utils import compute_pass_metrics _ids = count() @@ -81,10 +88,15 @@ def mk( def ep(*rollouts, env_name: str = "env", errors=(), group_id="g0", cls=Episode): - """One episode over these traces. ``model_construct`` skips validation so the duck-typed - stand-ins above can stand in for real ones.""" + """One episode over these traces, named and grouped the way the dispatcher stamps it. + ``model_construct`` skips validation so the duck-typed stand-ins above can stand in for real + ones.""" return cls.model_construct( - id=f"e{next(_ids)}", traces=list(rollouts), env_name=env_name, errors=list(errors), group_id=group_id + id=f"e{next(_ids)}", + traces=list(rollouts), + env=vf.EnvInfo(name=env_name), + group=vf.GroupInfo(id=group_id), + errors=list(errors), ) @@ -307,9 +319,9 @@ def test_traceless_episode_keeps_its_reason(): """The failure needs no type of its own: ``failed`` is just "no traces", and the reason is the error vf (or the dispatcher) already put on the episode.""" episode = ep(errors=[vf.Error(type="Cancelled", message="Off-policy cancel")]) - assert episode.is_empty and episode.rollouts == [] + assert not episode.traces and rollouts_of(episode) == [] assert episode.last_error is not None and episode.last_error.type == "Cancelled" - assert not ep(mk()).is_empty + assert ep(mk()).traces # A traceless episode still counts as an episode, but contributes no rollouts. pool = TrainRollouts([ep(mk()), episode]) assert len(pool) == 1 and len(pool.episodes) == 2 @@ -335,12 +347,17 @@ def test_inflight_episode_stamps_what_lands(): kind="train", env_name="rt", group_id=uuid4(), policy_version=3, episodes_owed=1, off_policy_steps=2 ) train = inflight.stamp(wire, run_id="r", policy_version=7, eval_step=None) - assert train.run is not None and train.run.type == "train" - assert train.run.step is None # the batch window it lands in is not known yet - assert (train.env_name, train.policy_version, train.off_policy_steps) == ("rt", 7, 2) # group's version wins + run = run_of(train) + assert (run.kind, run.policy_version, run.off_policy_steps) == ("train", 7, 2) # group's wins + assert run.step is None # the batch window it lands in is not known yet + assert train.env.name == "rt" and train.group is not None - evaluation = replace(inflight, kind="eval").stamp(wire, run_id="r", policy_version=7, eval_step=12) - assert evaluation.run is not None and (evaluation.run.type, evaluation.run.step) == ("eval", 12) + evaluation = replace(inflight, kind="eval").stamp( + vf.WireEpisode.model_construct(id="e", traces=[]), run_id="r", policy_version=7, eval_step=12 + ) + # An online eval belongs to the same training run — same record, told apart by kind. + assert (run_of(evaluation).kind, run_of(evaluation).step) == ("eval", 12) + assert run_of(evaluation).off_policy_steps == 0 # only what trains can go stale with pytest.raises(AssertionError): # an eval episode without its step is not representable replace(inflight, kind="eval").stamp(wire, run_id="r", policy_version=7, eval_step=None) @@ -349,5 +366,6 @@ def test_empty_is_not_the_same_as_failed(): """An episode can error and still have traces — that failure rides the traces and must not be counted as an episode that produced nothing.""" errored = ep(mk(has_error=True), errors=[vf.Error(type="EnvError", message="boom")]) - assert not errored.is_empty and not errored.ok # failed, but something came back - assert errored.rollouts[0].has_error # the failure rides the trace, where the seat's rate sees it + assert errored.traces and not errored.ok # failed, but something came back + # the failure rides the trace, where the seat's rate sees it + assert rollouts_of(errored)[0].has_error From 5ee35aaae29b16e428b3e9195058bb0e52f8b9a7 Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Tue, 4 Aug 2026 21:34:16 +0000 Subject: [PATCH 42/58] refactor(orchestrator)!: the sample monitors take episodes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit log_samples/log_eval_samples took loose traces, which is why a trace had to carry the episode it came from and the env it ran in. They take episodes now: the platform monitor reads the episode id off the episode, and the wandb table gains agent and branch_idx columns so a multi-agent episode reads as its seats, one branch per row. Rollout is an alias for vf.Trace. TrainRollout is what remains — the one place prime-rl extends a verifiers type, for trainer-bound state that has nowhere else to live. Co-Authored-By: Claude Fable 5 --- src/prime_rl/orchestrator/envs.py | 1 - src/prime_rl/orchestrator/orchestrator.py | 4 +- src/prime_rl/orchestrator/types.py | 29 ++---- src/prime_rl/utils/monitor/base.py | 10 +- src/prime_rl/utils/monitor/file.py | 6 +- src/prime_rl/utils/monitor/multi.py | 10 +- src/prime_rl/utils/monitor/prime.py | 27 ++--- src/prime_rl/utils/monitor/wandb.py | 117 +++++++++++++--------- tests/unit/orchestrator/test_metrics.py | 6 +- tests/unit/utils/test_prime_monitor.py | 11 +- 10 files changed, 123 insertions(+), 98 deletions(-) diff --git a/src/prime_rl/orchestrator/envs.py b/src/prime_rl/orchestrator/envs.py index f796697d31..27f4d5c916 100644 --- a/src/prime_rl/orchestrator/envs.py +++ b/src/prime_rl/orchestrator/envs.py @@ -224,7 +224,6 @@ async def run( ) rollouts = [ROLLOUT_TYPE.model_construct(**dict(wire)) for wire in episode.traces] for rollout in rollouts: - rollout.episode_id = episode.id if not episode.ok and rollout.ok: error = episode.last_error or vf.Error( type="EpisodeFailed", message="A sibling trace in this episode failed" diff --git a/src/prime_rl/orchestrator/orchestrator.py b/src/prime_rl/orchestrator/orchestrator.py index 18d4522227..cc04af4e52 100644 --- a/src/prime_rl/orchestrator/orchestrator.py +++ b/src/prime_rl/orchestrator/orchestrator.py @@ -683,7 +683,7 @@ async def finalize_train_batch(self, batch: TrainBatch) -> None: metrics[f"dropped/all/{name}/rate"] = count / self.train_sink.pre_filter_seen self.monitor.log(metrics, step=step) self.wait_for_policy_time = 0.0 - self.monitor.log_samples(effective.rollouts, step=step) + self.monitor.log_samples(effective.episodes, step=step) self.monitor.log_distributions( distributions={ "rewards": [r.reward for r in effective], @@ -851,7 +851,7 @@ async def finalize_eval_batch(self, batch: EvalBatch) -> None: await asyncio.to_thread( save_episodes, records, get_trace_path(self.config.output_dir, batch.step, "eval", "effective") ) - self.monitor.log_eval_samples(batch.rollouts, env_name=batch.env_name, step=batch.step) + self.monitor.log_eval_samples(batch.rollouts.episodes, env_name=batch.env_name, step=batch.step) policy_versions = {run_of(e).policy_version for e in batch.rollouts.episodes} policy_version = min(policy_versions) if len(policy_versions) > 1: diff --git a/src/prime_rl/orchestrator/types.py b/src/prime_rl/orchestrator/types.py index 85227d11d0..1fcf07e298 100644 --- a/src/prime_rl/orchestrator/types.py +++ b/src/prime_rl/orchestrator/types.py @@ -58,30 +58,23 @@ class GroupState: policy_version_at_start: int = 0 -class Rollout(vf.Trace[DataT], Generic[DataT]): - """A completed rollout: the env's typed ``vf.Trace`` *is* the rollout, carrying only the links - a consumer that works in loose traces — the sample monitors — needs to place one back among its - peers. Anything episode-scoped is read off the ``Episode``, which is the atomic unit everything - else passes around. All added fields are ``exclude=True``, so dumping a Rollout yields a plain - trace on the wire; ``vf.Trace.record_run`` mirrors them into ``info`` on arrival so the on-disk - records stay fully placeable. +Rollout = vf.Trace +"""The env's own trace, unextended — what an eval rollout is, and the read side of a training one. +Where a trace sits (its episode, its group, its env) is the episode's to say.""" - It is also the currency the scoring hooks receive: a hook reads the trace directly - (``rollout.reward``, ``rollout.nodes``, ``rollout.num_turns``).""" - - env_name: str = Field(default="", exclude=True) - # Links the traces of one episode; stamped into ``info`` on arrival so - # saved records keep their grouping. - episode_id: str = Field(default="", exclude=True) +class TrainRollout(vf.Trace[DataT], Generic[DataT]): + """A rollout on the training path — the one place prime-rl extends a verifiers type, because + trainer-bound state has nowhere else to live: the samples built from its branches, what the + degeneracy detectors measured on it, and whether the drop policy kept it. All of it is + ``exclude=True``, so dumping one yields a plain trace on the wire. -class TrainRollout(Rollout[DataT], Generic[DataT]): - """A rollout on the training path, which alone carries training state: the trainer-bound - samples built from its branches, the credit assigned over them, and the filter verdicts. Eval - rollouts have none of this, so they are plain ``Rollout``\\ s and can't be asked for it.""" + ``env_name`` rides along because a sample is routed by it (the trainer's per-env loss config) + and the sink's pending batch is a trace list. Eval rollouts carry the fields unset.""" model_config = ConfigDict(arbitrary_types_allowed=True) # ``samples`` holds msgspec structs + env_name: str = Field(default="", exclude=True) samples: list[TrainingSample] = Field(default_factory=list, exclude=True) detections: dict[str, bool] = Field(default_factory=dict, exclude=True) """What each degeneracy detector measured on this trace — a measurement, not a verdict.""" diff --git a/src/prime_rl/utils/monitor/base.py b/src/prime_rl/utils/monitor/base.py index 7a0a51a6e4..cbe63df6ff 100644 --- a/src/prime_rl/utils/monitor/base.py +++ b/src/prime_rl/utils/monitor/base.py @@ -6,7 +6,7 @@ from typing import TYPE_CHECKING, Any if TYPE_CHECKING: - from prime_rl.orchestrator.types import Rollout + from prime_rl.orchestrator.types import Episode _DROPPED_JSON_VALUE = object() @@ -85,11 +85,11 @@ def log(self, metrics: dict[str, Any], step: int) -> None: pass @abstractmethod - def log_samples(self, rollouts: list[Rollout], step: int) -> None: + def log_samples(self, episodes: list[Episode], step: int) -> None: pass @abstractmethod - def log_eval_samples(self, rollouts: list[Rollout], env_name: str, step: int) -> None: + def log_eval_samples(self, episodes: list[Episode], env_name: str, step: int) -> None: pass @abstractmethod @@ -118,10 +118,10 @@ def log(self, metrics: dict[str, Any], step: int) -> None: else: self.history = [metrics] - def log_samples(self, rollouts: list[Rollout], step: int) -> None: + def log_samples(self, episodes: list[Episode], step: int) -> None: pass - def log_eval_samples(self, rollouts: list[Rollout], env_name: str, step: int) -> None: + def log_eval_samples(self, episodes: list[Episode], env_name: str, step: int) -> None: pass def save_final_summary(self, filename: str = "final_summary.json") -> None: diff --git a/src/prime_rl/utils/monitor/file.py b/src/prime_rl/utils/monitor/file.py index 08a5240bd9..a44e7a9e0a 100644 --- a/src/prime_rl/utils/monitor/file.py +++ b/src/prime_rl/utils/monitor/file.py @@ -12,7 +12,7 @@ from prime_rl.utils.monitor.base import Monitor, drop_non_finite_json_values if TYPE_CHECKING: - from prime_rl.orchestrator.types import Rollout + from prime_rl.orchestrator.types import Episode class FileMonitor(Monitor): @@ -79,10 +79,10 @@ def log(self, metrics: dict[str, Any], step: int) -> None: self._file.write(json.dumps(row)) self._file.write("\n") - def log_samples(self, rollouts: list[Rollout], step: int) -> None: + def log_samples(self, episodes: list[Episode], step: int) -> None: pass - def log_eval_samples(self, rollouts: list[Rollout], env_name: str, step: int) -> None: + def log_eval_samples(self, episodes: list[Episode], env_name: str, step: int) -> None: pass def log_distributions(self, distributions: dict[str, list[float]], step: int) -> None: diff --git a/src/prime_rl/utils/monitor/multi.py b/src/prime_rl/utils/monitor/multi.py index 043953546d..11c64e1e84 100644 --- a/src/prime_rl/utils/monitor/multi.py +++ b/src/prime_rl/utils/monitor/multi.py @@ -7,7 +7,7 @@ from prime_rl.utils.monitor.prime import PrimeMonitor if TYPE_CHECKING: - from prime_rl.orchestrator.types import Rollout + from prime_rl.orchestrator.types import Episode class MultiMonitor(Monitor): @@ -33,17 +33,17 @@ def log(self, metrics: dict[str, Any], step: int) -> None: except Exception as e: self.logger.warning(f"Failed to log metrics to {monitor.__class__.__name__}: {e}") - def log_samples(self, rollouts: list[Rollout], step: int) -> None: + def log_samples(self, episodes: list[Episode], step: int) -> None: for monitor in self.monitors: try: - monitor.log_samples(rollouts=rollouts, step=step) + monitor.log_samples(episodes=episodes, step=step) except Exception as e: self.logger.warning(f"Failed to log samples to {monitor.__class__.__name__}: {e}") - def log_eval_samples(self, rollouts: list[Rollout], env_name: str, step: int) -> None: + def log_eval_samples(self, episodes: list[Episode], env_name: str, step: int) -> None: for monitor in self.monitors: try: - monitor.log_eval_samples(rollouts=rollouts, env_name=env_name, step=step) + monitor.log_eval_samples(episodes=episodes, env_name=env_name, step=step) except Exception as e: self.logger.warning(f"Failed to log eval samples to {monitor.__class__.__name__}: {e}") diff --git a/src/prime_rl/utils/monitor/prime.py b/src/prime_rl/utils/monitor/prime.py index eee454a9aa..fe5460b858 100644 --- a/src/prime_rl/utils/monitor/prime.py +++ b/src/prime_rl/utils/monitor/prime.py @@ -23,7 +23,7 @@ from prime_rl.utils.monitor.base import Monitor, drop_non_finite_json_values, sample_items_for_logging if TYPE_CHECKING: - from prime_rl.orchestrator.types import Rollout + from prime_rl.orchestrator.types import Episode _SAMPLE_SCHEMA = pa.schema( @@ -248,7 +248,7 @@ def log(self, metrics: dict[str, Any], step: int) -> None: }, ) - def log_samples(self, rollouts: list[Rollout], step: int) -> None: + def log_samples(self, episodes: list[Episode], step: int) -> None: """Logs rollouts to Prime Intellect API using presigned URLs for direct R2 upload.""" if not self.is_master: return @@ -262,11 +262,11 @@ def log_samples(self, rollouts: list[Rollout], step: int) -> None: ): return - rollouts = sample_items_for_logging( - rollouts, + episodes = sample_items_for_logging( + episodes, self.config.log_extras.sample_ratio, ) - if not rollouts: + if not episodes: return assert self.last_log_samples_step <= step, "Step must be greater than last logged step" @@ -276,7 +276,7 @@ def log_samples(self, rollouts: list[Rollout], step: int) -> None: self.logger.info(f"Logging {len(rollouts)} samples to Prime Intellect API at step {step}") start_time = time.perf_counter() - parquet_bytes = self._rollouts_to_parquet_bytes(rollouts, step) + parquet_bytes = self._rollouts_to_parquet_bytes(episodes, step) if not parquet_bytes: self.logger.warning(f"No samples to log at step {step}") @@ -291,8 +291,9 @@ def log_samples(self, rollouts: list[Rollout], step: int) -> None: f"Initiated samples upload at step {step} to Prime Intellect API in {time.perf_counter() - start_time:.2f}s" ) - def _rollouts_to_parquet_bytes(self, rollouts: list[Rollout], step: int) -> bytes | None: - """Convert rollouts to Parquet bytes for upload. One row per rollout. The conversation + def _rollouts_to_parquet_bytes(self, episodes: list[Episode], step: int) -> bytes | None: + """Convert episodes to Parquet bytes for upload. One row per trace, carrying the episode + it belongs to. The conversation is the unit (no prompt/completion split — meaningless mid-branch): `completion` is the last branch's messages and `trajectory` is one message list per branch. Shares `verifiers.v1.utils.platform.trace_to_sample` with verifiers' eval `--push`, so a training-run @@ -301,8 +302,10 @@ def _rollouts_to_parquet_bytes(self, rollouts: list[Rollout], step: int) -> byte now = datetime.now(timezone.utc) rows = [] - for sample_id, rollout in enumerate(rollouts): - sample = trace_to_sample(rollout, rollout_number=sample_id + 1, episode_id=rollout.episode_id or None) + for sample_id, (episode, rollout) in enumerate( + (episode, trace) for episode in episodes for trace in episode.traces + ): + sample = trace_to_sample(rollout, rollout_number=sample_id + 1, episode_id=episode.id) trajectory = sample["trajectory"] if not trajectory: # no branches (e.g. a rollout that errored before any message) continue @@ -326,7 +329,7 @@ def _rollouts_to_parquet_bytes(self, rollouts: list[Rollout], step: int) -> byte "completion": json.dumps(sample["completion"]), "trajectory": json.dumps(trajectory), "answer": "", - "env_name": rollout.env_name, + "env_name": episode.env.name or "", "task": json.dumps(sample["task"]), "info": json.dumps(rollout.info), "reward": sample["reward"], @@ -447,7 +450,7 @@ async def _confirm_samples_upload(self, step: int, s3_key: str, max_retries: int await asyncio.sleep(delay) return False - def log_eval_samples(self, rollouts: list[Rollout], env_name: str, step: int) -> None: + def log_eval_samples(self, episodes: list[Episode], env_name: str, step: int) -> None: pass def log_distributions(self, distributions: dict[str, list[float]], step: int) -> None: diff --git a/src/prime_rl/utils/monitor/wandb.py b/src/prime_rl/utils/monitor/wandb.py index 7340b3c6de..c87e1274af 100644 --- a/src/prime_rl/utils/monitor/wandb.py +++ b/src/prime_rl/utils/monitor/wandb.py @@ -23,7 +23,7 @@ from prime_rl.utils.monitor.base import Monitor, sample_items_for_logging if TYPE_CHECKING: - from prime_rl.orchestrator.types import Rollout + from prime_rl.orchestrator.types import Episode def _loggable_task(task) -> str: @@ -159,13 +159,32 @@ def init_wandb(max_retries: int): if config is not None and isinstance(config, WandbWithExtrasConfig) and config.log_extras: if config.log_extras.samples: self.last_log_samples_step = -1 - self.samples_cols = ["step", "env_name", "task", "task_idx", "messages", "input_ids", "reward"] + self.samples_cols = [ + "step", + "env_name", + "task", + "task_idx", + "agent", + "branch_idx", + "messages", + "input_ids", + "reward", + ] self.samples_table = wandb.Table( columns=self.samples_cols, log_mode="INCREMENTAL", ) self.tokenizer = tokenizer - self.eval_samples_cols = ["step", "env", "task", "task_idx", "completion", "reward"] + self.eval_samples_cols = [ + "step", + "env", + "task", + "task_idx", + "agent", + "branch_idx", + "completion", + "reward", + ] self.eval_samples_table = wandb.Table( columns=self.eval_samples_cols, log_mode="INCREMENTAL", @@ -189,8 +208,10 @@ def log(self, metrics: dict[str, Any], step: int) -> None: return wandb.log({**metrics, "step": step}) - def log_samples(self, rollouts: list[Rollout], step: int) -> None: - """Logs rollouts to W&B table.""" + def log_samples(self, episodes: list[Episode], step: int) -> None: + """Log a sample of episodes to the W&B table, one row per branch — the agent that ran it + and which of its branches the row is are columns, so a multi-agent episode reads as its + seats rather than as one blurred cell.""" if not self.is_master: return if ( @@ -203,45 +224,47 @@ def log_samples(self, rollouts: list[Rollout], step: int) -> None: # Do not log samples if not enabled or not log interval step return - rollouts = sample_items_for_logging( - rollouts, + episodes = sample_items_for_logging( + episodes, self.config.log_extras.sample_ratio, ) - if not rollouts: + if not episodes: return assert self.tokenizer is not None, "Tokenizer is required for sample logging" assert self.last_log_samples_step <= step, "Step must be greater than last logged step" assert self.logger is not None, "Logger is required for sample logging" - self.logger.info(f"Logging {len(rollouts)} samples to W&B table at step {step}") + self.logger.info(f"Logging {len(episodes)} episodes to W&B table at step {step}") start_time = time.perf_counter() - for rollout in rollouts: - trace = rollout - for branch in trace.branches: - token_ids = branch.token_ids - if not token_ids: - continue - sample = { - "step": step, - "env_name": rollout.env_name, - "task": _loggable_task(trace.task.data), - "task_idx": trace.task.data.idx, - "messages": self.tokenizer.decode(token_ids), - "input_ids": str(token_ids), - "reward": trace.reward, - } - assert list(sample.keys()) == self.samples_cols, ( - "Order of columns in the table must be the same as order of the keys here" - ) - self.samples_table.add_data(*sample.values()) + for episode in episodes: + for trace in episode.traces: + for branch_idx, branch in enumerate(trace.branches): + token_ids = branch.token_ids + if not token_ids: + continue + sample = { + "step": step, + "env_name": episode.env.name or "", + "task": _loggable_task(trace.task.data), + "task_idx": trace.task.data.idx, + "agent": trace.agent.name, + "branch_idx": branch_idx, + "messages": self.tokenizer.decode(token_ids), + "input_ids": str(token_ids), + "reward": trace.reward, + } + assert list(sample.keys()) == self.samples_cols, ( + "Order of columns in the table must be the same as order of the keys here" + ) + self.samples_table.add_data(*sample.values()) wandb.log({"samples": self.samples_table, "step": step}) self.last_log_samples_step = step self.logger.debug(f"Logged samples at step {step} to W&B table in {time.perf_counter() - start_time:.2f}s") - def log_eval_samples(self, rollouts: list[Rollout], env_name: str, step: int) -> None: + def log_eval_samples(self, episodes: list[Episode], env_name: str, step: int) -> None: """Logs eval rollouts to a separate W&B table.""" if not self.is_master: return @@ -253,23 +276,25 @@ def log_eval_samples(self, rollouts: list[Rollout], env_name: str, step: int) -> ): return - for rollout in rollouts: - trace = rollout - for branch in trace.branches: - # Eval runs the openai client (no token ids), so show the assistant message - # content rather than decoded tokens. - completion = "".join(m.content or "" for m in branch.messages if m.role == "assistant") - if not completion: - continue - sample = { - "step": step, - "env": env_name, - "task": _loggable_task(trace.task.data), - "task_idx": trace.task.data.idx, - "completion": completion, - "reward": trace.reward, - } - self.eval_samples_table.add_data(*sample.values()) + for episode in episodes: + for trace in episode.traces: + for branch_idx, branch in enumerate(trace.branches): + # Eval runs the openai client (no token ids), so show the assistant message + # content rather than decoded tokens. + completion = "".join(m.content or "" for m in branch.messages if m.role == "assistant") + if not completion: + continue + sample = { + "step": step, + "env": env_name, + "task": _loggable_task(trace.task.data), + "task_idx": trace.task.data.idx, + "agent": trace.agent.name, + "branch_idx": branch_idx, + "completion": completion, + "reward": trace.reward, + } + self.eval_samples_table.add_data(*sample.values()) wandb.log({"eval/samples": self.eval_samples_table, "step": step}) diff --git a/tests/unit/orchestrator/test_metrics.py b/tests/unit/orchestrator/test_metrics.py index 11f5ec1f37..8f7d4c9b34 100644 --- a/tests/unit/orchestrator/test_metrics.py +++ b/tests/unit/orchestrator/test_metrics.py @@ -332,10 +332,10 @@ def test_training_state_is_train_only(): """Only a train rollout carries trainer-bound state; an eval trace has no field for it. Credit is not among them — it lives on the graph's nodes, which every trace has.""" assert {"samples", "is_filtered", "detections"} <= set(TrainRollout.model_fields) - assert not {"samples", "is_filtered", "detections"} & set(Rollout.model_fields) assert "advantages" not in TrainRollout.model_fields # derived from the nodes - assert {"env_name", "episode_id"} <= set(Rollout.model_fields) # links stay shared - assert "group_id" not in Rollout.model_fields # the example a trace answered is the episode's + # An eval trace is the env's own, unextended — where a trace sits is the episode's to say. + assert Rollout is vf.Trace + assert not {"samples", "is_filtered", "detections", "group_id", "episode_id"} & set(Rollout.model_fields) def test_inflight_episode_stamps_what_lands(): diff --git a/tests/unit/utils/test_prime_monitor.py b/tests/unit/utils/test_prime_monitor.py index 799493f549..92787f369f 100644 --- a/tests/unit/utils/test_prime_monitor.py +++ b/tests/unit/utils/test_prime_monitor.py @@ -46,14 +46,19 @@ def _build_rollout(*, example_id: int, reward: float, task: str) -> TrainRollout return rollout +def _episode(*rollouts: TrainRollout, env_name: str = "task-a") -> vf.WireEpisode: + """The unit the monitor uploads: one episode, whose traces become the rows.""" + return vf.WireEpisode.model_construct(traces=list(rollouts), env=vf.EnvInfo(name=env_name)) + + def test_rollouts_to_parquet_bytes_preserves_all_rollouts_and_ids(): monitor = _new_monitor() monitor.run_id = "run-123" parquet_bytes = monitor._rollouts_to_parquet_bytes( [ - _build_rollout(example_id=101, reward=1.0, task="task-a"), - _build_rollout(example_id=202, reward=0.0, task="task-b"), + _episode(_build_rollout(example_id=101, reward=1.0, task="task-a")), + _episode(_build_rollout(example_id=202, reward=0.0, task="task-b"), env_name="task-b"), ], step=7, ) @@ -86,7 +91,7 @@ def test_rollouts_to_parquet_bytes_skips_rollouts_without_trajectory(): assert rollout_without_branches.branches == [] parquet_bytes = monitor._rollouts_to_parquet_bytes( - [rollout_with_branches, rollout_without_branches], + [_episode(rollout_with_branches, rollout_without_branches)], step=3, ) From df6372cf34f28bc4ee1c727823975c4bf4ade739 Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Tue, 4 Aug 2026 21:39:41 +0000 Subject: [PATCH 43/58] fix(orchestrator): the eval sink places an episode by its train-run step MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Online eval moved onto the training run, but the sink still asserted on EvalRunInfo. Nothing unit-tested the sink, so only a live run caught it — covered now. Co-Authored-By: Claude Fable 5 --- src/prime_rl/orchestrator/eval_sink.py | 20 ++++++++++++++------ tests/unit/orchestrator/test_metrics.py | 20 ++++++++++++++++++++ 2 files changed, 34 insertions(+), 6 deletions(-) diff --git a/src/prime_rl/orchestrator/eval_sink.py b/src/prime_rl/orchestrator/eval_sink.py index aea0649983..1062bb5c63 100644 --- a/src/prime_rl/orchestrator/eval_sink.py +++ b/src/prime_rl/orchestrator/eval_sink.py @@ -16,19 +16,27 @@ from collections import defaultdict -import verifiers.v1 as vf - from prime_rl.orchestrator.envs import EvalEnvs from prime_rl.orchestrator.metrics import EvalRollouts -from prime_rl.orchestrator.types import Episode, EvalBatch, Rollout, env_name_of, group_id_of, rollouts_of +from prime_rl.orchestrator.types import ( + Episode, + EvalBatch, + Rollout, + env_name_of, + group_id_of, + rollouts_of, + run_of, +) from prime_rl.utils.logger import get_logger def eval_step_of(episode: Episode) -> int: """The eval epoch an episode belongs to, off the run the dispatcher recorded when it landed. - Only the eval path has one, which is why this lives here and not on ``Episode``.""" - assert isinstance(episode.run, vf.EvalRunInfo) and episode.run.step is not None - return episode.run.step + An online eval belongs to the training run, so its step is known from the start — unlike an + episode to train on, whose step is the window it lands in.""" + run = run_of(episode) + assert run.kind == "eval" and run.step is not None, "not an eval episode" + return run.step class EvalSink: diff --git a/tests/unit/orchestrator/test_metrics.py b/tests/unit/orchestrator/test_metrics.py index 8f7d4c9b34..2f70d2986c 100644 --- a/tests/unit/orchestrator/test_metrics.py +++ b/tests/unit/orchestrator/test_metrics.py @@ -7,6 +7,7 @@ import pytest import verifiers.v1 as vf +from prime_rl.orchestrator.eval_sink import eval_step_of from prime_rl.orchestrator.metrics import EvalRollouts, Stat, TrainRollouts from prime_rl.orchestrator.types import ( Episode, @@ -362,6 +363,25 @@ def test_inflight_episode_stamps_what_lands(): replace(inflight, kind="eval").stamp(wire, run_id="r", policy_version=7, eval_step=None) +def test_eval_sink_reads_the_epoch_a_stamped_episode_landed_in(): + """The eval sink places an episode by the step on its run. Regression: online eval moved onto + the training run (kind="eval"), and nothing unit-tested the sink, so the sink kept asserting on + the old record and only a live run caught it.""" + inflight = InflightEpisode( + kind="eval", env_name="rt", group_id=uuid4(), policy_version=3, episodes_owed=1, eval_step=12 + ) + landed = inflight.stamp( + vf.WireEpisode.model_construct(id="e", traces=[]), run_id="r", policy_version=3, eval_step=12 + ) + assert eval_step_of(landed) == 12 + + trained_on = replace(inflight, kind="train").stamp( + vf.WireEpisode.model_construct(id="e2", traces=[]), run_id="r", policy_version=3, eval_step=None + ) + with pytest.raises(AssertionError): # an episode to train on is not placed by an eval epoch + eval_step_of(trained_on) + + def test_empty_is_not_the_same_as_failed(): """An episode can error and still have traces — that failure rides the traces and must not be counted as an episode that produced nothing.""" From 5f41d94b77fc4b01a5925bdfaadef324697dad7d Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Tue, 4 Aug 2026 21:42:06 +0000 Subject: [PATCH 44/58] fix(orchestrator): read the episode's env and group through the accessors MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Five reads still went at the fields the subclass used to have, which a pydantic model answers with AttributeError at runtime — in the sinks and the ship path, none of which the unit suite covers. Co-Authored-By: Claude Fable 5 --- src/prime_rl/orchestrator/eval_sink.py | 4 ++-- src/prime_rl/orchestrator/orchestrator.py | 6 ++++-- src/prime_rl/orchestrator/train_sink.py | 4 ++-- 3 files changed, 8 insertions(+), 6 deletions(-) diff --git a/src/prime_rl/orchestrator/eval_sink.py b/src/prime_rl/orchestrator/eval_sink.py index 1062bb5c63..bb9cd98746 100644 --- a/src/prime_rl/orchestrator/eval_sink.py +++ b/src/prime_rl/orchestrator/eval_sink.py @@ -86,7 +86,7 @@ def batch_progress(self) -> list[tuple[str, int, int, int, int]]: for group_id, episodes in self.pending_groups.items(): if not episodes: continue - env_name = episodes[0].env_name + env_name = env_name_of(episodes[0]) if self.eval_envs.get(env_name).requires_group_scoring: continue bkey = (env_name, eval_step_of(episodes[0])) @@ -120,7 +120,7 @@ def process_group(self, group_id: str) -> None: return # Read the group's facts off an episode, not a trace: every episode in it may have # produced none (a whole group cancelled off-policy). - env_name = finished[0].env_name + env_name = env_name_of(finished[0]) eval_step = eval_step_of(finished[0]) group = [t for e in finished for t in rollouts_of(e)] task_idx = group[0].task.data.idx if group else -1 diff --git a/src/prime_rl/orchestrator/orchestrator.py b/src/prime_rl/orchestrator/orchestrator.py index cc04af4e52..5b2304abaf 100644 --- a/src/prime_rl/orchestrator/orchestrator.py +++ b/src/prime_rl/orchestrator/orchestrator.py @@ -61,6 +61,8 @@ Policy, Progress, TrainBatch, + env_name_of, + group_id_of, run_of, ) from prime_rl.orchestrator.utils import ( @@ -626,7 +628,7 @@ async def finalize_train_batch(self, batch: TrainBatch) -> None: # counter, which only sees weight updates during generation. Frozen- # sourced rollouts stay 0 (their sampler doesn't follow the policy). for train_episode in batch.rollouts.episodes: - if self.train_envs.get(train_episode.env_name).sampler.samples_from_live_policy: + if self.train_envs.get(env_name_of(train_episode)).sampler.samples_from_live_policy: run = run_of(train_episode) run.off_policy_steps = (step - 1) - (run.policy_version or 0) @@ -660,7 +662,7 @@ async def finalize_train_batch(self, batch: TrainBatch) -> None: num_input = sum(r.num_input_tokens for r in effective) num_output = sum(r.num_output_tokens for r in effective) num_rollouts = len(batch.rollouts) - num_unique_examples = len({e.group_id for e in batch.rollouts.episodes}) + num_unique_examples = len({group_id_of(e) for e in batch.rollouts.episodes}) metrics |= { "progress/tokens": num_tokens, "progress/input_tokens": num_input, diff --git a/src/prime_rl/orchestrator/train_sink.py b/src/prime_rl/orchestrator/train_sink.py index 5780580a99..c64172bc30 100644 --- a/src/prime_rl/orchestrator/train_sink.py +++ b/src/prime_rl/orchestrator/train_sink.py @@ -124,7 +124,7 @@ def buffered_count(self) -> int: return sum( self.pending_group_episodes.get(group_id, 0) for group_id, episodes in self.pending_groups.items() - if episodes and not self.train_envs.get(episodes[0].env_name).requires_group_scoring + if episodes and not self.train_envs.get(env_name_of(episodes[0])).requires_group_scoring ) def pending_batch_by_env(self) -> dict[str, int]: @@ -192,7 +192,7 @@ async def process_group(self, group_id: str) -> None: return # Read the group's facts off an episode, not a trace: every episode in it may have # produced none (a whole group cancelled off-policy). - env_name = episodes[0].env_name + env_name = env_name_of(episodes[0]) group = [t for e in episodes for t in rollouts_of(e)] # Window membership follows group finalization, not arrival: a rollout # only becomes observable (metrics / persistence) once its whole group From 7c11b46b1eff8a3ffd0d52cf12160a893bcaa441 Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Tue, 4 Aug 2026 21:58:06 +0000 Subject: [PATCH 45/58] refactor(orchestrator)!: staleness is derived from the policy span MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The dispatcher counted weight updates per in-flight episode, then the main loop threw that away and recomputed staleness at ship — two mechanisms writing one field. The dispatcher now records the span generation covered and every reading derives from it: the cancel check subtracts against the live version, the in-flight gauges do the same, and off_policy_steps is a property of the run. Eval is measured the same way rather than pinned to 0, and a frozen sampler records no span at all, so its staleness reads None instead of a zero indistinguishable from fresh. Co-Authored-By: Claude Fable 5 --- deps/verifiers | 2 +- src/prime_rl/orchestrator/dispatcher.py | 39 ++++++++++++++++------- src/prime_rl/orchestrator/orchestrator.py | 15 ++------- src/prime_rl/orchestrator/types.py | 11 ++----- tests/unit/orchestrator/test_metrics.py | 28 ++++++++-------- 5 files changed, 47 insertions(+), 48 deletions(-) diff --git a/deps/verifiers b/deps/verifiers index 70490aa3ef..8c318d0a4b 160000 --- a/deps/verifiers +++ b/deps/verifiers @@ -1 +1 @@ -Subproject commit 70490aa3efe7080e1a4b138da58eb63b295b1771 +Subproject commit 8c318d0a4bc9e8a145267b96701cf02bebb7fc01 diff --git a/src/prime_rl/orchestrator/dispatcher.py b/src/prime_rl/orchestrator/dispatcher.py index 25f9f15598..f8d85fea49 100644 --- a/src/prime_rl/orchestrator/dispatcher.py +++ b/src/prime_rl/orchestrator/dispatcher.py @@ -15,8 +15,8 @@ source's emptiness), so in-flight rollouts of the opposite kind drain naturally on either side of an eval boundary. - ``on_version_pending`` (called by the watcher before the engines pause for - the weight update) bumps ``off_policy_steps`` on in-flight train rollouts and - drops groups past ``max_off_policy_steps``. + the weight update) drops in-flight train groups whose generating policy would + fall further behind than ``max_off_policy_steps``. Eval rollouts are measurements for the policy version they started with, so they are allowed to finish even if training advances. Train rollouts sampled from a frozen model never age — their sampler doesn't change @@ -227,15 +227,19 @@ def disable_train_scheduling(self) -> None: triggered eval drain naturally.""" self.train_scheduling_disabled = True + def _inflight_lag(self) -> list[int]: + """How far behind the live policy each in-flight train dispatch has fallen.""" + return [self.policy.version - m.policy_version for m in self.inflight.values() if m.kind == "train"] + @property def max_off_policy_level(self) -> int: - steps = [m.off_policy_steps for m in self.inflight.values() if m.kind == "train"] - return max(steps) if steps else 0 + lag = self._inflight_lag() + return max(lag) if lag else 0 @property def mean_off_policy_level(self) -> float: - steps = [m.off_policy_steps for m in self.inflight.values() if m.kind == "train"] - return sum(steps) / len(steps) if steps else 0.0 + lag = self._inflight_lag() + return sum(lag) / len(lag) if lag else 0.0 # ── lifecycle ────────────────────────────────────────────────────────── @@ -271,8 +275,13 @@ async def stop(self) -> None: await safe_cancel(self.task) self.task = None + def samples_from_live_policy(self, meta: InflightEpisode) -> bool: + """Whether this dispatch follows the live policy at all. A frozen sampler does not, so it + never ages and has no version to be behind.""" + return meta.kind == "eval" or self.train_envs.get(meta.env_name).sampler.samples_from_live_policy + async def on_version_pending(self, step: int) -> None: - """Bump off-policy counters and drop groups past + """Drop groups whose generating policy is now further behind than ``max_off_policy_steps`` (drop_group emits ``Cancelled`` markers so the sink still finalizes the partial group). Eval rollouts are not aged because they are tied to their start-time policy version. @@ -288,10 +297,11 @@ async def on_version_pending(self, step: int) -> None: continue # Frozen-sourced rollouts never go stale — their sampler doesn't # change with policy updates. - if not self.train_envs.get(meta.env_name).sampler.samples_from_live_policy: + if not self.samples_from_live_policy(meta): continue - meta.off_policy_steps += 1 - if meta.off_policy_steps > self.max_off_policy_steps: + # The live version is about to become ``step``'s, so this dispatch will be one + # further behind than it is now. + if (self.policy.version + 1) - meta.policy_version > self.max_off_policy_steps: stale_groups.add(meta.group_id) for gid in stale_groups: @@ -564,9 +574,14 @@ async def emit_episode( for rollout in episode.traces: rollout.env_name = meta.env_name - await self.out_q.put( - meta.stamp(episode, run_id=self.run_id, policy_version=policy_version, eval_step=eval_step) + # Generation is over, so the span closes at whatever version is live now: an episode that + # outlived an update spans more than one. + policy = ( + vf.PolicySpan(start=policy_version, end=self.policy.version) + if self.samples_from_live_policy(meta) + else None ) + await self.out_q.put(meta.stamp(episode, run_id=self.run_id, policy=policy, eval_step=eval_step)) async def emit_failed_episodes(self, meta: InflightEpisode, group: GroupState | None, error: vf.Error) -> None: """Emit one traceless episode per rollout the task owed, so the sink still counts its way to diff --git a/src/prime_rl/orchestrator/orchestrator.py b/src/prime_rl/orchestrator/orchestrator.py index 5b2304abaf..7f020fa9e8 100644 --- a/src/prime_rl/orchestrator/orchestrator.py +++ b/src/prime_rl/orchestrator/orchestrator.py @@ -61,7 +61,6 @@ Policy, Progress, TrainBatch, - env_name_of, group_id_of, run_of, ) @@ -622,16 +621,6 @@ async def finalize_train_batch(self, batch: TrainBatch) -> None: await self.version_advanced.wait() self.wait_for_policy_time += time.perf_counter() - hold_start - # Stamp each rollout's true staleness: batch ``step`` trains on policy - # v{step-1}, so a rollout generated from v{k} is (step-1)-k versions - # off-policy — queue time included, unlike the dispatcher's in-flight - # counter, which only sees weight updates during generation. Frozen- - # sourced rollouts stay 0 (their sampler doesn't follow the policy). - for train_episode in batch.rollouts.episodes: - if self.train_envs.get(env_name_of(train_episode)).sampler.samples_from_live_policy: - run = run_of(train_episode) - run.off_policy_steps = (step - 1) - (run.policy_version or 0) - # The effective (clean, trained-on) subset lands in the per-step ``effective`` trace file # at ship time; the full arrival window already streamed into ``all`` on arrival. # to_record drops the per-node training tensors — they're for training, not the rollout @@ -809,7 +798,7 @@ def log_train_batch(self, batch: TrainBatch, *, step: int, step_time: float) -> n_effective = len(effective) n_trainable = sum(1 for r in effective if r.is_trainable) trainable_rate = (n_trainable / n_effective) if n_effective else 0.0 - max_off_policy = max((run_of(e).off_policy_steps for e in effective.episodes), default=0) + max_off_policy = max((run_of(e).off_policy_steps or 0 for e in effective.episodes), default=0) head = ( f"Step {step} | {format_time(step_time):>7} | Reward {eff.reward.mean():.4f} | " @@ -833,7 +822,7 @@ def log_train_batch(self, batch: TrainBatch, *, step: int, step_time: float) -> lines.append( f"╰─ {env_name:<{name_width}} | Ratio {ratio:.1%} | Reward {env_eff.reward.mean():.4f} | " f"Turns {env_eff.num_turns.mean():.1f} | Branches {env_eff.num_branches.mean():.1f} | " - f"Max Off-Policy {max((run_of(e).off_policy_steps for e in env_eff_pool.episodes), default=0)} | " + f"Max Off-Policy {max((run_of(e).off_policy_steps or 0 for e in env_eff_pool.episodes), default=0)} | " f"Error {pool.metrics.has_error.mean():.1%} | Truncation {env_eff.is_truncated.mean():.1%}" ) get_logger().success("\n\t\t ".join(lines)) diff --git a/src/prime_rl/orchestrator/types.py b/src/prime_rl/orchestrator/types.py index 1fcf07e298..d7be4eb9b4 100644 --- a/src/prime_rl/orchestrator/types.py +++ b/src/prime_rl/orchestrator/types.py @@ -176,10 +176,9 @@ class InflightEpisode: episodes_owed: int """How many episodes this dispatch owes the sink — one, except on the legacy group path.""" client_config: vf.ClientConfig | None = None - off_policy_steps: int = 0 eval_step: int | None = None - def stamp(self, episode: Episode, *, run_id: str, policy_version: int, eval_step: int | None) -> Episode: # noqa: E501 + def stamp(self, episode: Episode, *, run_id: str, policy: vf.PolicySpan | None, eval_step: int | None) -> Episode: """Write the dispatch's facts onto the landed episode, in the places the episode already has for them. The group's values win over this dispatch's when it is still alive, so they are passed in rather than read off ``self``. @@ -192,13 +191,7 @@ def stamp(self, episode: Episode, *, run_id: str, policy_version: int, eval_step episode.group = vf.GroupInfo(id=str(self.group_id)) if self.kind == "eval": assert eval_step is not None, "eval episode missing its step" - episode.run = vf.TrainRunInfo( - id=run_id, - kind=self.kind, - step=eval_step, - policy_version=policy_version, - off_policy_steps=0 if self.kind == "eval" else self.off_policy_steps, - ) + episode.run = vf.TrainRunInfo(id=run_id, kind=self.kind, step=eval_step, policy=policy) return episode diff --git a/tests/unit/orchestrator/test_metrics.py b/tests/unit/orchestrator/test_metrics.py index 2f70d2986c..be0bcf02fd 100644 --- a/tests/unit/orchestrator/test_metrics.py +++ b/tests/unit/orchestrator/test_metrics.py @@ -344,23 +344,27 @@ def test_inflight_episode_stamps_what_lands(): dispatch become facts of an episode, and the run it writes is what tells the rest of the orchestrator which path the episode is on.""" wire = vf.WireEpisode.model_construct(id="e", traces=[]) - inflight = InflightEpisode( - kind="train", env_name="rt", group_id=uuid4(), policy_version=3, episodes_owed=1, off_policy_steps=2 - ) - train = inflight.stamp(wire, run_id="r", policy_version=7, eval_step=None) + inflight = InflightEpisode(kind="train", env_name="rt", group_id=uuid4(), policy_version=3, episodes_owed=1) + span = vf.PolicySpan(start=3, end=4) # an update landed while it was generating + train = inflight.stamp(wire, run_id="r", policy=span, eval_step=None) run = run_of(train) - assert (run.kind, run.policy_version, run.off_policy_steps) == ("train", 7, 2) # group's wins + assert (run.kind, run.policy) == ("train", span) assert run.step is None # the batch window it lands in is not known yet + assert run.off_policy_steps is None # so there is nothing to be behind yet assert train.env.name == "rt" and train.group is not None + run.step = 6 # the window it landed in, which step 6 trains v5 from + assert run.off_policy_steps == 2 and run.policy.drift == 1 + evaluation = replace(inflight, kind="eval").stamp( - vf.WireEpisode.model_construct(id="e", traces=[]), run_id="r", policy_version=7, eval_step=12 + vf.WireEpisode.model_construct(id="e", traces=[]), run_id="r", policy=span, eval_step=12 ) - # An online eval belongs to the same training run — same record, told apart by kind. + # An online eval belongs to the same training run — same record, told apart by kind, and + # measured against the policy in training exactly as an episode trained on is. assert (run_of(evaluation).kind, run_of(evaluation).step) == ("eval", 12) - assert run_of(evaluation).off_policy_steps == 0 # only what trains can go stale + assert run_of(evaluation).off_policy_steps == 8 with pytest.raises(AssertionError): # an eval episode without its step is not representable - replace(inflight, kind="eval").stamp(wire, run_id="r", policy_version=7, eval_step=None) + replace(inflight, kind="eval").stamp(wire, run_id="r", policy=span, eval_step=None) def test_eval_sink_reads_the_epoch_a_stamped_episode_landed_in(): @@ -370,13 +374,11 @@ def test_eval_sink_reads_the_epoch_a_stamped_episode_landed_in(): inflight = InflightEpisode( kind="eval", env_name="rt", group_id=uuid4(), policy_version=3, episodes_owed=1, eval_step=12 ) - landed = inflight.stamp( - vf.WireEpisode.model_construct(id="e", traces=[]), run_id="r", policy_version=3, eval_step=12 - ) + landed = inflight.stamp(vf.WireEpisode.model_construct(id="e", traces=[]), run_id="r", policy=None, eval_step=12) assert eval_step_of(landed) == 12 trained_on = replace(inflight, kind="train").stamp( - vf.WireEpisode.model_construct(id="e2", traces=[]), run_id="r", policy_version=3, eval_step=None + vf.WireEpisode.model_construct(id="e2", traces=[]), run_id="r", policy=None, eval_step=None ) with pytest.raises(AssertionError): # an episode to train on is not placed by an eval epoch eval_step_of(trained_on) From f6e78e95b5f0a3a0844c51b26fc4ecba132698a3 Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Tue, 4 Aug 2026 22:01:38 +0000 Subject: [PATCH 46/58] fix(orchestrator): the eval summary reads the policy off the span Co-Authored-By: Claude Fable 5 --- src/prime_rl/orchestrator/orchestrator.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/prime_rl/orchestrator/orchestrator.py b/src/prime_rl/orchestrator/orchestrator.py index 7f020fa9e8..461669be7b 100644 --- a/src/prime_rl/orchestrator/orchestrator.py +++ b/src/prime_rl/orchestrator/orchestrator.py @@ -843,8 +843,8 @@ async def finalize_eval_batch(self, batch: EvalBatch) -> None: save_episodes, records, get_trace_path(self.config.output_dir, batch.step, "eval", "effective") ) self.monitor.log_eval_samples(batch.rollouts.episodes, env_name=batch.env_name, step=batch.step) - policy_versions = {run_of(e).policy_version for e in batch.rollouts.episodes} - policy_version = min(policy_versions) + policy_versions = {run.policy.start for e in batch.rollouts.episodes if (run := run_of(e)).policy} + policy_version = min(policy_versions, default=0) if len(policy_versions) > 1: get_logger().warning( f"Eval {batch.env_name} step {batch.step} had mixed policy versions: {sorted(policy_versions)}" From e22d373ac3bf11f4475ef03b805e9b07a319627f Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Tue, 4 Aug 2026 22:14:43 +0000 Subject: [PATCH 47/58] refactor(orchestrator): to_record is prime-rl's, not verifiers' Only prime-rl writes episodes carrying training tensors, so only it needs the exclusion; vf's own writer dumps the episode whole. Co-Authored-By: Claude Fable 5 --- deps/verifiers | 2 +- src/prime_rl/orchestrator/orchestrator.py | 7 ++++--- src/prime_rl/orchestrator/types.py | 7 +++++++ 3 files changed, 12 insertions(+), 4 deletions(-) diff --git a/deps/verifiers b/deps/verifiers index 8c318d0a4b..54eef2cf97 160000 --- a/deps/verifiers +++ b/deps/verifiers @@ -1 +1 @@ -Subproject commit 8c318d0a4bc9e8a145267b96701cf02bebb7fc01 +Subproject commit 54eef2cf9715418cfd1054c54d6d3803cd955c25 diff --git a/src/prime_rl/orchestrator/orchestrator.py b/src/prime_rl/orchestrator/orchestrator.py index 461669be7b..4553480977 100644 --- a/src/prime_rl/orchestrator/orchestrator.py +++ b/src/prime_rl/orchestrator/orchestrator.py @@ -63,6 +63,7 @@ TrainBatch, group_id_of, run_of, + to_record, ) from prime_rl.orchestrator.utils import ( get_weight_dir, @@ -537,7 +538,7 @@ async def main_loop(self) -> None: assert step is not None await asyncio.to_thread( save_episodes, - [episode.to_record()], + [to_record(episode)], get_trace_path(self.config.output_dir, step, run.kind, "all"), ) @@ -625,7 +626,7 @@ async def finalize_train_batch(self, batch: TrainBatch) -> None: # at ship time; the full arrival window already streamed into ``all`` on arrival. # to_record drops the per-node training tensors — they're for training, not the rollout # record, and can't round-trip json (raw numpy bytes). - records = [e.to_record() for e in effective.episodes] + records = [to_record(e) for e in effective.episodes] await asyncio.to_thread(save_episodes, records, get_trace_path(config.output_dir, step, "train", "effective")) await self.sender.send(TrainingBatch(examples=batch.samples, step=step)) @@ -838,7 +839,7 @@ async def finalize_eval_batch(self, batch: EvalBatch) -> None: # completion (multiple eval envs share the step file — each epoch appends its cohort # once, and every record carries ``env_name``); the full returned cohort already # streamed into ``all`` on arrival. - records = [e.to_record() for e in batch.rollouts.effective.episodes] + records = [to_record(e) for e in batch.rollouts.effective.episodes] await asyncio.to_thread( save_episodes, records, get_trace_path(self.config.output_dir, batch.step, "eval", "effective") ) diff --git a/src/prime_rl/orchestrator/types.py b/src/prime_rl/orchestrator/types.py index d7be4eb9b4..87667b39d7 100644 --- a/src/prime_rl/orchestrator/types.py +++ b/src/prime_rl/orchestrator/types.py @@ -137,6 +137,13 @@ def rollouts_of(episode: Episode) -> list[TrainRollout]: return cast(list[TrainRollout], episode.traces) +def to_record(episode: Episode) -> dict[str, Any]: + """JSON record without the per-node training tensors — the episode form of + ``Trace.to_record``, and the unit ``traces.jsonl`` stores: one episode per line. The tensors + are the trainer's, not the record's, and raw numpy bytes don't round-trip through json.""" + return episode.model_dump(mode="json", exclude={"traces": {"__all__": EXCLUDE_FIELDS}}) + + def group_id_of(episode: Episode) -> str: """The group an episode was planned in. The dispatcher plans every episode into one, so this is always set by the time anything downstream asks.""" From 16d3fccc0a9f2dcb1ddf24ed81d18a1b753dde5f Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Tue, 4 Aug 2026 22:24:43 +0000 Subject: [PATCH 48/58] refactor(orchestrator)!: route on the run's episode metadata Co-Authored-By: Claude Fable 5 --- deps/verifiers | 2 +- src/prime_rl/orchestrator/eval_sink.py | 10 ++++++---- src/prime_rl/orchestrator/orchestrator.py | 12 +++++++----- src/prime_rl/orchestrator/types.py | 13 ++++++++----- tests/unit/orchestrator/test_metrics.py | 13 +++++++------ 5 files changed, 29 insertions(+), 21 deletions(-) diff --git a/deps/verifiers b/deps/verifiers index 54eef2cf97..260b6f1426 160000 --- a/deps/verifiers +++ b/deps/verifiers @@ -1 +1 @@ -Subproject commit 54eef2cf9715418cfd1054c54d6d3803cd955c25 +Subproject commit 260b6f1426bb510d01d7a2441a5e389c30fbe224 diff --git a/src/prime_rl/orchestrator/eval_sink.py b/src/prime_rl/orchestrator/eval_sink.py index bb9cd98746..f1bc00eeea 100644 --- a/src/prime_rl/orchestrator/eval_sink.py +++ b/src/prime_rl/orchestrator/eval_sink.py @@ -16,6 +16,8 @@ from collections import defaultdict +import verifiers.v1 as vf + from prime_rl.orchestrator.envs import EvalEnvs from prime_rl.orchestrator.metrics import EvalRollouts from prime_rl.orchestrator.types import ( @@ -32,11 +34,11 @@ def eval_step_of(episode: Episode) -> int: """The eval epoch an episode belongs to, off the run the dispatcher recorded when it landed. - An online eval belongs to the training run, so its step is known from the start — unlike an + An online eval belongs to the training run, and its metadata requires the step, unlike an episode to train on, whose step is the window it lands in.""" - run = run_of(episode) - assert run.kind == "eval" and run.step is not None, "not an eval episode" - return run.step + metadata = run_of(episode).metadata + assert isinstance(metadata, vf.EvalMetadata), "not an eval episode" + return metadata.step class EvalSink: diff --git a/src/prime_rl/orchestrator/orchestrator.py b/src/prime_rl/orchestrator/orchestrator.py index 4553480977..a40bd7159f 100644 --- a/src/prime_rl/orchestrator/orchestrator.py +++ b/src/prime_rl/orchestrator/orchestrator.py @@ -39,6 +39,8 @@ from prime_rl.transport.base import TrainingBatchSender from prime_rl.utils.client import InferencePool from prime_rl.utils.monitor.base import Monitor +import verifiers.v1 as vf + import prime_rl._compat # noqa: F401 — patch ring_flash_attn compat before transitive imports from prime_rl.configs.orchestrator import OrchestratorConfig from prime_rl.orchestrator.ckpt import setup_ckpt_manager @@ -532,17 +534,17 @@ async def main_loop(self) -> None: # An eval episode already knows its step (the eval that triggered it); a train one # belongs to the batch window collecting right now, which only this loop knows. run = run_of(episode) - if run.kind == "train": - run.step = self.progress.step - step = run.step + if isinstance(run.metadata, vf.TrainMetadata): + run.metadata.step = self.progress.step + step = run.metadata.step assert step is not None await asyncio.to_thread( save_episodes, [to_record(episode)], - get_trace_path(self.config.output_dir, step, run.kind, "all"), + get_trace_path(self.config.output_dir, step, run.metadata.type, "all"), ) - if run.kind == "eval": + if run.metadata.type == "eval": assert self.eval_sink is not None # eval rollouts only emitted when eval is configured eval_batch = self.eval_sink.add(episode) if eval_batch is not None: diff --git a/src/prime_rl/orchestrator/types.py b/src/prime_rl/orchestrator/types.py index 87667b39d7..b8f0407f47 100644 --- a/src/prime_rl/orchestrator/types.py +++ b/src/prime_rl/orchestrator/types.py @@ -190,15 +190,18 @@ def stamp(self, episode: Episode, *, run_id: str, policy: vf.PolicySpan | None, has for them. The group's values win over this dispatch's when it is still alive, so they are passed in rather than read off ``self``. - The run's ``kind`` says which path the episode is on, which is all the rest of the - orchestrator needs to route it. An eval episode knows its step here; an episode to train on - does not — it belongs to whichever batch window is collecting when it lands, so the main - loop fills that in.""" + The run's metadata says what the episode is to the run, which is all the rest of the + orchestrator needs to route it. An eval knows its step here; an episode to train on does + not — it belongs to whichever batch window is collecting when it lands, so the main loop + fills that in.""" episode.env.name = self.env_name episode.group = vf.GroupInfo(id=str(self.group_id)) if self.kind == "eval": assert eval_step is not None, "eval episode missing its step" - episode.run = vf.TrainRunInfo(id=run_id, kind=self.kind, step=eval_step, policy=policy) + metadata: vf.EpisodeMetadata = vf.EvalMetadata(step=eval_step) + else: + metadata = vf.TrainMetadata() + episode.run = vf.TrainRunInfo(id=run_id, metadata=metadata, policy=policy) return episode diff --git a/tests/unit/orchestrator/test_metrics.py b/tests/unit/orchestrator/test_metrics.py index be0bcf02fd..07122070b3 100644 --- a/tests/unit/orchestrator/test_metrics.py +++ b/tests/unit/orchestrator/test_metrics.py @@ -348,21 +348,22 @@ def test_inflight_episode_stamps_what_lands(): span = vf.PolicySpan(start=3, end=4) # an update landed while it was generating train = inflight.stamp(wire, run_id="r", policy=span, eval_step=None) run = run_of(train) - assert (run.kind, run.policy) == ("train", span) - assert run.step is None # the batch window it lands in is not known yet + assert isinstance(run.metadata, vf.TrainMetadata) and run.policy == span + assert run.metadata.step is None # the batch window it lands in is not known yet assert run.off_policy_steps is None # so there is nothing to be behind yet assert train.env.name == "rt" and train.group is not None - run.step = 6 # the window it landed in, which step 6 trains v5 from + run.metadata.step = 6 # the window it landed in, which step 6 trains v5 from assert run.off_policy_steps == 2 and run.policy.drift == 1 evaluation = replace(inflight, kind="eval").stamp( vf.WireEpisode.model_construct(id="e", traces=[]), run_id="r", policy=span, eval_step=12 ) - # An online eval belongs to the same training run — same record, told apart by kind, and + # An online eval belongs to the same training run — same id, told apart by its metadata, and # measured against the policy in training exactly as an episode trained on is. - assert (run_of(evaluation).kind, run_of(evaluation).step) == ("eval", 12) - assert run_of(evaluation).off_policy_steps == 8 + eval_run = run_of(evaluation) + assert isinstance(eval_run.metadata, vf.EvalMetadata) and eval_run.metadata.step == 12 + assert eval_run.id == run.id and eval_run.off_policy_steps == 8 with pytest.raises(AssertionError): # an eval episode without its step is not representable replace(inflight, kind="eval").stamp(wire, run_id="r", policy=span, eval_step=None) From cdc49dc50379c1c84e6658d4316823f94dca6c1f Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Tue, 4 Aug 2026 22:28:32 +0000 Subject: [PATCH 49/58] fix(orchestrator): restore the EXCLUDE_FIELDS import to_record needs Co-Authored-By: Claude Fable 5 --- src/prime_rl/orchestrator/types.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/prime_rl/orchestrator/types.py b/src/prime_rl/orchestrator/types.py index b8f0407f47..f077a390b7 100644 --- a/src/prime_rl/orchestrator/types.py +++ b/src/prime_rl/orchestrator/types.py @@ -10,6 +10,7 @@ import verifiers.v1 as vf from pydantic import ConfigDict, Field from verifiers.v1.task import DataT +from verifiers.v1.trace import EXCLUDE_FIELDS from prime_rl.transport import TrainingSample From 44b8f5bd1dbd3a6c13f9845699012a55ed3341cd Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Tue, 4 Aug 2026 22:59:58 +0000 Subject: [PATCH 50/58] chore: an eval's off-policy reading is its drift Co-Authored-By: Claude Fable 5 --- deps/verifiers | 2 +- tests/unit/orchestrator/test_metrics.py | 7 ++++--- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/deps/verifiers b/deps/verifiers index 260b6f1426..9576a5a9e3 160000 --- a/deps/verifiers +++ b/deps/verifiers @@ -1 +1 @@ -Subproject commit 260b6f1426bb510d01d7a2441a5e389c30fbe224 +Subproject commit 9576a5a9e315e28e5154c1fdca1dd577df971380 diff --git a/tests/unit/orchestrator/test_metrics.py b/tests/unit/orchestrator/test_metrics.py index 07122070b3..c49ef7d409 100644 --- a/tests/unit/orchestrator/test_metrics.py +++ b/tests/unit/orchestrator/test_metrics.py @@ -359,11 +359,12 @@ def test_inflight_episode_stamps_what_lands(): evaluation = replace(inflight, kind="eval").stamp( vf.WireEpisode.model_construct(id="e", traces=[]), run_id="r", policy=span, eval_step=12 ) - # An online eval belongs to the same training run — same id, told apart by its metadata, and - # measured against the policy in training exactly as an episode trained on is. + # An online eval belongs to the same training run — same id, told apart by its metadata. Its + # off-policy reading is what drifted under it, not its distance from a step it was dispatched + # at: nothing trains on an eval, and a slow one would otherwise look on-policy forever. eval_run = run_of(evaluation) assert isinstance(eval_run.metadata, vf.EvalMetadata) and eval_run.metadata.step == 12 - assert eval_run.id == run.id and eval_run.off_policy_steps == 8 + assert eval_run.id == run.id and eval_run.off_policy_steps == span.drift == 1 with pytest.raises(AssertionError): # an eval episode without its step is not representable replace(inflight, kind="eval").stamp(wire, run_id="r", policy=span, eval_step=None) From 8639c903b4d57571c19c2adcf934092ea73a4c4a Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Tue, 4 Aug 2026 23:07:07 +0000 Subject: [PATCH 51/58] chore: bump verifiers to the run-types move Co-Authored-By: Claude Fable 5 --- deps/verifiers | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deps/verifiers b/deps/verifiers index 9576a5a9e3..eddc781f22 160000 --- a/deps/verifiers +++ b/deps/verifiers @@ -1 +1 @@ -Subproject commit 9576a5a9e315e28e5154c1fdca1dd577df971380 +Subproject commit eddc781f22174ebefb749d0babf2367fdd034af3 From b393ae72aee466891ac0a5b2d9fa8c0ab6685063 Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Tue, 4 Aug 2026 23:52:43 +0000 Subject: [PATCH 52/58] refactor(orchestrator)!: degeneracy is measured unconditionally Whether a trace is gibberish or stuck in a loop is a fact about it, not something a run opts into computing: measuring it behind config made the rate depend on what was configured. Both run on every trace with fixed thresholds, and the only choice left is whether to act on one. Co-Authored-By: Claude Fable 5 --- docs/algorithms.md | 40 ++- examples/advanced/glm-5.2/swe.toml | 3 +- .../src/prime_rl/configs/orchestrator.py | 38 +-- src/prime_rl/orchestrator/degeneracy.py | 84 ++++++ src/prime_rl/orchestrator/detectors.py | 129 --------- src/prime_rl/orchestrator/metrics.py | 6 +- src/prime_rl/orchestrator/orchestrator.py | 10 +- src/prime_rl/orchestrator/train_sink.py | 18 +- src/prime_rl/orchestrator/types.py | 6 +- tests/unit/orchestrator/test_degeneracy.py | 171 +++++++++++ tests/unit/orchestrator/test_detectors.py | 269 ------------------ tests/unit/orchestrator/test_metrics.py | 12 +- 12 files changed, 300 insertions(+), 486 deletions(-) create mode 100644 src/prime_rl/orchestrator/degeneracy.py delete mode 100644 src/prime_rl/orchestrator/detectors.py create mode 100644 tests/unit/orchestrator/test_degeneracy.py delete mode 100644 tests/unit/orchestrator/test_detectors.py diff --git a/docs/algorithms.md b/docs/algorithms.md index 64d38699ad..a324aad635 100644 --- a/docs/algorithms.md +++ b/docs/algorithms.md @@ -1,6 +1,6 @@ # Algorithms -This page covers the math and the configurable algorithmic components: the algorithm abstraction and its algorithms, how off-policy training works, the loss components and advantage functions, how to plug in your own, the degeneracy detectors and drop policy applied between rollout and training, and how multi-turn rollouts get merged into training samples. +This page covers the math and the configurable algorithmic components: the algorithm abstraction and its algorithms, how off-policy training works, the loss components and advantage functions, how to plug in your own, the degeneracy measurements and drop policy applied between rollout and training, and how multi-turn rollouts get merged into training samples. ## Table of Contents @@ -21,7 +21,7 @@ This page covers the math and the configurable algorithmic components: the algor - [Self-Play Advantage (RAE)](#self-play-advantage-rae) - [Authoring an Algorithm](#authoring-an-algorithm) - [Reference Scoring](#reference-scoring) -- [Detectors and the drop policy](#detectors-and-the-drop-policy) +- [Degeneracy and the drop policy](#degeneracy-and-the-drop-policy) - [Multi-Turn Trajectories](#multi-turn-trajectories) - [Extension Property](#extension-property) - [Best-Effort Interleaving](#best-effort-interleaving) @@ -458,48 +458,40 @@ demo_key = "demonstration" Scoring runs at arrival, *before* the drop policy, so a rollout that is later dropped still cost its reference compute — accepted for the simpler one-rollout-at-a-time shape (zero-advantage dropping never applies to opd/opsd anyway, since neither assigns an advantage). -## Detectors and the drop policy +## Degeneracy and the drop policy Two separate things: **measuring** what a rollout looks like, and **deciding** whether to train on it. -Detectors measure. Each one asks a single question of a trace's tokens, runs on every trace as soon -as it is tokenized, and reports its rate per agent alongside reward and truncation — whether or not -anything acts on it. Every detector measures every trace, so one detection never hides another. +Every trace is measured, unconditionally, as soon as it is tokenized, and each measurement is +reported per agent alongside reward and truncation. Nothing configures that — a rate is only a rate +if nothing decided in advance which traces to look at, and every measurement runs on every trace, so +one never hides another. -| Detector | Measures | +| Measurement | What it finds | |---|---| -| `gibberish` | rare tokens (high BPE id) generated at high entropy — degenerate output | +| `gibberish` | a rare token (high BPE id) generated at high entropy — degenerate output | | `repetition` | a long stretch of very-high-confidence tokens — a repetition loop | -```toml -[orchestrator.detectors] -drop = ["gibberish"] # measured always; this says which detections also drop - -[orchestrator.detectors.repetition] -window = 2000 -``` - -Set a detector to `false` to stop measuring it (`[orchestrator.detectors] repetition = false`). - The drop policy decides. It runs once, when a finalized group's credit is assigned and the rollouts -would enter the batch buffer, and it has two inputs: the detections the run listed in -`detectors.drop`, and **zero credit**, which drops on its own: +would enter the batch buffer, and it has two inputs — the measurements the run chose to act on, and +**zero credit**, which drops on its own: ```toml [orchestrator] +drop_degenerate = ["gibberish"] # measuring is unconditional; acting on it is opt-in drop_zero_advantage = true # the default ``` A scored rollout whose every token is worth nothing — a GRPO group where all rollouts earned the -same reward — produces no gradient, so training on it is a wasted forward pass. It is not a plugin -because it is not a property of the generation: it is only knowable after the group is scored, and -it is what `TrainRollout.is_trainable` already means. +same reward — produces no gradient, so training on it is a wasted forward pass. It is not one of the +measurements because it is not a property of the generation: it is only knowable after the group is +scored, and it is what `TrainRollout.is_trainable` already means. A rollout that was **never** scored is not zero-credit. `opd` / `opsd` assign no advantages at all and train through reference KL, so they are never dropped by this rule. Dropped rollouts still appear in the metrics window and the `all` trace file — they just don't ship. -`{scope}/{subset}//detected//mean` is each detector's rate, and +`{scope}/{subset}//detected//mean` is each measurement's rate, and `{scope}/{subset}//is_filtered/mean` is the share the policy dropped. ## Multi-Turn Trajectories diff --git a/examples/advanced/glm-5.2/swe.toml b/examples/advanced/glm-5.2/swe.toml index a9133b475b..6ff2d86f6b 100644 --- a/examples/advanced/glm-5.2/swe.toml +++ b/examples/advanced/glm-5.2/swe.toml @@ -60,6 +60,7 @@ weight_decay = 0.1 [orchestrator] +drop_degenerate = ["gibberish"] batch_size = 4096 group_size = 16 oversampling_factor = 3 @@ -84,8 +85,6 @@ id = "bash" type = "prime" labels = ["glm5-pd-disag", "swe-bench-verified"] -[orchestrator.detectors] -drop = ["gibberish"] [inference] enable_expert_parallel = true diff --git a/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py b/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py index 0b625999e0..0b816014fe 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py @@ -357,38 +357,6 @@ class CheckpointConfig(BaseConfig): # Flags rare tokens generated at high entropy (Section 5.2, https://arxiv.org/abs/2510.02387). -class GibberishDetectorConfig(BaseConfig): - """Flags rare tokens generated at high entropy (Section 5.2, https://arxiv.org/abs/2510.02387).""" - - token_id_threshold: int = 100_000 - """Token IDs above this are candidates for gibberish. BPE tokens are sorted by merge order.""" - - logprob_offset: float = 2.0 - """Offset from uniform-distribution logprob. Threshold = ``-log(vocab_size) - logprob_offset``.""" - - -class RepetitionDetectorConfig(BaseConfig): - """Flags rollouts stuck in a repetition loop: high-confidence tokens for an extended stretch - (Section 3.2, https://arxiv.org/abs/2506.13585).""" - - window: int = Field(3_000, ge=1) - """Consecutive high-probability steps required to flag the rollout.""" - - prob_threshold: float = Field(0.99, gt=0, le=1) - """Tokens sampled with probability above this are considered repetitive. Consecutive such tokens count toward the window.""" - - -class DetectorsConfig(BaseConfig): - """Degeneracy detectors. Each one measures every trace and reports its rate per agent; set - ``drop`` to also keep what it flags out of the training batch.""" - - gibberish: GibberishDetectorConfig | None = GibberishDetectorConfig() - repetition: RepetitionDetectorConfig | None = RepetitionDetectorConfig() - - drop: list[Literal["gibberish", "repetition"]] = [] - """Which detections keep a rollout out of the batch. Measuring is always on; dropping is not.""" - - class FileSystemWeightBroadcastConfig(BaseConfig): type: Literal["filesystem"] = "filesystem" @@ -462,9 +430,9 @@ class OrchestratorConfig(BaseConfig): eval: EvalConfig | None = None """Evaluation configuration.""" - detectors: DetectorsConfig = DetectorsConfig() - """Degeneracy detectors, measured on every trace and reported per agent. ``detectors.drop`` - says which detections also keep a rollout out of the training batch.""" + drop_degenerate: list[Literal["gibberish", "repetition"]] = [] + """Which degeneracy measurements also keep a rollout out of the training batch. Every trace is + measured for all of them and reports them per agent; acting on one is opt-in.""" drop_zero_advantage: bool = True """Keep scored rollouts whose credit is all zero out of the batch — a GRPO group that earned a diff --git a/src/prime_rl/orchestrator/degeneracy.py b/src/prime_rl/orchestrator/degeneracy.py new file mode 100644 index 0000000000..62102f9847 --- /dev/null +++ b/src/prime_rl/orchestrator/degeneracy.py @@ -0,0 +1,84 @@ +"""Degeneracy measurements: what a trace's tokens say about how it was generated. + +Each one asks a single question — is it gibberish, is it stuck in a repetition loop — of every +trace, unconditionally, as soon as it is tokenized. They are metrics, reported per agent beside +reward and truncation, so they say something whether or not a run acts on them. + +Acting on one is a separate decision, made once when the batch is assembled +(``prime_rl.orchestrator.train_sink``). Keeping the two apart is what lets every trace be measured +for all of them: a policy that stopped at the first hit would leave the rest unmeasured. +""" + +from __future__ import annotations + +import math +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from prime_rl.orchestrator.types import TrainRollout + +TOKEN_ID_THRESHOLD = 100_000 +"""Token IDs above this are candidates for gibberish. BPE tokens are sorted by merge order.""" + +LOGPROB_OFFSET = 2.0 +"""Offset from the uniform-distribution logprob: `-log(vocab_size) - LOGPROB_OFFSET`.""" + +REPETITION_WINDOW = 3_000 +"""Consecutive high-probability sampled tokens that count as a repetition loop.""" + +REPETITION_PROB = 0.99 +"""Tokens sampled above this probability count toward the window.""" + + +def is_gibberish(rollout: TrainRollout, vocab_size: int) -> bool: + """Whether the trace generated a rare token at high entropy — a rare BPE id sampled as if + the model had no idea (Section 5.2, https://arxiv.org/abs/2510.02387).""" + threshold = -math.log(vocab_size) - LOGPROB_OFFSET + for branch in rollout.branches: + # branch.{token_ids,logprobs,sampled_mask} are flat and mutually aligned; the raw node + # arrays are not (node.logprobs covers only the sampled suffix, not the generation-prompt + # scaffold that token_ids/mask also span). + for token_id, logprob, sampled in zip(branch.token_ids, branch.logprobs, branch.sampled_mask): + if sampled and token_id > TOKEN_ID_THRESHOLD and logprob < threshold: + return True + return False + + +def is_repetitive(rollout: TrainRollout) -> bool: + """Whether the trace held very high confidence for a long stretch — the signature of a + repetition loop (Section 3.2, https://arxiv.org/abs/2506.13585).""" + threshold = math.log(REPETITION_PROB) + for branch in rollout.branches: + # Aligned branch streams (see `is_gibberish`), and reset the streak per branch: flat + # rollout.nodes interleaves distinct root->leaf paths (compaction/subagents), so a per-node + # walk would run a streak across a branch boundary. + consecutive = 0 + for logprob, sampled in zip(branch.logprobs, branch.sampled_mask): + if not sampled: + continue + consecutive = consecutive + 1 if logprob > threshold else 0 + if consecutive >= REPETITION_WINDOW: + return True + return False + + +def measure(rollout: TrainRollout, vocab_size: int) -> None: + """Record every measurement on the trace. All of them, every time — a rate is only a rate if + nothing decided in advance which traces to look at.""" + rollout.degeneracy = { + "gibberish": is_gibberish(rollout, vocab_size), + "repetition": is_repetitive(rollout), + } + + +def drop_reasons(rollout: TrainRollout, *, drop: list[str], drop_zero_advantage: bool) -> list[str]: + """Why this rollout should not be trained on, if anything. + + A measurement only drops when the run asked it to — measuring is unconditional, acting is not. + Zero credit drops on its own: a scored rollout whose every token is worth nothing produces no + gradient, so the forward pass is wasted. A rollout that was never scored is *not* zero-credit — + opd/opsd train through reference KL and assign no advantages at all, so they must survive.""" + reasons = [name for name in drop if rollout.degeneracy.get(name)] + if drop_zero_advantage and rollout.advantages is not None and not rollout.is_trainable: + reasons.append("zero_advantage") + return reasons diff --git a/src/prime_rl/orchestrator/detectors.py b/src/prime_rl/orchestrator/detectors.py deleted file mode 100644 index 4a0f56d797..0000000000 --- a/src/prime_rl/orchestrator/detectors.py +++ /dev/null @@ -1,129 +0,0 @@ -"""Degeneracy detectors: per-trace measurements of pathological generation. - -A detector answers one question about one trace's tokens — is it gibberish, is it stuck in a -repetition loop — and nothing else. Every configured detector runs on every trace as soon as it is -tokenized, so its rate is a trace metric like reward or truncation, reported per agent whether or -not anything acts on it. - -What to *do* about a detection is a separate decision, made once when the batch is assembled -(``prime_rl.orchestrator.train_sink``). Keeping the two apart is what lets every detector measure -every trace: a policy that stops at the first hit would leave the rest unmeasured. -""" - -from __future__ import annotations - -import math -from dataclasses import dataclass -from typing import TYPE_CHECKING, Protocol - -from prime_rl.configs.orchestrator import DetectorsConfig -from prime_rl.utils.logger import get_logger - -if TYPE_CHECKING: - from prime_rl.orchestrator.types import TrainRollout - - -class Detector(Protocol): - name: str - - def detect(self, rollout: TrainRollout) -> bool: ... - - -@dataclass -class GibberishDetector: - """Rare tokens generated at high entropy. - - A token counts when both: - - id(token) > token_id_threshold (rare BPE token) - - logprob(token) < -log(vocab_size) - logprob_offset (high entropy) - - References: - Section 5.2, https://arxiv.org/abs/2510.02387 - """ - - name: str - token_id_threshold: int - logprob_threshold: float - - def detect(self, rollout: TrainRollout) -> bool: - for branch in rollout.branches: - # branch.{token_ids,logprobs,sampled_mask} are flat and mutually aligned; the raw - # node arrays are not (node.logprobs covers only the sampled suffix, not the - # generation-prompt scaffold that token_ids/mask also span). - for token_id, logprob, sampled in zip(branch.token_ids, branch.logprobs, branch.sampled_mask): - if sampled and token_id > self.token_id_threshold and logprob < self.logprob_threshold: - return True - return False - - -@dataclass -class RepetitionDetector: - """A repetition loop: a long stretch of very-high-confidence tokens. - - Counts consecutive tokens with logprob > log(prob_threshold); a streak reaching ``window`` - is a detection. - - References: - Section 3.2, https://arxiv.org/abs/2506.13585 - """ - - name: str - window: int - logprob_threshold: float - - def detect(self, rollout: TrainRollout) -> bool: - for branch in rollout.branches: - # Aligned branch streams (see GibberishDetector), and reset the streak per branch: - # flat rollout.nodes interleaves distinct root->leaf paths (compaction/subagents), - # so a per-node walk would run a streak across a branch boundary. - consecutive = 0 - for logprob, sampled in zip(branch.logprobs, branch.sampled_mask): - if not sampled: - continue - consecutive = consecutive + 1 if logprob > self.logprob_threshold else 0 - if consecutive >= self.window: - return True - return False - - -def setup_detectors(config: DetectorsConfig, vocab_size: int) -> list[Detector]: - detectors: list[Detector] = [] - if config.gibberish is not None: - detectors.append( - GibberishDetector( - name="gibberish", - token_id_threshold=config.gibberish.token_id_threshold, - logprob_threshold=-math.log(vocab_size) - config.gibberish.logprob_offset, - ) - ) - if config.repetition is not None: - detectors.append( - RepetitionDetector( - name="repetition", - window=config.repetition.window, - logprob_threshold=math.log(config.repetition.prob_threshold), - ) - ) - if detectors: - get_logger().info(f"Measuring {len(detectors)} degeneracy detector(s): {', '.join(d.name for d in detectors)}") - if config.drop: - get_logger().info(f"Dropping detected rollouts: {', '.join(sorted(config.drop))}") - return detectors - - -def detect(detectors: list[Detector], rollout: TrainRollout) -> None: - """Measure every detector on one trace, writing the verdicts to ``rollout.detections``.""" - rollout.detections = {d.name: d.detect(rollout) for d in detectors} - - -def drop_reasons(rollout: TrainRollout, *, drop_detections: list[str], drop_zero_advantage: bool) -> list[str]: - """Why this rollout should not be trained on, if anything. - - A detection only drops when the run asked it to — measuring is always on, acting is opt-in. - Zero credit drops on its own: a scored rollout whose every token is worth nothing produces no - gradient, so the forward pass is wasted. A rollout that was never scored is *not* zero-credit — - opd/opsd train through reference KL and assign no advantages at all, so they must survive.""" - reasons = [name for name in drop_detections if rollout.detections.get(name)] - if drop_zero_advantage and rollout.advantages is not None and not rollout.is_trainable: - reasons.append("zero_advantage") - return reasons diff --git a/src/prime_rl/orchestrator/metrics.py b/src/prime_rl/orchestrator/metrics.py index bcbb61865b..b71c3e5603 100644 --- a/src/prime_rl/orchestrator/metrics.py +++ b/src/prime_rl/orchestrator/metrics.py @@ -344,16 +344,16 @@ def reward(self) -> Stat: def to_wandb(self, *, prefix: str, subset: Subset) -> dict[str, float]: out = super().to_wandb(prefix=prefix, subset=subset) - # Detections are measured on every trace; the drop verdict is only ever reached for + # Degeneracy is measured on every trace; the drop verdict is only ever reached for # trainable survivors (an untrainable seat is 0.0 throughout). Both read per agent. for agent, traces in self.by_agent().items(): p = f"{prefix}/{subset}/{agent}" rollouts = traces.rollouts out[f"{p}/is_trainable/mean"] = sum(float(r.is_trainable) for r in rollouts) / len(rollouts) out[f"{p}/is_filtered/mean"] = sum(float(r.is_filtered) for r in rollouts) / len(rollouts) - names = sorted({name for r in rollouts for name in r.detections}) + names = sorted({name for r in rollouts for name in r.degeneracy}) out |= { - f"{p}/detected/{name}/mean": sum(1 for r in rollouts if r.detections.get(name)) / len(rollouts) + f"{p}/detected/{name}/mean": sum(1 for r in rollouts if r.degeneracy.get(name)) / len(rollouts) for name in names } return out diff --git a/src/prime_rl/orchestrator/orchestrator.py b/src/prime_rl/orchestrator/orchestrator.py index a40bd7159f..c50a2e2322 100644 --- a/src/prime_rl/orchestrator/orchestrator.py +++ b/src/prime_rl/orchestrator/orchestrator.py @@ -5,7 +5,7 @@ - ``RolloutDispatcher`` schedules rollouts; emits ``Episode`` (train/eval discriminated by ``kind``) on its queue — the env's own episode, or a traceless one saying why there is none. -- ``TrainSink`` ingests train rollouts (tokenize → detect → advantages → drop policy) +- ``TrainSink`` ingests train rollouts (tokenize → measure → advantages → drop policy) and returns a ``TrainBatch`` when the threshold is met. - ``EvalSink`` ingests eval rollouts and returns an ``EvalBatch`` (the full returned cohort) on epoch completion. @@ -44,7 +44,6 @@ import prime_rl._compat # noqa: F401 — patch ring_flash_attn compat before transitive imports from prime_rl.configs.orchestrator import OrchestratorConfig from prime_rl.orchestrator.ckpt import setup_ckpt_manager -from prime_rl.orchestrator.detectors import setup_detectors from prime_rl.orchestrator.dispatcher import DispatcherMetrics, DispatcherMode, RolloutDispatcher from prime_rl.orchestrator.envs import EvalEnvs, TrainEnvs from prime_rl.orchestrator.eval_sink import EvalSink @@ -258,7 +257,6 @@ async def setup(self) -> None: self.usage_reporter = UsageReporter() # Filters apply to train rollouts only - detectors = setup_detectors(config.detectors, vocab_size=self.tokenizer.vocab_size) get_logger().info("Loading training environments") self.train_envs = TrainEnvs( @@ -420,8 +418,8 @@ async def setup(self) -> None: mm_token_type_ids_mapping=self.mm_token_type_ids_mapping, batch_size=config.batch_size, token_batch_size=config.token_batch_size, - detectors=detectors, - drop_detections=config.detectors.drop, + vocab_size=self.tokenizer.vocab_size, + drop_degenerate=config.drop_degenerate, drop_zero_advantage=config.drop_zero_advantage, ) self.eval_sink = EvalSink(eval_envs=self.eval_envs) if self.eval_envs is not None else None @@ -592,7 +590,7 @@ async def finalize_train_batch(self, batch: TrainBatch) -> None: if self.consecutive_empty_batches >= MAX_CONSECUTIVE_EMPTY_BATCHES: raise RuntimeError( f"{self.consecutive_empty_batches} consecutive empty train batches — " - "check the drop policy (detectors.drop / drop_zero_advantage) or task difficulty." + "check the drop policy (drop_degenerate / drop_zero_advantage) or task difficulty." ) return self.consecutive_empty_batches = 0 diff --git a/src/prime_rl/orchestrator/train_sink.py b/src/prime_rl/orchestrator/train_sink.py index c64172bc30..812dff285b 100644 --- a/src/prime_rl/orchestrator/train_sink.py +++ b/src/prime_rl/orchestrator/train_sink.py @@ -1,8 +1,8 @@ """TrainSink: three-level rollout sink for the training side. 1. ``process_rollout`` — eager per-rollout tokenization (overlaps with - dispatcher producing more rollouts), the degeneracy detectors, then the env - algorithm's ``finalize_rollout`` (rollout-local scoring + any reference + dispatcher producing more rollouts), the degeneracy measurements, then the + env algorithm's ``finalize_rollout`` (rollout-local scoring + any reference I/O). Errored and untrainable rollouts skip this. 2. ``process_group`` — filters errored rollouts, hands the episodes narrowed to their trainable survivors to the env algorithm's ``finalize_group`` @@ -23,7 +23,7 @@ from collections import defaultdict from prime_rl.configs.orchestrator import OrchestratorConfig -from prime_rl.orchestrator.detectors import Detector, detect, drop_reasons +from prime_rl.orchestrator.degeneracy import drop_reasons, measure from prime_rl.orchestrator.envs import TrainEnvs from prime_rl.orchestrator.metrics import TrainRollouts from prime_rl.orchestrator.trajectories import trace_to_samples @@ -66,8 +66,8 @@ def __init__( mm_token_type_ids_mapping: dict[int, int] | None, batch_size: int | None, token_batch_size: int | None, - detectors: list[Detector], - drop_detections: list[str], + vocab_size: int, + drop_degenerate: list[str], drop_zero_advantage: bool, ) -> None: assert (batch_size is None) != (token_batch_size is None), ( @@ -79,8 +79,8 @@ def __init__( self.mm_token_type_ids_mapping = mm_token_type_ids_mapping self.batch_size = batch_size self.token_batch_size = token_batch_size - self.detectors = detectors - self.drop_detections = drop_detections + self.vocab_size = vocab_size + self.drop_degenerate = drop_degenerate self.drop_zero_advantage = drop_zero_advantage # Observation window for the next shipped batch: rollouts of groups @@ -176,7 +176,7 @@ async def process_rollout(self, rollout: TrainRollout) -> None: mm_token_type_ids_mapping=self.mm_token_type_ids_mapping, ) rollout.samples = samples or [] - detect(self.detectors, rollout) + measure(rollout, self.vocab_size) # Arrival phase: rollout-local scoring (raw reward, echo observation # weighting, opd/opsd reference logprobs) runs as soon as the rollout is # tokenized — before its group is complete. @@ -242,7 +242,7 @@ async def process_group(self, group_id: str) -> None: for r in survivors: self.pre_filter_seen += 1 reasons = drop_reasons( - r, drop_detections=self.drop_detections, drop_zero_advantage=self.drop_zero_advantage + r, drop=self.drop_degenerate, drop_zero_advantage=self.drop_zero_advantage ) r.is_filtered = bool(reasons) if reasons: diff --git a/src/prime_rl/orchestrator/types.py b/src/prime_rl/orchestrator/types.py index f077a390b7..a1b865ca3e 100644 --- a/src/prime_rl/orchestrator/types.py +++ b/src/prime_rl/orchestrator/types.py @@ -67,7 +67,7 @@ class GroupState: class TrainRollout(vf.Trace[DataT], Generic[DataT]): """A rollout on the training path — the one place prime-rl extends a verifiers type, because trainer-bound state has nowhere else to live: the samples built from its branches, what the - degeneracy detectors measured on it, and whether the drop policy kept it. All of it is + degeneracy measurements found on it, and whether the drop policy kept it. All of it is ``exclude=True``, so dumping one yields a plain trace on the wire. ``env_name`` rides along because a sample is routed by it (the trainer's per-env loss config) @@ -77,8 +77,8 @@ class TrainRollout(vf.Trace[DataT], Generic[DataT]): env_name: str = Field(default="", exclude=True) samples: list[TrainingSample] = Field(default_factory=list, exclude=True) - detections: dict[str, bool] = Field(default_factory=dict, exclude=True) - """What each degeneracy detector measured on this trace — a measurement, not a verdict.""" + degeneracy: dict[str, bool] = Field(default_factory=dict, exclude=True) + """What each degeneracy measurement found on this trace — a measurement, not a verdict.""" is_filtered: bool = Field(default=False, exclude=True) """The sink's verdict: this rollout is not trained on. Kept for the metrics window, which reports what came back as well as what shipped.""" diff --git a/tests/unit/orchestrator/test_degeneracy.py b/tests/unit/orchestrator/test_degeneracy.py new file mode 100644 index 0000000000..975ec33789 --- /dev/null +++ b/tests/unit/orchestrator/test_degeneracy.py @@ -0,0 +1,171 @@ +import math + +import verifiers.v1 as vf + +from prime_rl.orchestrator.degeneracy import ( + REPETITION_PROB, + REPETITION_WINDOW, + TOKEN_ID_THRESHOLD, + drop_reasons, + is_gibberish, + is_repetitive, + measure, +) +from prime_rl.orchestrator.types import TrainRollout + +VOCAB_SIZE = 128_000 +GIBBERISH_LOGPROB = -math.log(VOCAB_SIZE) - 2.0 - 1.0 # comfortably under the entropy threshold +RARE_TOKEN = TOKEN_ID_THRESHOLD + 1 +REPEAT_LOGPROB = math.log(REPETITION_PROB) + 0.001 # just above the confidence threshold + + +def _assistant_node(token_ids: list[int], logprobs: list[float]) -> vf.MessageNode: + """An assistant node whose tokens are all model-sampled (the measurements read each node's + masked-True tokens + logprobs).""" + return vf.MessageNode( + message=vf.AssistantMessage(content="x"), + token_ids=token_ids, + mask=[True] * len(token_ids), + logprobs=logprobs, + ) + + +def _make_rollout(nodes: list[vf.MessageNode]) -> TrainRollout: + rollout = TrainRollout[vf.TaskData]( + task=vf.TraceTask(type="Task", data=vf.TaskData(idx=0, prompt="")), + agent=vf.AgentInfo(config=vf.AgentConfig()), + nodes=nodes, + rewards={"reward": vf.Reward(score=1.0)}, + ) + rollout.env_name = "test" + return rollout + + +def _sampled(token_ids: list[int], logprobs: list[float], *, multi_step: bool = False) -> TrainRollout: + """A rollout carrying these completion tokens, optionally split across two model turns.""" + if not multi_step: + return _make_rollout([_assistant_node(token_ids, logprobs)]) + mid = len(token_ids) // 2 + return _make_rollout( + [ + _assistant_node(token_ids[:mid], logprobs[:mid]), + _assistant_node(token_ids[mid:], logprobs[mid:]), + ] + ) + + +# --- gibberish --- + + +def test_gibberish_detects_rare_low_prob_token(): + assert is_gibberish(_sampled([50, RARE_TOKEN, 80], [-1.0, GIBBERISH_LOGPROB, -0.5]), VOCAB_SIZE) + + +def test_gibberish_ignores_normal_tokens(): + assert not is_gibberish(_sampled([10, 200, 5000], [-1.0, -2.0, -3.0]), VOCAB_SIZE) + + +def test_gibberish_ignores_high_prob_rare_token(): + assert not is_gibberish(_sampled([RARE_TOKEN], [-0.5]), VOCAB_SIZE) + + +def test_gibberish_works_across_trajectory_steps(): + rollout = _sampled([50, 60, RARE_TOKEN, 80], [-1.0, -0.5, GIBBERISH_LOGPROB, -0.5], multi_step=True) + assert is_gibberish(rollout, VOCAB_SIZE) + + +def test_gibberish_aligns_logprobs_under_generation_prompt_scaffold(): + """Regression: the assistant node carries a generation-prompt scaffold (mask=False) and + suffix-only logprobs, and the gibberish token is the LAST completion token. A per-node + ``zip(token_ids, logprobs, mask)`` truncates at len(logprobs) and never examines it; reading + the aligned branch streams finds it.""" + node = vf.MessageNode( + message=vf.AssistantMessage(content="x"), + token_ids=[1, 1, 50, 80, RARE_TOKEN], + mask=[False, False, True, True, True], + logprobs=[-1.0, -0.5, GIBBERISH_LOGPROB], # the sampled suffix only, as vLLM returns it + ) + assert is_gibberish(_make_rollout([node]), VOCAB_SIZE) + + +# --- repetition --- + + +def test_repetition_triggers_at_the_window(): + n = REPETITION_WINDOW + assert is_repetitive(_sampled(list(range(n)), [REPEAT_LOGPROB] * n)) + + +def test_repetition_no_trigger_below_the_window(): + n = REPETITION_WINDOW - 1 + assert not is_repetitive(_sampled(list(range(n)), [REPEAT_LOGPROB] * n)) + + +def test_repetition_resets_on_a_low_probability_token(): + """The streak has to be consecutive: one unconfident token in the middle breaks it, so nearly + twice the window's worth of confident tokens either side of it is not a loop.""" + half = [REPEAT_LOGPROB] * (REPETITION_WINDOW - 1) + logprobs = [*half, -2.0, *half] + assert not is_repetitive(_sampled(list(range(len(logprobs))), logprobs)) + + +def test_repetition_does_not_run_a_streak_across_branches(): + """Each branch is measured on its own — a flat walk over ``nodes`` would join two turns' + streaks into one that never happened.""" + half = [REPEAT_LOGPROB] * (REPETITION_WINDOW - 1) + logprobs = [*half, *half] + assert not is_repetitive(_sampled(list(range(len(logprobs))), logprobs, multi_step=True)) + + +# --- measuring is unconditional --- + + +def test_measure_records_every_measurement(): + """Every trace is measured for all of them, so a rate is a rate: one hit does not shadow the + others, and nothing decided in advance which traces to look at.""" + rollout = _sampled([RARE_TOKEN] * 5, [GIBBERISH_LOGPROB] * 5) + measure(rollout, VOCAB_SIZE) + assert rollout.degeneracy == {"gibberish": True, "repetition": False} + + +def test_measure_handles_a_trace_with_no_sampled_tokens(): + rollout = _make_rollout([]) + measure(rollout, VOCAB_SIZE) + assert rollout.degeneracy == {"gibberish": False, "repetition": False} + + +# --- the drop policy --- + + +def _drop(rollout, drop=(), zero_advantage=True): + return drop_reasons(rollout, drop=list(drop), drop_zero_advantage=zero_advantage) + + +def test_a_measurement_only_drops_when_asked(): + rollout = _sampled([1], [-1.0]) + rollout.degeneracy = {"gibberish": True} + rollout.assign_advantages(0.5) + assert _drop(rollout) == [] # measured, but nothing asked for it to drop + assert _drop(rollout, drop=["gibberish"]) == ["gibberish"] + + +def test_zero_credit_drops_by_default(): + rollout = _sampled([1], [-1.0]) + rollout.assign_advantages(0.0) + assert _drop(rollout) == ["zero_advantage"] + assert _drop(rollout, zero_advantage=False) == [] + + +def test_unscored_rollout_is_not_zero_credit(): + """opd/opsd assign no credit at all and train through reference KL — dropping them as + zero-advantage would ship an empty batch for every distillation run.""" + rollout = _sampled([1], [-1.0]) + assert rollout.advantages is None + assert _drop(rollout) == [] + + +def test_reasons_accumulate(): + rollout = _sampled([1], [-1.0]) + rollout.degeneracy = {"gibberish": True, "repetition": True} + rollout.assign_advantages(0.0) + assert _drop(rollout, drop=["gibberish", "repetition"]) == ["gibberish", "repetition", "zero_advantage"] diff --git a/tests/unit/orchestrator/test_detectors.py b/tests/unit/orchestrator/test_detectors.py deleted file mode 100644 index 84e3730b51..0000000000 --- a/tests/unit/orchestrator/test_detectors.py +++ /dev/null @@ -1,269 +0,0 @@ -import math - -import verifiers.v1 as vf - -from prime_rl.configs.orchestrator import DetectorsConfig, GibberishDetectorConfig, RepetitionDetectorConfig -from prime_rl.orchestrator.detectors import ( - GibberishDetector, - RepetitionDetector, - detect, - drop_reasons, - setup_detectors, -) -from prime_rl.orchestrator.types import TrainRollout - - -def _assistant_node(token_ids: list[int], logprobs: list[float]) -> vf.MessageNode: - """An assistant node whose tokens are all model-sampled (the detectors read each node's - masked-True tokens + logprobs).""" - return vf.MessageNode( - message=vf.AssistantMessage(content="x"), - token_ids=token_ids, - mask=[True] * len(token_ids), - logprobs=logprobs, - ) - - -def _scaffold_assistant_node( - completion_ids: list[int], completion_logprobs: list[float], *, scaffold: int = 2 -) -> vf.MessageNode: - """A realistic v1 assistant node: a leading generation-prompt scaffold (mask=False, not - model-sampled) then the sampled completion. ``logprobs`` cover only the completion suffix - (vLLM returns logprobs for generated tokens only) — the exact layout where per-node - ``zip(token_ids, logprobs, mask)`` mispairs and the branch streams normalize.""" - return vf.MessageNode( - message=vf.AssistantMessage(content="x"), - token_ids=[1] * scaffold + completion_ids, - mask=[False] * scaffold + [True] * len(completion_ids), - logprobs=completion_logprobs, - ) - - -def _make_rollout( - completion_ids: list[int], - completion_logprobs: list[float], - *, - reward: float = 1.0, - multi_step: bool = False, -) -> TrainRollout: - """Build a ``TrainRollout`` (a message-graph trace) carrying the completion tokens — enough for - the detectors to inspect each node's sampled tokens / logprobs.""" - if multi_step: - mid = len(completion_ids) // 2 - nodes = [ - _assistant_node(completion_ids[:mid], completion_logprobs[:mid]), - _assistant_node(completion_ids[mid:], completion_logprobs[mid:]), - ] - else: - nodes = [_assistant_node(completion_ids, completion_logprobs)] - rollout = TrainRollout[vf.TaskData]( - task=vf.TraceTask(type="Task", data=vf.TaskData(idx=0, prompt="")), - agent=vf.AgentInfo(config=vf.AgentConfig()), - nodes=nodes, - rewards={"reward": vf.Reward(score=reward)}, - ) - rollout.env_name = "test" - return rollout - - -def _make_gibberish_detector(vocab_size=128_000, token_id_threshold=100_000, logprob_offset=2.0): - return GibberishDetector( - name="gibberish", - token_id_threshold=token_id_threshold, - logprob_threshold=-math.log(vocab_size) - logprob_offset, - ) - - -def _make_repetition_detector(window=5, prob_threshold=0.99): - return RepetitionDetector(name="repetition", window=window, logprob_threshold=math.log(prob_threshold)) - - -# --- GibberishDetector --- - - -def test_gibberish_detects_rare_low_prob_token(): - gibberish = _make_gibberish_detector() - - detected = gibberish.detect( - _make_rollout( - completion_ids=[50, 120_000, 80], - completion_logprobs=[-1.0, gibberish.logprob_threshold - 1.0, -0.5], - ) - ) - assert detected is True - - -def test_gibberish_ignores_normal_tokens(): - gibberish = _make_gibberish_detector() - - detected = gibberish.detect( - _make_rollout( - completion_ids=[10, 200, 5000], - completion_logprobs=[-1.0, -2.0, -3.0], - ) - ) - assert detected is False - - -def test_gibberish_ignores_high_prob_rare_token(): - gibberish = _make_gibberish_detector() - - detected = gibberish.detect( - _make_rollout( - completion_ids=[120_000], - completion_logprobs=[-0.5], - ) - ) - assert detected is False - - -def test_gibberish_works_across_trajectory_steps(): - gibberish = _make_gibberish_detector() - - detected = gibberish.detect( - _make_rollout( - completion_ids=[50, 60, 120_000, 80], - completion_logprobs=[-1.0, -0.5, gibberish.logprob_threshold - 1.0, -0.5], - multi_step=True, - ) - ) - assert detected is True - - -def test_gibberish_aligns_logprobs_under_generation_prompt_scaffold(): - """Regression: the assistant node carries a generation-prompt scaffold (mask=False) and - suffix-only logprobs, and the gibberish token is the LAST completion token. The old - per-node ``zip(token_ids, logprobs, mask)`` truncated at len(logprobs) and never examined - it; reading the aligned branch streams detects it.""" - gibberish = _make_gibberish_detector() - - rollout = TrainRollout[vf.TaskData]( - task=vf.TraceTask(type="Task", data=vf.TaskData(idx=0, prompt="")), - agent=vf.AgentInfo(config=vf.AgentConfig()), - nodes=[_scaffold_assistant_node([50, 80, 120_000], [-1.0, -0.5, gibberish.logprob_threshold - 1.0])], - rewards={"reward": vf.Reward(score=1.0)}, - ) - - detected = gibberish.detect(rollout) - assert detected is True - - -# --- RepetitionDetector --- - - -def test_repetition_triggers_after_window(): - repetition = _make_repetition_detector(window=5) - - detected = repetition.detect( - _make_rollout( - completion_ids=list(range(5)), - completion_logprobs=[-0.001] * 5, - ) - ) - assert detected is True - - -def test_repetition_no_trigger_below_window(): - repetition = _make_repetition_detector(window=5) - - detected = repetition.detect( - _make_rollout( - completion_ids=list(range(4)), - completion_logprobs=[-0.001] * 4, - ) - ) - assert detected is False - - -def test_repetition_resets_on_low_prob(): - repetition = _make_repetition_detector(window=5) - - logprobs = [-0.001] * 3 + [-2.0] + [-0.001] * 3 - detected = repetition.detect( - _make_rollout( - completion_ids=list(range(7)), - completion_logprobs=logprobs, - ) - ) - assert detected is False - - -def test_repetition_varied_probs_no_trigger(): - repetition = _make_repetition_detector(window=3) - - detected = repetition.detect( - _make_rollout( - completion_ids=list(range(6)), - completion_logprobs=[-0.001, -3.0, -0.001, -3.0, -0.001, -3.0], - ) - ) - assert detected is False - - -# --- setup + drop policy --- - - -def test_setup_detectors_builds_both_and_derives_thresholds(): - detectors = setup_detectors(DetectorsConfig(), vocab_size=128_000) - assert [d.name for d in detectors] == ["gibberish", "repetition"] - gibberish, repetition = detectors - assert gibberish.token_id_threshold == 100_000 - assert gibberish.logprob_threshold == -math.log(128_000) - 2.0 - assert repetition.logprob_threshold == math.log(0.99) - - -def test_setup_detectors_can_turn_one_off(): - config = DetectorsConfig(gibberish=None, repetition=RepetitionDetectorConfig(window=7)) - assert [d.name for d in setup_detectors(config, vocab_size=128_000)] == ["repetition"] - - -def test_setup_detectors_thresholds_follow_config(): - config = DetectorsConfig(gibberish=GibberishDetectorConfig(token_id_threshold=50_000, logprob_offset=1.0)) - gibberish = setup_detectors(config, vocab_size=1_000)[0] - assert (gibberish.token_id_threshold, gibberish.logprob_threshold) == (50_000, -math.log(1_000) - 1.0) - - -def test_detect_measures_every_detector(): - """Every detector measures every trace — one hit does not shadow the others, which is what - makes a detection rate a metric rather than a by-product of the drop order.""" - gibberish = _make_gibberish_detector() - rollout = _make_rollout( - completion_ids=[120_000] * 5, - completion_logprobs=[gibberish.logprob_threshold - 1.0] * 5, - ) - detect([gibberish, _make_repetition_detector(window=5)], rollout) - assert rollout.detections == {"gibberish": True, "repetition": False} - - -def _drop(rollout, drop=(), zero_advantage=True): - return drop_reasons(rollout, drop_detections=list(drop), drop_zero_advantage=zero_advantage) - - -def test_detection_only_drops_when_asked(): - rollout = _make_rollout(completion_ids=[1], completion_logprobs=[-1.0]) - rollout.detections = {"gibberish": True} - rollout.assign_advantages(0.5) - assert _drop(rollout) == [] # measured, but nothing asked for it to drop - assert _drop(rollout, drop=["gibberish"]) == ["gibberish"] - - -def test_zero_credit_drops_by_default(): - rollout = _make_rollout(completion_ids=[1], completion_logprobs=[-1.0]) - rollout.assign_advantages(0.0) - assert _drop(rollout) == ["zero_advantage"] - assert _drop(rollout, zero_advantage=False) == [] - - -def test_unscored_rollout_is_not_zero_credit(): - """opd/opsd assign no credit at all and train through reference KL — dropping them as - zero-advantage would ship an empty batch for every distillation run.""" - rollout = _make_rollout(completion_ids=[1], completion_logprobs=[-1.0]) - assert rollout.advantages is None - assert _drop(rollout) == [] - - -def test_reasons_accumulate(): - rollout = _make_rollout(completion_ids=[1], completion_logprobs=[-1.0]) - rollout.detections = {"gibberish": True, "repetition": True} - rollout.assign_advantages(0.0) - assert _drop(rollout, drop=["gibberish", "repetition"]) == ["gibberish", "repetition", "zero_advantage"] diff --git a/tests/unit/orchestrator/test_metrics.py b/tests/unit/orchestrator/test_metrics.py index c49ef7d409..332f21c7c5 100644 --- a/tests/unit/orchestrator/test_metrics.py +++ b/tests/unit/orchestrator/test_metrics.py @@ -43,7 +43,7 @@ def mk( trainable: bool = True, is_trainable: bool = True, is_filtered: bool = False, - detections: dict | None = None, + degeneracy: dict | None = None, setup: float = 0.0, agent: float = 0.0, agent_model: float = 0.0, @@ -74,7 +74,7 @@ def mk( agent=SimpleNamespace(trainable=trainable, name=agent_name), is_trainable=is_trainable, is_filtered=is_filtered, - detections=detections or {}, + degeneracy=degeneracy or {}, timing=SimpleNamespace( setup=SimpleNamespace(duration=setup), agent=SimpleNamespace( @@ -283,8 +283,8 @@ def test_nested_timing(): def test_train_only_metrics_absent_from_eval(): rollouts = [ - mk(is_trainable=True, is_filtered=True, detections={"gibberish": True}), - mk(is_trainable=False, detections={"gibberish": False}), + mk(is_trainable=True, is_filtered=True, degeneracy={"gibberish": True}), + mk(is_trainable=False, degeneracy={"gibberish": False}), ] out = train_wandb(rollouts) assert out["train/agg/all/agent/is_trainable/mean"] == 0.5 @@ -332,11 +332,11 @@ def test_traceless_episode_keeps_its_reason(): def test_training_state_is_train_only(): """Only a train rollout carries trainer-bound state; an eval trace has no field for it. Credit is not among them — it lives on the graph's nodes, which every trace has.""" - assert {"samples", "is_filtered", "detections"} <= set(TrainRollout.model_fields) + assert {"samples", "is_filtered", "degeneracy"} <= set(TrainRollout.model_fields) assert "advantages" not in TrainRollout.model_fields # derived from the nodes # An eval trace is the env's own, unextended — where a trace sits is the episode's to say. assert Rollout is vf.Trace - assert not {"samples", "is_filtered", "detections", "group_id", "episode_id"} & set(Rollout.model_fields) + assert not {"samples", "is_filtered", "degeneracy", "group_id", "episode_id"} & set(Rollout.model_fields) def test_inflight_episode_stamps_what_lands(): From 8c1e2f70019d58d9eb5a8aed7827c0a15373faea Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Wed, 5 Aug 2026 00:00:08 +0000 Subject: [PATCH 53/58] fix(orchestrator): the undefined names main's F821 rule surfaced Three lost imports and one stale reference in the prime monitor, which would have raised on the first sample upload. Co-Authored-By: Claude Fable 5 --- src/prime_rl/orchestrator/dispatcher.py | 1 + src/prime_rl/orchestrator/types.py | 2 +- src/prime_rl/utils/monitor/prime.py | 2 +- 3 files changed, 3 insertions(+), 2 deletions(-) diff --git a/src/prime_rl/orchestrator/dispatcher.py b/src/prime_rl/orchestrator/dispatcher.py index f8d85fea49..a65cd62f51 100644 --- a/src/prime_rl/orchestrator/dispatcher.py +++ b/src/prime_rl/orchestrator/dispatcher.py @@ -41,6 +41,7 @@ from prime_rl.orchestrator.eval_source import EvalSource from prime_rl.orchestrator.train_source import TrainSource from prime_rl.orchestrator.types import ( + Episode, GroupState, InflightEpisode, Policy, diff --git a/src/prime_rl/orchestrator/types.py b/src/prime_rl/orchestrator/types.py index a1b865ca3e..758f3a2908 100644 --- a/src/prime_rl/orchestrator/types.py +++ b/src/prime_rl/orchestrator/types.py @@ -5,7 +5,7 @@ import uuid from collections.abc import Callable, Iterable from dataclasses import dataclass -from typing import TYPE_CHECKING, Generic, Literal, Protocol, cast +from typing import TYPE_CHECKING, Any, Generic, Literal, Protocol, cast import verifiers.v1 as vf from pydantic import ConfigDict, Field diff --git a/src/prime_rl/utils/monitor/prime.py b/src/prime_rl/utils/monitor/prime.py index fe5460b858..468a63653d 100644 --- a/src/prime_rl/utils/monitor/prime.py +++ b/src/prime_rl/utils/monitor/prime.py @@ -273,7 +273,7 @@ def log_samples(self, episodes: list[Episode], step: int) -> None: assert step not in self._pending_sample_steps, f"Step {step} upload already in progress" assert self.logger is not None, "Logger is required for sample logging" - self.logger.info(f"Logging {len(rollouts)} samples to Prime Intellect API at step {step}") + self.logger.info(f"Logging {len(episodes)} episodes to Prime Intellect API at step {step}") start_time = time.perf_counter() parquet_bytes = self._rollouts_to_parquet_bytes(episodes, step) From 7bb1c5a811597013f0cfe15e4963782306074258 Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Wed, 5 Aug 2026 00:10:59 +0000 Subject: [PATCH 54/58] style: format with the ruff version CI pins Co-Authored-By: Claude Fable 5 --- src/prime_rl/orchestrator/train_sink.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/src/prime_rl/orchestrator/train_sink.py b/src/prime_rl/orchestrator/train_sink.py index 812dff285b..9feb1934ef 100644 --- a/src/prime_rl/orchestrator/train_sink.py +++ b/src/prime_rl/orchestrator/train_sink.py @@ -241,9 +241,7 @@ async def process_group(self, group_id: str) -> None: num_dropped = 0 for r in survivors: self.pre_filter_seen += 1 - reasons = drop_reasons( - r, drop=self.drop_degenerate, drop_zero_advantage=self.drop_zero_advantage - ) + reasons = drop_reasons(r, drop=self.drop_degenerate, drop_zero_advantage=self.drop_zero_advantage) r.is_filtered = bool(reasons) if reasons: self.pre_filter_dropped += 1 From afcf78bc06f644442d511c0db37d3b49b001a5ae Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Wed, 5 Aug 2026 00:38:13 +0000 Subject: [PATCH 55/58] refactor(orchestrator)!: read the policy span off the episode's metadata Co-Authored-By: Claude Fable 5 --- deps/verifiers | 2 +- src/prime_rl/orchestrator/orchestrator.py | 6 +++--- src/prime_rl/orchestrator/types.py | 6 +++--- tests/unit/orchestrator/test_metrics.py | 12 ++++++------ 4 files changed, 13 insertions(+), 13 deletions(-) diff --git a/deps/verifiers b/deps/verifiers index eddc781f22..9313f4b88e 160000 --- a/deps/verifiers +++ b/deps/verifiers @@ -1 +1 @@ -Subproject commit eddc781f22174ebefb749d0babf2367fdd034af3 +Subproject commit 9313f4b88e15e566db3dfb90254ef8da92f54b7f diff --git a/src/prime_rl/orchestrator/orchestrator.py b/src/prime_rl/orchestrator/orchestrator.py index 88f2604772..60a18e6b8b 100644 --- a/src/prime_rl/orchestrator/orchestrator.py +++ b/src/prime_rl/orchestrator/orchestrator.py @@ -794,7 +794,7 @@ def log_train_batch(self, batch: TrainBatch, *, step: int, step_time: float) -> n_effective = len(effective) n_trainable = sum(1 for r in effective if r.is_trainable) trainable_rate = (n_trainable / n_effective) if n_effective else 0.0 - max_off_policy = max((run_of(e).off_policy_steps or 0 for e in effective.episodes), default=0) + max_off_policy = max((run_of(e).metadata.off_policy_steps or 0 for e in effective.episodes), default=0) head = ( f"Step {step} | {format_time(step_time):>7} | Reward {eff.reward.mean():.4f} | " @@ -818,7 +818,7 @@ def log_train_batch(self, batch: TrainBatch, *, step: int, step_time: float) -> lines.append( f"╰─ {env_name:<{name_width}} | Ratio {ratio:.1%} | Reward {env_eff.reward.mean():.4f} | " f"Turns {env_eff.num_turns.mean():.1f} | Branches {env_eff.num_branches.mean():.1f} | " - f"Max Off-Policy {max((run_of(e).off_policy_steps or 0 for e in env_eff_pool.episodes), default=0)} | " + f"Max Off-Policy {max((run_of(e).metadata.off_policy_steps or 0 for e in env_eff_pool.episodes), default=0)} | " f"Error {pool.metrics.has_error.mean():.1%} | Truncation {env_eff.is_truncated.mean():.1%}" ) get_logger().success("\n\t\t ".join(lines)) @@ -839,7 +839,7 @@ async def finalize_eval_batch(self, batch: EvalBatch) -> None: save_episodes, records, get_trace_path(self.config.output_dir, batch.step, "eval", "effective") ) self.monitor.log_eval_samples(batch.rollouts.episodes, env_name=batch.env_name, step=batch.step) - policy_versions = {run.policy.start for e in batch.rollouts.episodes if (run := run_of(e)).policy} + policy_versions = {m.policy.start for e in batch.rollouts.episodes if (m := run_of(e).metadata).policy} policy_version = min(policy_versions, default=0) if len(policy_versions) > 1: get_logger().warning( diff --git a/src/prime_rl/orchestrator/types.py b/src/prime_rl/orchestrator/types.py index 758f3a2908..5262912615 100644 --- a/src/prime_rl/orchestrator/types.py +++ b/src/prime_rl/orchestrator/types.py @@ -199,10 +199,10 @@ def stamp(self, episode: Episode, *, run_id: str, policy: vf.PolicySpan | None, episode.group = vf.GroupInfo(id=str(self.group_id)) if self.kind == "eval": assert eval_step is not None, "eval episode missing its step" - metadata: vf.EpisodeMetadata = vf.EvalMetadata(step=eval_step) + metadata: vf.EpisodeMetadata = vf.EvalMetadata(step=eval_step, policy=policy) else: - metadata = vf.TrainMetadata() - episode.run = vf.TrainRunInfo(id=run_id, metadata=metadata, policy=policy) + metadata = vf.TrainMetadata(policy=policy) + episode.run = vf.TrainRunInfo(id=run_id, metadata=metadata) return episode diff --git a/tests/unit/orchestrator/test_metrics.py b/tests/unit/orchestrator/test_metrics.py index 332f21c7c5..c2d96262da 100644 --- a/tests/unit/orchestrator/test_metrics.py +++ b/tests/unit/orchestrator/test_metrics.py @@ -348,13 +348,13 @@ def test_inflight_episode_stamps_what_lands(): span = vf.PolicySpan(start=3, end=4) # an update landed while it was generating train = inflight.stamp(wire, run_id="r", policy=span, eval_step=None) run = run_of(train) - assert isinstance(run.metadata, vf.TrainMetadata) and run.policy == span + assert isinstance(run.metadata, vf.TrainMetadata) and run.metadata.policy == span assert run.metadata.step is None # the batch window it lands in is not known yet - assert run.off_policy_steps is None # so there is nothing to be behind yet + assert run.metadata.off_policy_steps is None # so there is nothing to be behind yet assert train.env.name == "rt" and train.group is not None run.metadata.step = 6 # the window it landed in, which step 6 trains v5 from - assert run.off_policy_steps == 2 and run.policy.drift == 1 + assert run.metadata.off_policy_steps == 2 and run.metadata.policy.drift == 1 evaluation = replace(inflight, kind="eval").stamp( vf.WireEpisode.model_construct(id="e", traces=[]), run_id="r", policy=span, eval_step=12 @@ -362,9 +362,9 @@ def test_inflight_episode_stamps_what_lands(): # An online eval belongs to the same training run — same id, told apart by its metadata. Its # off-policy reading is what drifted under it, not its distance from a step it was dispatched # at: nothing trains on an eval, and a slow one would otherwise look on-policy forever. - eval_run = run_of(evaluation) - assert isinstance(eval_run.metadata, vf.EvalMetadata) and eval_run.metadata.step == 12 - assert eval_run.id == run.id and eval_run.off_policy_steps == span.drift == 1 + eval_meta = run_of(evaluation).metadata + assert isinstance(eval_meta, vf.EvalMetadata) and eval_meta.step == 12 + assert run_of(evaluation).id == run.id and eval_meta.off_policy_steps == span.drift == 1 with pytest.raises(AssertionError): # an eval episode without its step is not representable replace(inflight, kind="eval").stamp(wire, run_id="r", policy=span, eval_step=None) From 84e2daf006211cbc15f06e88ce5eb472d5893086 Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Wed, 5 Aug 2026 00:56:13 +0000 Subject: [PATCH 56/58] fix(orchestrator): a rollout is the wire trace an episode actually holds MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Rollout aliased the bare generic while Episode aliased the wire form, and TrainRollout resolved its agent config to the strict default — so the traces prime-rl declared were a different specialization from the ones WireEpisode carries. model_construct skipped validation, so nothing ever said so. Co-Authored-By: Claude Fable 5 --- src/prime_rl/orchestrator/types.py | 6 ++++-- tests/unit/orchestrator/test_advantage.py | 5 +++-- tests/unit/orchestrator/test_algorithms.py | 7 ++++--- tests/unit/orchestrator/test_degeneracy.py | 3 ++- tests/unit/orchestrator/test_metrics.py | 5 +++-- tests/unit/utils/test_prime_monitor.py | 5 +++-- 6 files changed, 19 insertions(+), 12 deletions(-) diff --git a/src/prime_rl/orchestrator/types.py b/src/prime_rl/orchestrator/types.py index 5262912615..68c93c6215 100644 --- a/src/prime_rl/orchestrator/types.py +++ b/src/prime_rl/orchestrator/types.py @@ -9,6 +9,8 @@ import verifiers.v1 as vf from pydantic import ConfigDict, Field +from verifiers.v1.configs.agent import WireAgentConfig +from verifiers.v1.state import State from verifiers.v1.task import DataT from verifiers.v1.trace import EXCLUDE_FIELDS @@ -59,12 +61,12 @@ class GroupState: policy_version_at_start: int = 0 -Rollout = vf.Trace +Rollout = vf.WireTrace """The env's own trace, unextended — what an eval rollout is, and the read side of a training one. Where a trace sits (its episode, its group, its env) is the episode's to say.""" -class TrainRollout(vf.Trace[DataT], Generic[DataT]): +class TrainRollout(vf.Trace[DataT, State, WireAgentConfig], Generic[DataT]): """A rollout on the training path — the one place prime-rl extends a verifiers type, because trainer-bound state has nowhere else to live: the samples built from its branches, what the degeneracy measurements found on it, and whether the drop policy kept it. All of it is diff --git a/tests/unit/orchestrator/test_advantage.py b/tests/unit/orchestrator/test_advantage.py index 95505a28f1..6e7cb269d8 100644 --- a/tests/unit/orchestrator/test_advantage.py +++ b/tests/unit/orchestrator/test_advantage.py @@ -2,6 +2,7 @@ import pytest import verifiers.v1 as vf +from verifiers.v1.configs.agent import WireAgentConfig from prime_rl.configs.algorithm import ( GRPOAlgoConfig, @@ -100,7 +101,7 @@ def _take(n: int) -> list[int]: rollout = TrainRollout[vf.TaskData]( task=vf.TraceTask(type="Task", data=vf.TaskData(idx=0, prompt=None)), - agent=vf.AgentInfo(config=vf.AgentConfig()), + agent=vf.AgentInfo(config=WireAgentConfig()), nodes=nodes, calls=calls, rewards={"reward": vf.Reward(score=reward)}, @@ -298,7 +299,7 @@ def test_stamp_advantages_zeros_a_shared_node_in_the_later_branch(): ] rollout = TrainRollout[vf.TaskData]( task=vf.TraceTask(type="Task", data=vf.TaskData(idx=0, prompt=None)), - agent=vf.AgentInfo(config=vf.AgentConfig()), + agent=vf.AgentInfo(config=WireAgentConfig()), nodes=[root, *leaves], rewards={"reward": vf.Reward(score=0.0)}, ) diff --git a/tests/unit/orchestrator/test_algorithms.py b/tests/unit/orchestrator/test_algorithms.py index 2e219360b0..aa5a2da772 100644 --- a/tests/unit/orchestrator/test_algorithms.py +++ b/tests/unit/orchestrator/test_algorithms.py @@ -4,6 +4,7 @@ import pydantic import pytest import verifiers.v1 as vf +from verifiers.v1.configs.agent import WireAgentConfig from verifiers.v1.graph import MessageNode from verifiers.v1.types import AssistantMessage, ToolMessage, UserMessage @@ -173,7 +174,7 @@ def _make_rollout(samples: list[TrainingSample]) -> TrainRollout: ] rollout = TrainRollout( task=vf.TraceTask(type="Task", data=vf.TaskData(idx=0, prompt=None)), - agent=vf.AgentInfo(config=vf.AgentConfig()), + agent=vf.AgentInfo(config=WireAgentConfig()), nodes=nodes, rewards={}, env_name="test-env", @@ -252,7 +253,7 @@ def _two_turn_rollout(observation_role: str = "tool") -> TrainRollout: ] rollout = TrainRollout( task=vf.TraceTask(type="Task", data=vf.TaskData(idx=0, prompt=None)), - agent=vf.AgentInfo(config=vf.AgentConfig()), + agent=vf.AgentInfo(config=WireAgentConfig()), nodes=nodes, rewards={"r": vf.Reward(score=1.0)}, env_name="test-env", @@ -304,7 +305,7 @@ def test_echo_weights_only_content_tokens_when_is_content_present(): ] rollout = TrainRollout( task=vf.TraceTask(type="Task", data=vf.TaskData(idx=0, prompt=None)), - agent=vf.AgentInfo(config=vf.AgentConfig()), + agent=vf.AgentInfo(config=WireAgentConfig()), nodes=nodes, rewards={"r": vf.Reward(score=1.0)}, env_name="test-env", diff --git a/tests/unit/orchestrator/test_degeneracy.py b/tests/unit/orchestrator/test_degeneracy.py index 975ec33789..c21b955b3c 100644 --- a/tests/unit/orchestrator/test_degeneracy.py +++ b/tests/unit/orchestrator/test_degeneracy.py @@ -1,6 +1,7 @@ import math import verifiers.v1 as vf +from verifiers.v1.configs.agent import WireAgentConfig from prime_rl.orchestrator.degeneracy import ( REPETITION_PROB, @@ -33,7 +34,7 @@ def _assistant_node(token_ids: list[int], logprobs: list[float]) -> vf.MessageNo def _make_rollout(nodes: list[vf.MessageNode]) -> TrainRollout: rollout = TrainRollout[vf.TaskData]( task=vf.TraceTask(type="Task", data=vf.TaskData(idx=0, prompt="")), - agent=vf.AgentInfo(config=vf.AgentConfig()), + agent=vf.AgentInfo(config=WireAgentConfig()), nodes=nodes, rewards={"reward": vf.Reward(score=1.0)}, ) diff --git a/tests/unit/orchestrator/test_metrics.py b/tests/unit/orchestrator/test_metrics.py index c2d96262da..4ad05973ec 100644 --- a/tests/unit/orchestrator/test_metrics.py +++ b/tests/unit/orchestrator/test_metrics.py @@ -334,8 +334,9 @@ def test_training_state_is_train_only(): Credit is not among them — it lives on the graph's nodes, which every trace has.""" assert {"samples", "is_filtered", "degeneracy"} <= set(TrainRollout.model_fields) assert "advantages" not in TrainRollout.model_fields # derived from the nodes - # An eval trace is the env's own, unextended — where a trace sits is the episode's to say. - assert Rollout is vf.Trace + # An eval trace is the env's own, unextended — where a trace sits is the episode's to say, + # and it is the same wire specialization the episode's ``traces`` hold. + assert Rollout is vf.WireTrace assert not {"samples", "is_filtered", "degeneracy", "group_id", "episode_id"} & set(Rollout.model_fields) diff --git a/tests/unit/utils/test_prime_monitor.py b/tests/unit/utils/test_prime_monitor.py index 92787f369f..e83ef19800 100644 --- a/tests/unit/utils/test_prime_monitor.py +++ b/tests/unit/utils/test_prime_monitor.py @@ -4,6 +4,7 @@ import pyarrow.parquet as pq import verifiers.v1 as vf +from verifiers.v1.configs.agent import WireAgentConfig from prime_rl.orchestrator.types import TrainRollout from prime_rl.utils.monitor.prime import PrimeMonitor @@ -37,7 +38,7 @@ def _build_rollout(*, example_id: int, reward: float, task: str) -> TrainRollout ] rollout = TrainRollout[vf.TaskData]( task=vf.TraceTask(type="Task", data=vf.TaskData(idx=example_id, prompt=f"prompt-{example_id}")), - agent=vf.AgentInfo(config=vf.AgentConfig()), + agent=vf.AgentInfo(config=WireAgentConfig()), nodes=nodes, rewards={"reward": vf.Reward(score=reward)}, ) @@ -86,7 +87,7 @@ def test_rollouts_to_parquet_bytes_skips_rollouts_without_trajectory(): rollout_with_branches = _build_rollout(example_id=1, reward=1.0, task="task-a") rollout_without_branches = TrainRollout[vf.TaskData]( task=vf.TraceTask(type="Task", data=vf.TaskData(idx=2, prompt="missing-trajectory")), - agent=vf.AgentInfo(config=vf.AgentConfig()), + agent=vf.AgentInfo(config=WireAgentConfig()), ) assert rollout_without_branches.branches == [] From 890db838afa8576c6daadfeb22141153ddb3510d Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Wed, 5 Aug 2026 21:54:24 +0000 Subject: [PATCH 57/58] refactor(orchestrator): the comparison group is prime-rl's, on episode.info Verifiers has no notion of a group, so carrying one there was a field it held for a single consumer. It rides in info, which is where an episode takes a consumer's metadata, and it still lands on the saved record so a row stays placeable. Co-Authored-By: Claude Fable 5 --- deps/verifiers | 2 +- src/prime_rl/orchestrator/types.py | 12 +++++++++--- tests/unit/orchestrator/test_metrics.py | 6 ++++-- 3 files changed, 14 insertions(+), 6 deletions(-) diff --git a/deps/verifiers b/deps/verifiers index 9313f4b88e..3a38df548e 160000 --- a/deps/verifiers +++ b/deps/verifiers @@ -1 +1 @@ -Subproject commit 9313f4b88e15e566db3dfb90254ef8da92f54b7f +Subproject commit 3a38df548e27ef58c3a2e7b151fc7962a64f8d8a diff --git a/src/prime_rl/orchestrator/types.py b/src/prime_rl/orchestrator/types.py index 68c93c6215..0a81081ccf 100644 --- a/src/prime_rl/orchestrator/types.py +++ b/src/prime_rl/orchestrator/types.py @@ -147,11 +147,17 @@ def to_record(episode: Episode) -> dict[str, Any]: return episode.model_dump(mode="json", exclude={"traces": {"__all__": EXCLUDE_FIELDS}}) +GROUP_ID = "group_id" +"""``Episode.info`` key for the comparison group. Verifiers has no notion of one — a group is the +consumer's: the episodes it planned from one task and scores against each other.""" + + def group_id_of(episode: Episode) -> str: """The group an episode was planned in. The dispatcher plans every episode into one, so this is always set by the time anything downstream asks.""" - assert episode.group is not None, "the dispatcher plans every episode into a group" - return episode.group.id + group_id = episode.info.get(GROUP_ID) + assert isinstance(group_id, str), "the dispatcher plans every episode into a group" + return group_id def env_name_of(episode: Episode) -> str: @@ -198,7 +204,7 @@ def stamp(self, episode: Episode, *, run_id: str, policy: vf.PolicySpan | None, not — it belongs to whichever batch window is collecting when it lands, so the main loop fills that in.""" episode.env.name = self.env_name - episode.group = vf.GroupInfo(id=str(self.group_id)) + episode.info[GROUP_ID] = str(self.group_id) if self.kind == "eval": assert eval_step is not None, "eval episode missing its step" metadata: vf.EpisodeMetadata = vf.EvalMetadata(step=eval_step, policy=policy) diff --git a/tests/unit/orchestrator/test_metrics.py b/tests/unit/orchestrator/test_metrics.py index 4ad05973ec..d30ec2cfb4 100644 --- a/tests/unit/orchestrator/test_metrics.py +++ b/tests/unit/orchestrator/test_metrics.py @@ -10,10 +10,12 @@ from prime_rl.orchestrator.eval_sink import eval_step_of from prime_rl.orchestrator.metrics import EvalRollouts, Stat, TrainRollouts from prime_rl.orchestrator.types import ( + GROUP_ID, Episode, InflightEpisode, Rollout, TrainRollout, + group_id_of, rollouts_of, run_of, ) @@ -96,7 +98,7 @@ def ep(*rollouts, env_name: str = "env", errors=(), group_id="g0", cls=Episode): id=f"e{next(_ids)}", traces=list(rollouts), env=vf.EnvInfo(name=env_name), - group=vf.GroupInfo(id=group_id), + info={GROUP_ID: group_id}, errors=list(errors), ) @@ -352,7 +354,7 @@ def test_inflight_episode_stamps_what_lands(): assert isinstance(run.metadata, vf.TrainMetadata) and run.metadata.policy == span assert run.metadata.step is None # the batch window it lands in is not known yet assert run.metadata.off_policy_steps is None # so there is nothing to be behind yet - assert train.env.name == "rt" and train.group is not None + assert train.env.name == "rt" and group_id_of(train) is not None run.metadata.step = 6 # the window it landed in, which step 6 trains v5 from assert run.metadata.off_policy_steps == 2 and run.metadata.policy.drift == 1 From 937c94c84eb6b12a4ed501fcc8c65bcb6bc3fc02 Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Wed, 5 Aug 2026 22:23:14 +0000 Subject: [PATCH 58/58] revert(orchestrator): restore the filters as they were MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The detector/drop-policy split is a change of its own and does not need to ride along with the episode work — it goes in a follow-up. filters.py, its config, its tests and its docs are back to main; TrainRollout keeps is_filtered and filter_results. Co-Authored-By: Claude Fable 5 --- configs/ci/nightly-fft/wiki-search.toml | 4 + configs/debug/algo/echo.toml | 10 +- docs/algorithms.md | 59 +-- examples/advanced/glm-5.2/swe.toml | 4 +- examples/advanced/intellect-3.1/rl.toml | 4 + examples/basic/wiki-search/rl.toml | 4 + .../src/prime_rl/configs/orchestrator.py | 78 +++- src/prime_rl/orchestrator/algo/base.py | 6 +- src/prime_rl/orchestrator/degeneracy.py | 84 ---- src/prime_rl/orchestrator/filters.py | 172 ++++++++ src/prime_rl/orchestrator/metrics.py | 8 +- src/prime_rl/orchestrator/orchestrator.py | 23 +- src/prime_rl/orchestrator/train_sink.py | 66 +-- src/prime_rl/orchestrator/types.py | 11 +- tests/unit/orchestrator/test_degeneracy.py | 172 -------- tests/unit/orchestrator/test_filters.py | 407 ++++++++++++++++++ tests/unit/orchestrator/test_metrics.py | 16 +- 17 files changed, 760 insertions(+), 368 deletions(-) delete mode 100644 src/prime_rl/orchestrator/degeneracy.py create mode 100644 src/prime_rl/orchestrator/filters.py delete mode 100644 tests/unit/orchestrator/test_degeneracy.py create mode 100644 tests/unit/orchestrator/test_filters.py diff --git a/configs/ci/nightly-fft/wiki-search.toml b/configs/ci/nightly-fft/wiki-search.toml index 62b40a2767..9814e6351b 100644 --- a/configs/ci/nightly-fft/wiki-search.toml +++ b/configs/ci/nightly-fft/wiki-search.toml @@ -19,6 +19,10 @@ batch_size = 512 group_size = 16 oversampling_factor = 2.0 +[[orchestrator.pre_batch_filters]] +type = "zero_advantage" +enforce = true + [[orchestrator.train.source]] name = "wiki-search" diff --git a/configs/debug/algo/echo.toml b/configs/debug/algo/echo.toml index 1e66648c8d..c05beb865e 100644 --- a/configs/debug/algo/echo.toml +++ b/configs/debug/algo/echo.toml @@ -16,10 +16,6 @@ name = "debug-echo" batch_size = 32 group_size = 4 -# ECHO learns from observation tokens even when the GRPO advantage collapses -# to zero — keep zero-advantage rollouts in the batch. -drop_zero_advantage = false - # alphabet-sort's feedback arrives as user messages, so train the user role # instead of echo's tool default. [orchestrator.algo] @@ -48,6 +44,12 @@ type = "subprocess" [orchestrator.train.sampling] max_completion_tokens = 512 +# ECHO learns from observation tokens even when the GRPO advantage collapses +# to zero — keep zero-advantage rollouts in the batch. +[[orchestrator.post_batch_filters]] +type = "zero_advantage" +enforce = false + # Fine-tune inherits the PrimeIntellect Qwen3 template byte-for-byte. [orchestrator.renderer] name = "prime-qwen3" diff --git a/docs/algorithms.md b/docs/algorithms.md index a324aad635..a2bcd1e7db 100644 --- a/docs/algorithms.md +++ b/docs/algorithms.md @@ -1,6 +1,6 @@ # Algorithms -This page covers the math and the configurable algorithmic components: the algorithm abstraction and its algorithms, how off-policy training works, the loss components and advantage functions, how to plug in your own, the degeneracy measurements and drop policy applied between rollout and training, and how multi-turn rollouts get merged into training samples. +This page covers the math and the configurable algorithmic components: the algorithm abstraction and its algorithms, how off-policy training works, the loss components and advantage functions, how to plug in your own, the filters applied between rollout and training, and how multi-turn rollouts get merged into training samples. ## Table of Contents @@ -21,7 +21,7 @@ This page covers the math and the configurable algorithmic components: the algor - [Self-Play Advantage (RAE)](#self-play-advantage-rae) - [Authoring an Algorithm](#authoring-an-algorithm) - [Reference Scoring](#reference-scoring) -- [Degeneracy and the drop policy](#degeneracy-and-the-drop-policy) +- [Filters](#filters) - [Multi-Turn Trajectories](#multi-turn-trajectories) - [Extension Property](#extension-property) - [Best-Effort Interleaving](#best-effort-interleaving) @@ -164,12 +164,12 @@ At runtime, each env's resolved config builds two objects: a `Sampler` (`prime_r | `hierarchical_grpo` | `HierarchicalGRPOAlgorithm` | `score_group`: GRPO baseline per episode for solvers, per group for the proposer | | `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 the drop policy) | +| `sft` | `SFTDistillAlgorithm` | `score_group`: group-norm credit (feeds filters) | Each class owns its hooks outright — reading one top to bottom reads the algorithm, and everything on the class is an override point. The two hooks are one scope-and-timing ladder — the wider scope is unlocked by a later barrier, so the two axes coincide. Each is handed the env's own data directly — a `Rollout` (the typed trace: `reward`, `nodes`, `num_turns`, ... with `samples` attached) on arrival, the group's `Episode`s at group time — plus `assign_advantages` to write credit: -- `async score_rollout(rollout)` — one rollout, **on arrival** (as it's tokenized, before its group is complete): rollout-local credit (`rollout.assign_advantages(...)`), observation ce weights, **or** model I/O — query a reference pool (e.g. `self.teacher_pool`, connected in `setup()` via `self.connect(...)`, or the live `self.policy_pool` for opsd) and attach per-token results (e.g. teacher logprobs) with bounded concurrency. No siblings. `echo` weights observation tokens here, identifying env-provided observation nodes by their non-sampled status and source step role attribution, applying the optional user filter, and writing the `ce_weights` stream. Model I/O runs *before* the drop policy, so it pays compute on rollouts that may then be dropped. -- `score_group(group)` — the cohort, **before the drop policy** (which reads the streams), synchronous: group-relative credit (GRPO/MaxRL baselines). `group` is a list of `TrainEpisode`, so an algorithm can compare within an episode as well as across them (`hierarchical_grpo` does); `group_rollouts(group)` flattens it for the ones that only compare across. +- `async score_rollout(rollout)` — one rollout, **on arrival** (as it's tokenized, before its group is complete): rollout-local credit (`rollout.assign_advantages(...)`), observation ce weights, **or** model I/O — query a reference pool (e.g. `self.teacher_pool`, connected in `setup()` via `self.connect(...)`, or the live `self.policy_pool` for opsd) and attach per-token results (e.g. teacher logprobs) with bounded concurrency. No siblings. `echo` weights observation tokens here, identifying env-provided observation nodes by their non-sampled status and source step role attribution, applying the optional user filter, and writing the `ce_weights` stream. Model I/O runs *before* the pre-batch filters, so it pays compute on rollouts that may then be filtered out. +- `score_group(group)` — the cohort, **before filtering** (filters read the streams), synchronous: group-relative credit (GRPO/MaxRL baselines). `group` is a list of `TrainEpisode`, so an algorithm can compare within an episode as well as across them (`hierarchical_grpo` does); `group_rollouts(group)` flattens it for the ones that only compare across. The pipeline drives the hooks through two non-virtual methods it never looks inside: `algorithm.finalize_rollout(rollout)` per arrival (rollout-local scoring + reference I/O) and `algorithm.finalize_group(episodes)` per group (scoring + wire stamping; after this the records are frozen — groups die at stamping). Sample construction (interleaving) is pure pipeline — observation-token provenance is available through structural attribution (`node.sampled`, `node.is_content`) for any algorithm that trains on env-provided tokens. @@ -305,8 +305,8 @@ The per-token training signal is set by `algo.type` and the [algorithm](#the-alg | `rae` | `rl` | Reward minus a per-agent EMA baseline (SPIRAL's role-conditioned advantage estimation) — for multi-agent self-play envs. | | `hierarchical_grpo` | `rl` | GRPO for proposer-solver envs: solvers are compared within one proposed problem, while proposers are compared across proposals. | | `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 (`teacher`, an inline frozen hosted model), evaluated in the trainer from shipped reference logprobs. No credit — rollouts keep `advantages = None` (so zero-advantage dropping never applies) 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` (so zero-advantage dropping never applies) and ship no advantage stream. | +| `opd` | `ref_kl` | On-policy distillation: per-token reverse KL to a reference model (`teacher`, 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. | | `sft` | `ce` | Cross-entropy on the sampled tokens. Assigns no advantage — trains on every sampled token. | ### Default Advantage @@ -441,7 +441,7 @@ class MyAlgorithm(Algorithm): Add a typed `MyAlgoConfig` to `prime_rl.configs.algorithm` and its discriminated union, then register `"my_algo": MyAlgorithm` in `ALGORITHM_CLASSES`. Pick the hook by *when* your signal is ready: `score_rollout` for per-arrival credit or credit that needs a model call (it's `async`), `score_group` for group-relative credit (GRPO/MaxRL). `assign_advantages` takes a scalar and broadcasts it over the rollout's trainable tokens, writing it onto the graph's nodes so branches sharing a node cannot disagree about its credit. -The drop policy and metrics derive from the streams (zero-advantage dropping checks for an all-zero stream; logged distributions use per-rollout means). Signals that depend on the live policy's weights (like OPD's reverse KL) cannot be precomputed here; those are reference-scoring algorithms, evaluated in the trainer. +Advantage-based filters and metrics derive from the streams (the zero-advantage filter checks for all-zero streams; logged distributions use per-rollout means). Signals that depend on the live policy's weights (like OPD's reverse KL) cannot be precomputed here; those are reference-scoring algorithms, evaluated in the trainer. ### Reference Scoring @@ -456,43 +456,30 @@ type = "opsd" demo_key = "demonstration" ``` -Scoring runs at arrival, *before* the drop policy, so a rollout that is later dropped still cost its reference compute — accepted for the simpler one-rollout-at-a-time shape (zero-advantage dropping never applies to opd/opsd anyway, since neither assigns an advantage). +Scoring runs at arrival, *before* the pre-batch filters, so a rollout that is later filtered still cost its reference compute — accepted for the simpler one-rollout-at-a-time shape (advantage-based filters never fire for opd/opsd anyway, since neither assigns an advantage). -## Degeneracy and the drop policy +## Filters -Two separate things: **measuring** what a rollout looks like, and **deciding** whether to train on it. +Filters drop rollouts between scoring and training. Built-ins (composable): -Every trace is measured, unconditionally, as soon as it is tokenized, and each measurement is -reported per agent alongside reward and truncation. Nothing configures that — a rate is only a rate -if nothing decided in advance which traces to look at, and every measurement runs on every trace, so -one never hides another. - -| Measurement | What it finds | +| Filter | Effect | |---|---| -| `gibberish` | a rare token (high BPE id) generated at high entropy — degenerate output | -| `repetition` | a long stretch of very-high-confidence tokens — a repetition loop | +| `gibberish` | Drops rollouts whose mean log-prob fall below a threshold — usually a sign of degenerate output. | +| `repetition` | Drops rollouts with high n-gram repetition. | +| `zero_advantage` | Drops rollouts whose advantage is zero, so the trainer doesn't waste tokens on them. | -The drop policy decides. It runs once, when a finalized group's credit is assigned and the rollouts -would enter the batch buffer, and it has two inputs — the measurements the run chose to act on, and -**zero credit**, which drops on its own: +The default `[orchestrator]` config registers all three in both filter slots: `post_batch_filters` enforce by default (flagged rollouts are recorded but not shipped to the trainer), while `pre_batch_filters` run in monitor mode (`enforce = false`); flip `enforce = true` there to drop matching rollouts before they consume a slot in the batch. Setting a slot replaces its defaults wholesale: ```toml -[orchestrator] -drop_degenerate = ["gibberish"] # measuring is unconditional; acting on it is opt-in -drop_zero_advantage = true # the default -``` +[[orchestrator.post_batch_filters]] +type = "zero_advantage" -A scored rollout whose every token is worth nothing — a GRPO group where all rollouts earned the -same reward — produces no gradient, so training on it is a wasted forward pass. It is not one of the -measurements because it is not a property of the generation: it is only knowable after the group is -scored, and it is what `TrainRollout.is_trainable` already means. - -A rollout that was **never** scored is not zero-credit. `opd` / `opsd` assign no advantages at all -and train through reference KL, so they are never dropped by this rule. +[[orchestrator.post_batch_filters]] +type = "repetition" +threshold = 0.4 +``` -Dropped rollouts still appear in the metrics window and the `all` trace file — they just don't ship. -`{scope}/{subset}//detected//mean` is each measurement's rate, and -`{scope}/{subset}//is_filtered/mean` is the share the policy dropped. +Filtered rollouts still appear in W&B distributions, just not in the trainer batch — useful for spotting whether filtering is doing its job. ## Multi-Turn Trajectories diff --git a/examples/advanced/glm-5.2/swe.toml b/examples/advanced/glm-5.2/swe.toml index 6ff2d86f6b..c397f717f6 100644 --- a/examples/advanced/glm-5.2/swe.toml +++ b/examples/advanced/glm-5.2/swe.toml @@ -60,7 +60,6 @@ weight_decay = 0.1 [orchestrator] -drop_degenerate = ["gibberish"] batch_size = 4096 group_size = 16 oversampling_factor = 3 @@ -85,6 +84,9 @@ id = "bash" type = "prime" labels = ["glm5-pd-disag", "swe-bench-verified"] +[[orchestrator.post_batch_filters]] +type = "gibberish" +enforce = true [inference] enable_expert_parallel = true diff --git a/examples/advanced/intellect-3.1/rl.toml b/examples/advanced/intellect-3.1/rl.toml index a22f5ab7fd..174639662f 100644 --- a/examples/advanced/intellect-3.1/rl.toml +++ b/examples/advanced/intellect-3.1/rl.toml @@ -119,6 +119,10 @@ id = "null" [orchestrator.train.source.env.agent.runtime] type = "subprocess" +[[orchestrator.pre_batch_filters]] +type = "zero_advantage" +enforce = true + [orchestrator.eval] interval = 25 diff --git a/examples/basic/wiki-search/rl.toml b/examples/basic/wiki-search/rl.toml index 068f0868ba..8f9d37945a 100644 --- a/examples/basic/wiki-search/rl.toml +++ b/examples/basic/wiki-search/rl.toml @@ -40,6 +40,10 @@ name = "qwen3-4b-wiki-search" [orchestrator.train.sampling] max_completion_tokens = 512 +[[orchestrator.pre_batch_filters]] +type = "zero_advantage" +enforce = true + [[orchestrator.train.source]] name = "wiki-search" diff --git a/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py b/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py index 767e196fa6..4030ccbd44 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py @@ -362,6 +362,49 @@ class CheckpointConfig(BaseConfig): # Flags rare tokens generated at high entropy (Section 5.2, https://arxiv.org/abs/2510.02387). +class GibberishFilterConfig(BaseConfig): + type: Literal["gibberish"] = "gibberish" + + enforce: bool = False + """When True, skip detected rollouts entirely so they are not sent to the trainer. When False, only track detection metrics.""" + + token_id_threshold: int = 100_000 + """Token IDs above this are candidates for gibberish. BPE tokens are sorted by merge order.""" + + logprob_offset: float = 2.0 + """Offset from uniform-distribution logprob. Threshold = ``-log(vocab_size) - logprob_offset``.""" + + +# Flags rollouts stuck in a repetition loop: emits high-confidence tokens for an extended stretch. +# Flagged when `window` consecutive tokens are each sampled with probability above `prob_threshold`. +# (Section 3.2, https://arxiv.org/abs/2506.13585) +class RepetitionFilterConfig(BaseConfig): + type: Literal["repetition"] = "repetition" + + enforce: bool = False + """When True, skip detected rollouts entirely so they are not sent to the trainer. When False, only track detection metrics.""" + + window: int = Field(3_000, ge=1) + """Consecutive high-probability steps required to flag the rollout.""" + + prob_threshold: float = Field(0.99, gt=0, le=1) + """Tokens sampled with probability above this are considered repetitive. Consecutive such tokens count toward the window.""" + + +# Flags rollouts with zero advantage. +class ZeroAdvantageFilterConfig(BaseConfig): + type: Literal["zero_advantage"] = "zero_advantage" + + enforce: bool = True + """When True, skip detected rollouts entirely so they are not sent to the trainer. When False, only track detection metrics.""" + + +FilterConfig: TypeAlias = Annotated[ + GibberishFilterConfig | RepetitionFilterConfig | ZeroAdvantageFilterConfig, + Field(discriminator="type"), +] + + class FileSystemWeightBroadcastConfig(BaseConfig): type: Literal["filesystem"] = "filesystem" @@ -435,14 +478,23 @@ class OrchestratorConfig(BaseConfig): eval: EvalConfig | None = None """Evaluation configuration.""" - drop_degenerate: list[Literal["gibberish", "repetition"]] = [] - """Which degeneracy measurements also keep a rollout out of the training batch. Every trace is - measured for all of them and reports them per agent; acting on one is opt-in.""" - - drop_zero_advantage: bool = True - """Keep scored rollouts whose credit is all zero out of the batch — a GRPO group that earned a - uniform reward carries no gradient, so training on it is wasted compute. Rollouts that were - never scored (opd/opsd train through reference KL, not credit) are unaffected.""" + pre_batch_filters: list[FilterConfig] = [ + GibberishFilterConfig(enforce=False), + RepetitionFilterConfig(enforce=False), + ZeroAdvantageFilterConfig(enforce=False), + ] + """Filters applied *before* a rollout enters the training batch buffer. + All three filter types are registered in monitor mode by default; flip ``enforce=true`` per type + to drop matching rollouts before they consume a slot in the batch (e.g. a zero-advantage group + never makes it into a training batch).""" + + post_batch_filters: list[FilterConfig] = [ + GibberishFilterConfig(), + RepetitionFilterConfig(), + ZeroAdvantageFilterConfig(), + ] + """Filters applied *after* a batch has been assembled. Each filter annotates each rollout; + rollouts flagged by an enforcing filter are still recorded but not shipped to the trainer.""" log: LogConfig = LogConfig() @@ -538,6 +590,16 @@ def auto_setup_prime_monitor_run_name(self): self.prime_monitor.run_name = self.wandb.name return self + @model_validator(mode="after") + def validate_unique_filter_types(self): + for slot_name in ("pre_batch_filters", "post_batch_filters"): + types = [f.type for f in getattr(self, slot_name)] + if len(types) != len(set(types)): + raise ValueError( + f"Duplicate filter types in {slot_name}: {types}. Each filter type may only appear once per slot." + ) + return self + @model_validator(mode="after") def inherit_env_algorithms(self): """Envs without their own algorithm inherit the top-level one. diff --git a/src/prime_rl/orchestrator/algo/base.py b/src/prime_rl/orchestrator/algo/base.py index 3d578f8be6..a7962c0327 100644 --- a/src/prime_rl/orchestrator/algo/base.py +++ b/src/prime_rl/orchestrator/algo/base.py @@ -107,11 +107,11 @@ class Algorithm: nothing. - :meth:`score_group` — the cohort of episodes, *before* filtering (filters read the streams): group-relative credit. Default: nothing — - rollouts keep ``advantages=None``, which the drop policy reads as unscored, - not as zero credit. + rollouts keep ``advantages=None``, so advantage-based filters skip them. Model I/O lives in :meth:`score_rollout`: it runs at arrival, *before* the - drop policy, so it pays compute on rollouts that may then be dropped — accepted for the simpler one-rollout-at-a-time shape. + pre-batch filters, so it pays compute on rollouts that may then be filtered + out — accepted for the simpler one-rollout-at-a-time shape. Constructed with the algorithm config it interprets plus the live policy pool (``self.policy_pool`` — always available, never closed by the diff --git a/src/prime_rl/orchestrator/degeneracy.py b/src/prime_rl/orchestrator/degeneracy.py deleted file mode 100644 index 62102f9847..0000000000 --- a/src/prime_rl/orchestrator/degeneracy.py +++ /dev/null @@ -1,84 +0,0 @@ -"""Degeneracy measurements: what a trace's tokens say about how it was generated. - -Each one asks a single question — is it gibberish, is it stuck in a repetition loop — of every -trace, unconditionally, as soon as it is tokenized. They are metrics, reported per agent beside -reward and truncation, so they say something whether or not a run acts on them. - -Acting on one is a separate decision, made once when the batch is assembled -(``prime_rl.orchestrator.train_sink``). Keeping the two apart is what lets every trace be measured -for all of them: a policy that stopped at the first hit would leave the rest unmeasured. -""" - -from __future__ import annotations - -import math -from typing import TYPE_CHECKING - -if TYPE_CHECKING: - from prime_rl.orchestrator.types import TrainRollout - -TOKEN_ID_THRESHOLD = 100_000 -"""Token IDs above this are candidates for gibberish. BPE tokens are sorted by merge order.""" - -LOGPROB_OFFSET = 2.0 -"""Offset from the uniform-distribution logprob: `-log(vocab_size) - LOGPROB_OFFSET`.""" - -REPETITION_WINDOW = 3_000 -"""Consecutive high-probability sampled tokens that count as a repetition loop.""" - -REPETITION_PROB = 0.99 -"""Tokens sampled above this probability count toward the window.""" - - -def is_gibberish(rollout: TrainRollout, vocab_size: int) -> bool: - """Whether the trace generated a rare token at high entropy — a rare BPE id sampled as if - the model had no idea (Section 5.2, https://arxiv.org/abs/2510.02387).""" - threshold = -math.log(vocab_size) - LOGPROB_OFFSET - for branch in rollout.branches: - # branch.{token_ids,logprobs,sampled_mask} are flat and mutually aligned; the raw node - # arrays are not (node.logprobs covers only the sampled suffix, not the generation-prompt - # scaffold that token_ids/mask also span). - for token_id, logprob, sampled in zip(branch.token_ids, branch.logprobs, branch.sampled_mask): - if sampled and token_id > TOKEN_ID_THRESHOLD and logprob < threshold: - return True - return False - - -def is_repetitive(rollout: TrainRollout) -> bool: - """Whether the trace held very high confidence for a long stretch — the signature of a - repetition loop (Section 3.2, https://arxiv.org/abs/2506.13585).""" - threshold = math.log(REPETITION_PROB) - for branch in rollout.branches: - # Aligned branch streams (see `is_gibberish`), and reset the streak per branch: flat - # rollout.nodes interleaves distinct root->leaf paths (compaction/subagents), so a per-node - # walk would run a streak across a branch boundary. - consecutive = 0 - for logprob, sampled in zip(branch.logprobs, branch.sampled_mask): - if not sampled: - continue - consecutive = consecutive + 1 if logprob > threshold else 0 - if consecutive >= REPETITION_WINDOW: - return True - return False - - -def measure(rollout: TrainRollout, vocab_size: int) -> None: - """Record every measurement on the trace. All of them, every time — a rate is only a rate if - nothing decided in advance which traces to look at.""" - rollout.degeneracy = { - "gibberish": is_gibberish(rollout, vocab_size), - "repetition": is_repetitive(rollout), - } - - -def drop_reasons(rollout: TrainRollout, *, drop: list[str], drop_zero_advantage: bool) -> list[str]: - """Why this rollout should not be trained on, if anything. - - A measurement only drops when the run asked it to — measuring is unconditional, acting is not. - Zero credit drops on its own: a scored rollout whose every token is worth nothing produces no - gradient, so the forward pass is wasted. A rollout that was never scored is *not* zero-credit — - opd/opsd train through reference KL and assign no advantages at all, so they must survive.""" - reasons = [name for name in drop if rollout.degeneracy.get(name)] - if drop_zero_advantage and rollout.advantages is not None and not rollout.is_trainable: - reasons.append("zero_advantage") - return reasons diff --git a/src/prime_rl/orchestrator/filters.py b/src/prime_rl/orchestrator/filters.py new file mode 100644 index 0000000000..ad023fd928 --- /dev/null +++ b/src/prime_rl/orchestrator/filters.py @@ -0,0 +1,172 @@ +"""Orchestrator-side rollout filters for detecting degenerate generations. + +Filters run after rollouts complete, inspecting token IDs and logprobs to +detect gibberish or repetition. Detection metrics are always tracked. +When enforce=True, detected rollouts are skipped entirely during training and +are not sent to the trainer. Reward is kept as-is for baseline calculation. +""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import TYPE_CHECKING, Protocol + +from prime_rl.configs.orchestrator import FilterConfig +from prime_rl.utils.logger import get_logger + +if TYPE_CHECKING: + from prime_rl.orchestrator.types import Rollout + + +@dataclass +class FilterResult: + detected: bool + + +class RolloutFilter(Protocol): + name: str + enforce: bool + + def check(self, rollout: Rollout) -> FilterResult: ... + + +@dataclass +class GibberishFilter: + """Flags rollouts containing rare tokens generated at high entropy. + + A token is flagged when both: + - id(token) > token_id_threshold (rare BPE token) + - logprob(token) < -log(vocab_size) - logprob_offset (high entropy) + + References: + Section 5.2, https://arxiv.org/abs/2510.02387 + """ + + name: str + token_id_threshold: int + logprob_threshold: float + enforce: bool = False + + def check(self, rollout: Rollout) -> FilterResult: + for branch in rollout.branches: + # branch.{token_ids,logprobs,sampled_mask} are flat and mutually aligned; the raw + # node arrays are not (node.logprobs covers only the sampled suffix, not the + # generation-prompt scaffold that token_ids/mask also span). + for token_id, logprob, sampled in zip(branch.token_ids, branch.logprobs, branch.sampled_mask): + if not sampled: + continue + if token_id > self.token_id_threshold and logprob < self.logprob_threshold: + return FilterResult(detected=True) + return FilterResult(detected=False) + + +@dataclass +class RepetitionFilter: + """Flags rollouts with pathological repetition loops. + + Counts consecutive tokens where logprob > log(prob_threshold), indicating + the model is generating with very high confidence. When the streak reaches + the window size, the rollout is flagged. + + References: + Section 3.2, https://arxiv.org/abs/2506.13585 + """ + + name: str + window: int + logprob_threshold: float + enforce: bool = False + + def check(self, rollout: Rollout) -> FilterResult: + for branch in rollout.branches: + # Aligned branch streams (see GibberishFilter), and reset the streak per branch: + # flat rollout.nodes interleaves distinct root->leaf paths (compaction/subagents), + # so a per-node walk would run a streak across a branch boundary. + consecutive = 0 + for logprob, sampled in zip(branch.logprobs, branch.sampled_mask): + if not sampled: + continue + if logprob > self.logprob_threshold: + consecutive += 1 + else: + consecutive = 0 + if consecutive >= self.window: + return FilterResult(detected=True) + return FilterResult(detected=False) + + +@dataclass +class ZeroAdvantageFilter: + """Flags rollouts whose advantage stream is all zero (e.g. all rollouts in + a GRPO group earned the same reward, so the centered advantage collapses).""" + + name: str + enforce: bool = True + + def check(self, rollout: Rollout) -> FilterResult: + if rollout.advantages is not None and all(a == 0.0 for a in rollout.advantages): + return FilterResult(detected=True) + return FilterResult(detected=False) + + +def setup_filter(config: FilterConfig, vocab_size: int) -> RolloutFilter: + """Create a RolloutFilter from a filter config.""" + if config.type == "gibberish": + return GibberishFilter( + name="gibberish", + token_id_threshold=config.token_id_threshold, + logprob_threshold=-math.log(vocab_size) - config.logprob_offset, + enforce=config.enforce, + ) + elif config.type == "repetition": + return RepetitionFilter( + name="repetition", + window=config.window, + logprob_threshold=math.log(config.prob_threshold), + enforce=config.enforce, + ) + elif config.type == "zero_advantage": + return ZeroAdvantageFilter( + name="zero_advantage", + enforce=config.enforce, + ) + raise ValueError(f"Unknown filter type: {config.type}") + + +def setup_filters(configs: list[FilterConfig], vocab_size: int, *, kind: str) -> list[RolloutFilter]: + """Create RolloutFilters from a list of filter configs.""" + filters = [setup_filter(config, vocab_size) for config in configs] + if filters: + get_logger().info(f"Configured {len(filters)} {kind} rollout filter(s):") + for config, filt in zip(configs, filters): + mode = "Enforcing" if filt.enforce else "Monitoring" + params = ", ".join(f"{k}={v}" for k, v in config.model_dump().items()) + get_logger().info(f" {mode} {filt.name} filter ({params})") + return filters + + +def apply_filters(filters: list[RolloutFilter], rollouts: list[Rollout]) -> None: + """Flag ``Rollout``\\ s in place with per-filter detection + drop decision. + + Each rollout's ``filter_results`` dict records per-filter detection bools; + ``is_filtered`` is True iff an enforcing filter detected it. First matching + filter wins per rollout (no double-counting). Reward and trajectory tokens + are left untouched so the rollout can still contribute to baseline + calculations and metric aggregation. + """ + for rollout in rollouts: + rollout.filter_results = {f.name: False for f in filters} + rollout.is_filtered = False + + if not filters: + return + + for rollout in rollouts: + for filt in filters: + result = filt.check(rollout) + if result.detected: + rollout.filter_results[filt.name] = True + if filt.enforce: + rollout.is_filtered = True + break diff --git a/src/prime_rl/orchestrator/metrics.py b/src/prime_rl/orchestrator/metrics.py index b71c3e5603..c4f1eca34c 100644 --- a/src/prime_rl/orchestrator/metrics.py +++ b/src/prime_rl/orchestrator/metrics.py @@ -344,16 +344,16 @@ def reward(self) -> Stat: def to_wandb(self, *, prefix: str, subset: Subset) -> dict[str, float]: out = super().to_wandb(prefix=prefix, subset=subset) - # Degeneracy is measured on every trace; the drop verdict is only ever reached for - # trainable survivors (an untrainable seat is 0.0 throughout). Both read per agent. + # The pipeline verdicts are per-trace (an untrainable seat is 0.0 throughout, and filters + # only ever run on trainable survivors), so they read per agent like the rest. for agent, traces in self.by_agent().items(): p = f"{prefix}/{subset}/{agent}" rollouts = traces.rollouts out[f"{p}/is_trainable/mean"] = sum(float(r.is_trainable) for r in rollouts) / len(rollouts) out[f"{p}/is_filtered/mean"] = sum(float(r.is_filtered) for r in rollouts) / len(rollouts) - names = sorted({name for r in rollouts for name in r.degeneracy}) + names = sorted({name for r in rollouts for name in r.filter_results}) out |= { - f"{p}/detected/{name}/mean": sum(1 for r in rollouts if r.degeneracy.get(name)) / len(rollouts) + f"{p}/filters/{name}/mean": sum(1 for r in rollouts if r.filter_results.get(name)) / len(rollouts) for name in names } return out diff --git a/src/prime_rl/orchestrator/orchestrator.py b/src/prime_rl/orchestrator/orchestrator.py index 60a18e6b8b..6db9e0853f 100644 --- a/src/prime_rl/orchestrator/orchestrator.py +++ b/src/prime_rl/orchestrator/orchestrator.py @@ -5,7 +5,7 @@ - ``RolloutDispatcher`` schedules rollouts; emits ``Episode`` (train/eval discriminated by ``kind``) on its queue — the env's own episode, or a traceless one saying why there is none. -- ``TrainSink`` ingests train rollouts (tokenize → measure → advantages → drop policy) +- ``TrainSink`` ingests train rollouts (tokenize → advantages → filters) and returns a ``TrainBatch`` when the threshold is met. - ``EvalSink`` ingests eval rollouts and returns an ``EvalBatch`` (the full returned cohort) on epoch completion. @@ -48,6 +48,7 @@ from prime_rl.orchestrator.envs import EvalEnvs, TrainEnvs from prime_rl.orchestrator.eval_sink import EvalSink from prime_rl.orchestrator.eval_source import EvalSource +from prime_rl.orchestrator.filters import setup_filters from prime_rl.orchestrator.inference_metrics import InferenceMetricsCollector from prime_rl.orchestrator.patches import ( monkey_patch_chat_completion_logprobs, @@ -99,8 +100,9 @@ # shutdown wedges (env-server ZMQ recv, vLLM admin aclose, etc) SHUTDOWN_TIMEOUT_S = 300 -# Abort after this many consecutive train batches drop every rollout — usually an over-eager -# drop policy or a homogeneous-reward dataset; fail loudly instead of spinning +# Abort after this many consecutive train batches drop all rollouts to +# post-batch filters — usually a misconfigured filter or homogeneous-reward +# dataset; fail loudly instead of spinning MAX_CONSECUTIVE_EMPTY_BATCHES = 10 # Maximum batches the orchestrator may run ahead of the trainer. The @@ -257,6 +259,8 @@ async def setup(self) -> None: self.usage_reporter = UsageReporter() # Filters apply to train rollouts only + pre_filters = setup_filters(config.pre_batch_filters, vocab_size=self.tokenizer.vocab_size, kind="pre-batch") + post_filters = setup_filters(config.post_batch_filters, vocab_size=self.tokenizer.vocab_size, kind="post-batch") get_logger().info("Loading training environments") self.train_envs = TrainEnvs( @@ -413,9 +417,8 @@ async def setup(self) -> None: mm_token_type_ids_mapping=self.mm_token_type_ids_mapping, batch_size=config.batch_size, token_batch_size=config.token_batch_size, - vocab_size=self.tokenizer.vocab_size, - drop_degenerate=config.drop_degenerate, - drop_zero_advantage=config.drop_zero_advantage, + pre_filters=pre_filters, + post_filters=post_filters, ) self.eval_sink = EvalSink(eval_envs=self.eval_envs) if self.eval_envs is not None else None self.watcher = WeightWatcher( @@ -585,7 +588,7 @@ async def finalize_train_batch(self, batch: TrainBatch) -> None: if self.consecutive_empty_batches >= MAX_CONSECUTIVE_EMPTY_BATCHES: raise RuntimeError( f"{self.consecutive_empty_batches} consecutive empty train batches — " - "check the drop policy (drop_degenerate / drop_zero_advantage) or task difficulty." + "check filter config (pre_batch_filters / post_batch_filters) or task difficulty." ) return self.consecutive_empty_batches = 0 @@ -665,9 +668,11 @@ async def finalize_train_batch(self, batch: TrainBatch) -> None: for env_name, env_pool in batch.rollouts.by_env().items(): metrics[f"batch/{env_name}"] = len(env_pool) / len(batch.rollouts) if self.train_sink.pre_filter_seen > 0: - metrics["dropped/all/rate"] = self.train_sink.pre_filter_dropped / self.train_sink.pre_filter_seen + metrics["pre_filters/all/dropped_rate"] = ( + self.train_sink.pre_filter_dropped / self.train_sink.pre_filter_seen + ) for name, count in self.train_sink.pre_filter_dropped_by_name.items(): - metrics[f"dropped/all/{name}/rate"] = count / self.train_sink.pre_filter_seen + metrics[f"pre_filters/all/{name}/rate"] = count / self.train_sink.pre_filter_seen self.monitor.log(metrics, step=step) self.wait_for_policy_time = 0.0 self.monitor.log_samples(effective.episodes, step=step) diff --git a/src/prime_rl/orchestrator/train_sink.py b/src/prime_rl/orchestrator/train_sink.py index 9feb1934ef..b340404a7b 100644 --- a/src/prime_rl/orchestrator/train_sink.py +++ b/src/prime_rl/orchestrator/train_sink.py @@ -1,15 +1,14 @@ """TrainSink: three-level rollout sink for the training side. 1. ``process_rollout`` — eager per-rollout tokenization (overlaps with - dispatcher producing more rollouts), the degeneracy measurements, then the - env algorithm's ``finalize_rollout`` (rollout-local scoring + any reference + dispatcher producing more rollouts), then the env algorithm's + ``finalize_rollout`` (rollout-local scoring + any reference I/O). Errored and untrainable rollouts skip this. 2. ``process_group`` — filters errored rollouts, hands the episodes narrowed to their trainable survivors to the env algorithm's ``finalize_group`` - (advantages + per-sample wire stamping), then applies the drop policy — - the one place a rollout is kept out of training. -3. ``process_batch`` — pops a cohort and flattens it into the trainer-bound - ``TrainingSample`` list. Returns a ``TrainBatch``. + (advantages + per-sample wire stamping), runs the pre-batch filter pass. +3. ``process_batch`` — applies post-batch filter annotations and assembles the + trainer-bound ``TrainingSample`` list. Returns a ``TrainBatch``. ``add()`` takes one ``Episode`` and returns ``TrainBatch | None``; group accounting counts episodes, never loose traces. @@ -23,8 +22,8 @@ from collections import defaultdict from prime_rl.configs.orchestrator import OrchestratorConfig -from prime_rl.orchestrator.degeneracy import drop_reasons, measure from prime_rl.orchestrator.envs import TrainEnvs +from prime_rl.orchestrator.filters import RolloutFilter, apply_filters from prime_rl.orchestrator.metrics import TrainRollouts from prime_rl.orchestrator.trajectories import trace_to_samples from prime_rl.orchestrator.types import ( @@ -66,9 +65,8 @@ def __init__( mm_token_type_ids_mapping: dict[int, int] | None, batch_size: int | None, token_batch_size: int | None, - vocab_size: int, - drop_degenerate: list[str], - drop_zero_advantage: bool, + pre_filters: list[RolloutFilter], + post_filters: list[RolloutFilter], ) -> None: assert (batch_size is None) != (token_batch_size is None), ( "Exactly one of batch_size / token_batch_size must be set" @@ -79,9 +77,8 @@ def __init__( self.mm_token_type_ids_mapping = mm_token_type_ids_mapping self.batch_size = batch_size self.token_batch_size = token_batch_size - self.vocab_size = vocab_size - self.drop_degenerate = drop_degenerate - self.drop_zero_advantage = drop_zero_advantage + self.pre_filters = pre_filters + self.post_filters = post_filters # Observation window for the next shipped batch: rollouts of groups # finalized since the last ship (errored + filtered + survivors). @@ -176,7 +173,6 @@ async def process_rollout(self, rollout: TrainRollout) -> None: mm_token_type_ids_mapping=self.mm_token_type_ids_mapping, ) rollout.samples = samples or [] - measure(rollout, self.vocab_size) # Arrival phase: rollout-local scoring (raw reward, echo observation # weighting, opd/opsd reference logprobs) runs as soon as the rollout is # tokenized — before its group is complete. @@ -185,7 +181,7 @@ async def process_rollout(self, rollout: TrainRollout) -> None: async def process_group(self, group_id: str) -> None: """Finalize one GRPO group: drop errored rollouts (the whole group when ``requires_group_scoring`` and any failed), assign advantages, - apply the drop policy, append what survives to ``pending_batch``.""" + run pre-batch filters, append survivors to ``pending_batch``.""" episodes = self.pending_groups.pop(group_id, []) self.pending_group_episodes.pop(group_id, None) if not episodes: @@ -235,32 +231,36 @@ async def process_group(self, group_id: str) -> None: for sample in r.samples: sample.temperatures = [temperature] * len(sample.token_ids) - # Credit is assigned, so the drop decision can be made now: every detection was already - # measured at tokenization, and zero credit is only knowable after the group scored. - dropped_by_reason: dict[str, int] = {} - num_dropped = 0 + if self.pre_filters: + apply_filters(self.pre_filters, survivors) + filtered_by_name: dict[str, int] = {} + num_filtered = 0 for r in survivors: self.pre_filter_seen += 1 - reasons = drop_reasons(r, drop=self.drop_degenerate, drop_zero_advantage=self.drop_zero_advantage) - r.is_filtered = bool(reasons) - if reasons: + if r.is_filtered: self.pre_filter_dropped += 1 - num_dropped += 1 - for reason in reasons: - self.pre_filter_dropped_by_name[reason] = self.pre_filter_dropped_by_name.get(reason, 0) + 1 - dropped_by_reason[reason] = dropped_by_reason.get(reason, 0) + 1 + num_filtered += 1 + for name, hit in r.filter_results.items(): + if hit: + self.pre_filter_dropped_by_name[name] = self.pre_filter_dropped_by_name.get(name, 0) + 1 + filtered_by_name[name] = filtered_by_name.get(name, 0) + 1 continue + # Reset annotations so the post-batch filter pass starts clean + r.filter_results = {} + r.is_filtered = False self.pending_batch.append(r) if self.token_batch_size is not None: self.pending_tokens += payload_tokens(r) + # Per-group summary. One line per finalized group; per-filter + # detection breakdown lives at debug level in ``apply_filters`` rewards = [r.reward for r in survivors] avg_reward = sum(rewards) / len(rewards) if rewards else 0.0 - drop_str = ", ".join(f"{n}={c}" for n, c in dropped_by_reason.items()) if dropped_by_reason else "—" + filter_str = ", ".join(f"{n}={c}" for n, c in filtered_by_name.items()) if filtered_by_name else "—" get_logger().debug( f"Finished group | env={env_name} task_idx={task_idx} | " - f"rollouts={len(group)} (errored={num_errored}, dropped={num_dropped}) | " - f"reward={avg_reward:.4f} | dropped: {drop_str}" + f"rollouts={len(group)} (errored={num_errored}, filtered={num_filtered}) | " + f"reward={avg_reward:.4f} | filters: {filter_str}" ) def process_batch(self) -> TrainBatch: @@ -283,10 +283,12 @@ def process_batch(self) -> TrainBatch: self.pending_batch = self.pending_batch[cut:] self.pending_tokens -= running + if self.post_filters: + apply_filters(self.post_filters, cohort) + # Samples are pre-built by ``process_rollout``; ``process_group`` already stamped the - # advantage stream and loss routing on each sample, and decided what ships. Past this line - # the batch is samples — the episode it came from has served its purpose. - samples: list[TrainingSample] = [sample for r in cohort for sample in r.samples] + # advantage stream and loss routing on each sample. Filtered rollouts don't ship. + samples: list[TrainingSample] = [sample for r in cohort if not r.is_filtered for sample in r.samples] # ``rollouts`` is the observation window — every rollout of every group finalized since the # last ship (errored + filtered + survivors) — while ``samples`` is the shipped cohort's diff --git a/src/prime_rl/orchestrator/types.py b/src/prime_rl/orchestrator/types.py index 0a81081ccf..eadf05fe3f 100644 --- a/src/prime_rl/orchestrator/types.py +++ b/src/prime_rl/orchestrator/types.py @@ -68,9 +68,9 @@ class GroupState: class TrainRollout(vf.Trace[DataT, State, WireAgentConfig], Generic[DataT]): """A rollout on the training path — the one place prime-rl extends a verifiers type, because - trainer-bound state has nowhere else to live: the samples built from its branches, what the - degeneracy measurements found on it, and whether the drop policy kept it. All of it is - ``exclude=True``, so dumping one yields a plain trace on the wire. + trainer-bound state has nowhere else to live: the samples built from its branches and the + filter verdicts over them. All of it is ``exclude=True``, so dumping one yields a plain trace + on the wire. ``env_name`` rides along because a sample is routed by it (the trainer's per-env loss config) and the sink's pending batch is a trace list. Eval rollouts carry the fields unset.""" @@ -79,11 +79,8 @@ class TrainRollout(vf.Trace[DataT, State, WireAgentConfig], Generic[DataT]): env_name: str = Field(default="", exclude=True) samples: list[TrainingSample] = Field(default_factory=list, exclude=True) - degeneracy: dict[str, bool] = Field(default_factory=dict, exclude=True) - """What each degeneracy measurement found on this trace — a measurement, not a verdict.""" is_filtered: bool = Field(default=False, exclude=True) - """The sink's verdict: this rollout is not trained on. Kept for the metrics window, which - reports what came back as well as what shipped.""" + filter_results: dict[str, bool] = Field(default_factory=dict, exclude=True) def assign_advantages(self, value: float) -> None: """Write ``value`` as the credit for every trainable token, node by node. Credit lives on diff --git a/tests/unit/orchestrator/test_degeneracy.py b/tests/unit/orchestrator/test_degeneracy.py deleted file mode 100644 index c21b955b3c..0000000000 --- a/tests/unit/orchestrator/test_degeneracy.py +++ /dev/null @@ -1,172 +0,0 @@ -import math - -import verifiers.v1 as vf -from verifiers.v1.configs.agent import WireAgentConfig - -from prime_rl.orchestrator.degeneracy import ( - REPETITION_PROB, - REPETITION_WINDOW, - TOKEN_ID_THRESHOLD, - drop_reasons, - is_gibberish, - is_repetitive, - measure, -) -from prime_rl.orchestrator.types import TrainRollout - -VOCAB_SIZE = 128_000 -GIBBERISH_LOGPROB = -math.log(VOCAB_SIZE) - 2.0 - 1.0 # comfortably under the entropy threshold -RARE_TOKEN = TOKEN_ID_THRESHOLD + 1 -REPEAT_LOGPROB = math.log(REPETITION_PROB) + 0.001 # just above the confidence threshold - - -def _assistant_node(token_ids: list[int], logprobs: list[float]) -> vf.MessageNode: - """An assistant node whose tokens are all model-sampled (the measurements read each node's - masked-True tokens + logprobs).""" - return vf.MessageNode( - message=vf.AssistantMessage(content="x"), - token_ids=token_ids, - mask=[True] * len(token_ids), - logprobs=logprobs, - ) - - -def _make_rollout(nodes: list[vf.MessageNode]) -> TrainRollout: - rollout = TrainRollout[vf.TaskData]( - task=vf.TraceTask(type="Task", data=vf.TaskData(idx=0, prompt="")), - agent=vf.AgentInfo(config=WireAgentConfig()), - nodes=nodes, - rewards={"reward": vf.Reward(score=1.0)}, - ) - rollout.env_name = "test" - return rollout - - -def _sampled(token_ids: list[int], logprobs: list[float], *, multi_step: bool = False) -> TrainRollout: - """A rollout carrying these completion tokens, optionally split across two model turns.""" - if not multi_step: - return _make_rollout([_assistant_node(token_ids, logprobs)]) - mid = len(token_ids) // 2 - return _make_rollout( - [ - _assistant_node(token_ids[:mid], logprobs[:mid]), - _assistant_node(token_ids[mid:], logprobs[mid:]), - ] - ) - - -# --- gibberish --- - - -def test_gibberish_detects_rare_low_prob_token(): - assert is_gibberish(_sampled([50, RARE_TOKEN, 80], [-1.0, GIBBERISH_LOGPROB, -0.5]), VOCAB_SIZE) - - -def test_gibberish_ignores_normal_tokens(): - assert not is_gibberish(_sampled([10, 200, 5000], [-1.0, -2.0, -3.0]), VOCAB_SIZE) - - -def test_gibberish_ignores_high_prob_rare_token(): - assert not is_gibberish(_sampled([RARE_TOKEN], [-0.5]), VOCAB_SIZE) - - -def test_gibberish_works_across_trajectory_steps(): - rollout = _sampled([50, 60, RARE_TOKEN, 80], [-1.0, -0.5, GIBBERISH_LOGPROB, -0.5], multi_step=True) - assert is_gibberish(rollout, VOCAB_SIZE) - - -def test_gibberish_aligns_logprobs_under_generation_prompt_scaffold(): - """Regression: the assistant node carries a generation-prompt scaffold (mask=False) and - suffix-only logprobs, and the gibberish token is the LAST completion token. A per-node - ``zip(token_ids, logprobs, mask)`` truncates at len(logprobs) and never examines it; reading - the aligned branch streams finds it.""" - node = vf.MessageNode( - message=vf.AssistantMessage(content="x"), - token_ids=[1, 1, 50, 80, RARE_TOKEN], - mask=[False, False, True, True, True], - logprobs=[-1.0, -0.5, GIBBERISH_LOGPROB], # the sampled suffix only, as vLLM returns it - ) - assert is_gibberish(_make_rollout([node]), VOCAB_SIZE) - - -# --- repetition --- - - -def test_repetition_triggers_at_the_window(): - n = REPETITION_WINDOW - assert is_repetitive(_sampled(list(range(n)), [REPEAT_LOGPROB] * n)) - - -def test_repetition_no_trigger_below_the_window(): - n = REPETITION_WINDOW - 1 - assert not is_repetitive(_sampled(list(range(n)), [REPEAT_LOGPROB] * n)) - - -def test_repetition_resets_on_a_low_probability_token(): - """The streak has to be consecutive: one unconfident token in the middle breaks it, so nearly - twice the window's worth of confident tokens either side of it is not a loop.""" - half = [REPEAT_LOGPROB] * (REPETITION_WINDOW - 1) - logprobs = [*half, -2.0, *half] - assert not is_repetitive(_sampled(list(range(len(logprobs))), logprobs)) - - -def test_repetition_does_not_run_a_streak_across_branches(): - """Each branch is measured on its own — a flat walk over ``nodes`` would join two turns' - streaks into one that never happened.""" - half = [REPEAT_LOGPROB] * (REPETITION_WINDOW - 1) - logprobs = [*half, *half] - assert not is_repetitive(_sampled(list(range(len(logprobs))), logprobs, multi_step=True)) - - -# --- measuring is unconditional --- - - -def test_measure_records_every_measurement(): - """Every trace is measured for all of them, so a rate is a rate: one hit does not shadow the - others, and nothing decided in advance which traces to look at.""" - rollout = _sampled([RARE_TOKEN] * 5, [GIBBERISH_LOGPROB] * 5) - measure(rollout, VOCAB_SIZE) - assert rollout.degeneracy == {"gibberish": True, "repetition": False} - - -def test_measure_handles_a_trace_with_no_sampled_tokens(): - rollout = _make_rollout([]) - measure(rollout, VOCAB_SIZE) - assert rollout.degeneracy == {"gibberish": False, "repetition": False} - - -# --- the drop policy --- - - -def _drop(rollout, drop=(), zero_advantage=True): - return drop_reasons(rollout, drop=list(drop), drop_zero_advantage=zero_advantage) - - -def test_a_measurement_only_drops_when_asked(): - rollout = _sampled([1], [-1.0]) - rollout.degeneracy = {"gibberish": True} - rollout.assign_advantages(0.5) - assert _drop(rollout) == [] # measured, but nothing asked for it to drop - assert _drop(rollout, drop=["gibberish"]) == ["gibberish"] - - -def test_zero_credit_drops_by_default(): - rollout = _sampled([1], [-1.0]) - rollout.assign_advantages(0.0) - assert _drop(rollout) == ["zero_advantage"] - assert _drop(rollout, zero_advantage=False) == [] - - -def test_unscored_rollout_is_not_zero_credit(): - """opd/opsd assign no credit at all and train through reference KL — dropping them as - zero-advantage would ship an empty batch for every distillation run.""" - rollout = _sampled([1], [-1.0]) - assert rollout.advantages is None - assert _drop(rollout) == [] - - -def test_reasons_accumulate(): - rollout = _sampled([1], [-1.0]) - rollout.degeneracy = {"gibberish": True, "repetition": True} - rollout.assign_advantages(0.0) - assert _drop(rollout, drop=["gibberish", "repetition"]) == ["gibberish", "repetition", "zero_advantage"] diff --git a/tests/unit/orchestrator/test_filters.py b/tests/unit/orchestrator/test_filters.py new file mode 100644 index 0000000000..914d753793 --- /dev/null +++ b/tests/unit/orchestrator/test_filters.py @@ -0,0 +1,407 @@ +import math + +import verifiers.v1 as vf +from verifiers.v1.configs.agent import WireAgentConfig + +from prime_rl.configs.orchestrator import GibberishFilterConfig, RepetitionFilterConfig +from prime_rl.orchestrator.filters import ( + GibberishFilter, + RepetitionFilter, + apply_filters, + setup_filter, + setup_filters, +) +from prime_rl.orchestrator.types import TrainRollout + + +def _assistant_node(token_ids: list[int], logprobs: list[float]) -> vf.MessageNode: + """An assistant node whose tokens are all model-sampled (the filters read each node's + masked-True tokens + logprobs).""" + return vf.MessageNode( + message=vf.AssistantMessage(content="x"), + token_ids=token_ids, + mask=[True] * len(token_ids), + logprobs=logprobs, + ) + + +def _scaffold_assistant_node( + completion_ids: list[int], completion_logprobs: list[float], *, scaffold: int = 2 +) -> vf.MessageNode: + """A realistic v1 assistant node: a leading generation-prompt scaffold (mask=False, not + model-sampled) then the sampled completion. ``logprobs`` cover only the completion suffix + (vLLM returns logprobs for generated tokens only) — the exact layout where per-node + ``zip(token_ids, logprobs, mask)`` mispairs and the branch streams normalize.""" + return vf.MessageNode( + message=vf.AssistantMessage(content="x"), + token_ids=[1] * scaffold + completion_ids, + mask=[False] * scaffold + [True] * len(completion_ids), + logprobs=completion_logprobs, + ) + + +def _make_rollout( + completion_ids: list[int], + completion_logprobs: list[float], + *, + reward: float = 1.0, + multi_step: bool = False, +) -> TrainRollout: + """Build a ``TrainRollout`` (a message-graph trace) carrying the completion tokens — enough for + the filters to inspect each node's sampled tokens / logprobs.""" + if multi_step: + mid = len(completion_ids) // 2 + nodes = [ + _assistant_node(completion_ids[:mid], completion_logprobs[:mid]), + _assistant_node(completion_ids[mid:], completion_logprobs[mid:]), + ] + else: + nodes = [_assistant_node(completion_ids, completion_logprobs)] + rollout = TrainRollout[vf.TaskData]( + task=vf.TraceTask(type="Task", data=vf.TaskData(idx=0, prompt="")), + agent=vf.AgentInfo(config=WireAgentConfig()), + nodes=nodes, + rewards={"reward": vf.Reward(score=reward)}, + ) + rollout.env_name = "test" + return rollout + + +def _make_gibberish_filter(vocab_size=128_000, token_id_threshold=100_000, logprob_offset=2.0, enforce=False): + logprob_threshold = -math.log(vocab_size) - logprob_offset + return GibberishFilter( + name="gibberish", token_id_threshold=token_id_threshold, logprob_threshold=logprob_threshold, enforce=enforce + ) + + +def _make_repetition_filter(window=5, prob_threshold=0.99, enforce=False): + return RepetitionFilter( + name="repetition", window=window, logprob_threshold=math.log(prob_threshold), enforce=enforce + ) + + +# --- GibberishFilter tests --- + + +def test_gibberish_detects_rare_low_prob_token(): + gibberish_filter = _make_gibberish_filter() + + result = gibberish_filter.check( + _make_rollout( + completion_ids=[50, 120_000, 80], + completion_logprobs=[-1.0, gibberish_filter.logprob_threshold - 1.0, -0.5], + ) + ) + assert result.detected is True + + +def test_gibberish_ignores_normal_tokens(): + gibberish_filter = _make_gibberish_filter() + + result = gibberish_filter.check( + _make_rollout( + completion_ids=[10, 200, 5000], + completion_logprobs=[-1.0, -2.0, -3.0], + ) + ) + assert result.detected is False + + +def test_gibberish_ignores_high_prob_rare_token(): + gibberish_filter = _make_gibberish_filter() + + result = gibberish_filter.check( + _make_rollout( + completion_ids=[120_000], + completion_logprobs=[-0.5], + ) + ) + assert result.detected is False + + +def test_gibberish_works_across_trajectory_steps(): + gibberish_filter = _make_gibberish_filter() + + result = gibberish_filter.check( + _make_rollout( + completion_ids=[50, 60, 120_000, 80], + completion_logprobs=[-1.0, -0.5, gibberish_filter.logprob_threshold - 1.0, -0.5], + multi_step=True, + ) + ) + assert result.detected is True + + +def test_gibberish_aligns_logprobs_under_generation_prompt_scaffold(): + """Regression: the assistant node carries a generation-prompt scaffold (mask=False) and + suffix-only logprobs, and the gibberish token is the LAST completion token. The old + per-node ``zip(token_ids, logprobs, mask)`` truncated at len(logprobs) and never examined + it; reading the aligned branch streams detects it.""" + gibberish_filter = _make_gibberish_filter() + + rollout = TrainRollout[vf.TaskData]( + task=vf.TraceTask(type="Task", data=vf.TaskData(idx=0, prompt="")), + agent=vf.AgentInfo(config=WireAgentConfig()), + nodes=[_scaffold_assistant_node([50, 80, 120_000], [-1.0, -0.5, gibberish_filter.logprob_threshold - 1.0])], + rewards={"reward": vf.Reward(score=1.0)}, + ) + + result = gibberish_filter.check(rollout) + assert result.detected is True + + +# --- RepetitionFilter tests --- + + +def test_repetition_triggers_after_window(): + repetition_filter = _make_repetition_filter(window=5) + + result = repetition_filter.check( + _make_rollout( + completion_ids=list(range(5)), + completion_logprobs=[-0.001] * 5, + ) + ) + assert result.detected is True + + +def test_repetition_no_trigger_below_window(): + repetition_filter = _make_repetition_filter(window=5) + + result = repetition_filter.check( + _make_rollout( + completion_ids=list(range(4)), + completion_logprobs=[-0.001] * 4, + ) + ) + assert result.detected is False + + +def test_repetition_resets_on_low_prob(): + repetition_filter = _make_repetition_filter(window=5) + + logprobs = [-0.001] * 3 + [-2.0] + [-0.001] * 3 + result = repetition_filter.check( + _make_rollout( + completion_ids=list(range(7)), + completion_logprobs=logprobs, + ) + ) + assert result.detected is False + + +def test_repetition_varied_probs_no_trigger(): + repetition_filter = _make_repetition_filter(window=3) + + result = repetition_filter.check( + _make_rollout( + completion_ids=list(range(6)), + completion_logprobs=[-0.001, -3.0, -0.001, -3.0, -0.001, -3.0], + ) + ) + assert result.detected is False + + +# --- setup_filter / setup_filters tests --- + + +def test_setup_filter_gibberish(): + config = GibberishFilterConfig(token_id_threshold=100_000, logprob_offset=2.0) + gibberish_filter = setup_filter(config, vocab_size=128_000) + assert isinstance(gibberish_filter, GibberishFilter) + assert gibberish_filter.name == "gibberish" + assert gibberish_filter.token_id_threshold == 100_000 + assert abs(gibberish_filter.logprob_threshold - (-math.log(128_000) - 2.0)) < 1e-10 + assert gibberish_filter.enforce is False + + +def test_setup_filter_gibberish_enforce(): + config = GibberishFilterConfig(enforce=True) + gibberish_filter = setup_filter(config, vocab_size=128_000) + assert gibberish_filter.enforce is True + + +def test_setup_filter_repetition(): + config = RepetitionFilterConfig(window=3_000, prob_threshold=0.99) + repetition_filter = setup_filter(config, vocab_size=128_000) + assert isinstance(repetition_filter, RepetitionFilter) + assert repetition_filter.name == "repetition" + assert repetition_filter.window == 3_000 + assert abs(repetition_filter.logprob_threshold - math.log(0.99)) < 1e-10 + assert repetition_filter.enforce is False + + +def test_setup_filter_repetition_enforce(): + config = RepetitionFilterConfig(enforce=True) + repetition_filter = setup_filter(config, vocab_size=128_000) + assert repetition_filter.enforce is True + + +def test_setup_filters_multiple(): + configs = [ + GibberishFilterConfig(), + RepetitionFilterConfig(), + ] + filters = setup_filters(configs, vocab_size=128_000, kind="post-batch") + assert len(filters) == 2 + assert filters[0].name == "gibberish" + assert filters[1].name == "repetition" + + +# --- apply_filters tests (enforce=True) --- + + +def test_apply_filters_enforced_flags_rollout(): + gibberish_filter = _make_gibberish_filter(enforce=True) + + rollout = _make_rollout( + completion_ids=[120_000], + completion_logprobs=[gibberish_filter.logprob_threshold - 1.0], + reward=1.0, + ) + + apply_filters([gibberish_filter], [rollout]) + + assert rollout.reward == 1.0 + assert rollout.nodes[0].token_ids == [120_000] + assert rollout.nodes[0].mask == [True] + assert rollout.stop_condition is None + assert rollout.filter_results == {"gibberish": True} + assert rollout.is_filtered is True + + +def test_apply_filters_preserves_clean_rollouts(): + gibberish_filter = _make_gibberish_filter(enforce=True) + + rollout = _make_rollout( + completion_ids=[50, 60, 70], + completion_logprobs=[-1.0, -2.0, -1.5], + reward=1.0, + ) + + apply_filters([gibberish_filter], [rollout]) + + assert rollout.reward == 1.0 + assert rollout.nodes[0].token_ids == [50, 60, 70] + assert all(rollout.nodes[0].mask) + assert rollout.stop_condition is None + assert rollout.filter_results == {"gibberish": False} + assert rollout.is_filtered is False + + +def test_apply_filters_first_filter_wins(): + gibberish_filter = _make_gibberish_filter(enforce=True) + repetition_filter = _make_repetition_filter(window=2, enforce=True) + + rollout = _make_rollout( + completion_ids=[120_000, 1, 2], + completion_logprobs=[gibberish_filter.logprob_threshold - 1.0, -0.001, -0.001], + reward=1.0, + ) + + apply_filters([gibberish_filter, repetition_filter], [rollout]) + + assert rollout.stop_condition is None + assert rollout.filter_results == {"gibberish": True, "repetition": False} + assert rollout.is_filtered is True + + +def test_apply_filters_empty_list(): + rollout = _make_rollout( + completion_ids=[1, 2, 3], + completion_logprobs=[-1.0, -1.0, -1.0], + ) + apply_filters([], [rollout]) + assert rollout.filter_results == {} + assert rollout.is_filtered is False + assert rollout.reward == 1.0 + + +def test_apply_filters_mixed_batch(): + gibberish_filter = _make_gibberish_filter(enforce=True) + + clean = _make_rollout(completion_ids=[50], completion_logprobs=[-1.0], reward=1.0) + dirty = _make_rollout( + completion_ids=[120_000], completion_logprobs=[gibberish_filter.logprob_threshold - 1.0], reward=1.0 + ) + + apply_filters([gibberish_filter], [clean, dirty]) + + assert clean.reward == 1.0 + assert dirty.reward == 1.0 + assert clean.is_filtered is False + assert dirty.is_filtered is True + + +def test_apply_filters_enforced_preserves_rollout_tokens(): + gibberish_filter = _make_gibberish_filter(enforce=True) + + rollout = _make_rollout( + completion_ids=[10, 120_000, 30], + completion_logprobs=[-1.0, gibberish_filter.logprob_threshold - 1.0, -0.5], + reward=1.0, + ) + + apply_filters([gibberish_filter], [rollout]) + + assert rollout.nodes[0].token_ids == [10, 120_000, 30] + assert rollout.nodes[0].logprobs == [ + -1.0, + gibberish_filter.logprob_threshold - 1.0, + -0.5, + ] + assert rollout.nodes[0].mask == [True, True, True] + assert rollout.is_filtered is True + + +def test_apply_filters_preserves_existing_stop_condition(): + gibberish_filter = _make_gibberish_filter(enforce=True) + + rollout = _make_rollout( + completion_ids=[120_000], + completion_logprobs=[gibberish_filter.logprob_threshold - 1.0], + reward=1.0, + ) + rollout.stop_condition = "generation_truncated" + + apply_filters([gibberish_filter], [rollout]) + + assert rollout.stop_condition == "generation_truncated" + assert rollout.is_filtered is True + + +# --- apply_filters tests (monitor-only, enforce=False) --- + + +def test_apply_filters_monitor_only_tracks_detection(): + gibberish_filter = _make_gibberish_filter(enforce=False) + + rollout = _make_rollout( + completion_ids=[120_000], + completion_logprobs=[gibberish_filter.logprob_threshold - 1.0], + reward=1.0, + ) + + apply_filters([gibberish_filter], [rollout]) + + assert rollout.reward == 1.0 + assert all(rollout.nodes[0].mask) + assert rollout.stop_condition is None + assert rollout.filter_results == {"gibberish": True} + assert rollout.is_filtered is False + + +def test_apply_filters_monitor_only_mixed_batch(): + gibberish_filter = _make_gibberish_filter(enforce=False) + + clean = _make_rollout(completion_ids=[50], completion_logprobs=[-1.0], reward=1.0) + dirty = _make_rollout( + completion_ids=[120_000], completion_logprobs=[gibberish_filter.logprob_threshold - 1.0], reward=1.0 + ) + + apply_filters([gibberish_filter], [clean, dirty]) + + assert clean.reward == 1.0 + assert dirty.reward == 1.0 + assert clean.is_filtered is False + assert dirty.is_filtered is False diff --git a/tests/unit/orchestrator/test_metrics.py b/tests/unit/orchestrator/test_metrics.py index d30ec2cfb4..6b8ba82db6 100644 --- a/tests/unit/orchestrator/test_metrics.py +++ b/tests/unit/orchestrator/test_metrics.py @@ -45,7 +45,7 @@ def mk( trainable: bool = True, is_trainable: bool = True, is_filtered: bool = False, - degeneracy: dict | None = None, + filter_results: dict | None = None, setup: float = 0.0, agent: float = 0.0, agent_model: float = 0.0, @@ -76,7 +76,7 @@ def mk( agent=SimpleNamespace(trainable=trainable, name=agent_name), is_trainable=is_trainable, is_filtered=is_filtered, - degeneracy=degeneracy or {}, + filter_results=filter_results or {}, timing=SimpleNamespace( setup=SimpleNamespace(duration=setup), agent=SimpleNamespace( @@ -285,16 +285,16 @@ def test_nested_timing(): def test_train_only_metrics_absent_from_eval(): rollouts = [ - mk(is_trainable=True, is_filtered=True, degeneracy={"gibberish": True}), - mk(is_trainable=False, degeneracy={"gibberish": False}), + mk(is_trainable=True, is_filtered=True, filter_results={"gibberish": True}), + mk(is_trainable=False, filter_results={"gibberish": False}), ] out = train_wandb(rollouts) assert out["train/agg/all/agent/is_trainable/mean"] == 0.5 assert out["train/agg/all/agent/is_filtered/mean"] == 0.5 - assert out["train/agg/all/agent/detected/gibberish/mean"] == 0.5 + assert out["train/agg/all/agent/filters/gibberish/mean"] == 0.5 assert "train/agg/all/is_trainable/mean" not in out # pipeline verdicts are per-trace eval_out = EvalRollouts(solo(rollouts)).metrics.to_wandb(prefix="eval/x", subset="all") - assert not any("is_trainable" in k or "is_filtered" in k or "/detected/" in k for k in eval_out) + assert not any("is_trainable" in k or "is_filtered" in k or "/filters/" in k for k in eval_out) def test_eval_avg_at_k_and_pass_k(): @@ -334,12 +334,12 @@ def test_traceless_episode_keeps_its_reason(): def test_training_state_is_train_only(): """Only a train rollout carries trainer-bound state; an eval trace has no field for it. Credit is not among them — it lives on the graph's nodes, which every trace has.""" - assert {"samples", "is_filtered", "degeneracy"} <= set(TrainRollout.model_fields) + assert {"samples", "is_filtered", "filter_results"} <= set(TrainRollout.model_fields) assert "advantages" not in TrainRollout.model_fields # derived from the nodes # An eval trace is the env's own, unextended — where a trace sits is the episode's to say, # and it is the same wire specialization the episode's ``traces`` hold. assert Rollout is vf.WireTrace - assert not {"samples", "is_filtered", "degeneracy", "group_id", "episode_id"} & set(Rollout.model_fields) + assert not {"samples", "is_filtered", "filter_results", "group_id", "episode_id"} & set(Rollout.model_fields) def test_inflight_episode_stamps_what_lands():