diff --git a/configs/ci/nightly-fft/wiki-search.toml b/configs/ci/nightly-fft/wiki-search.toml index c3977c0adb..cb5536c453 100644 --- a/configs/ci/nightly-fft/wiki-search.toml +++ b/configs/ci/nightly-fft/wiki-search.toml @@ -19,9 +19,8 @@ batch_size = 512 group_size = 16 oversampling_factor = 2.0 -[[orchestrator.pre_batch_filters]] -type = "zero_advantage" -enforce = true +[orchestrator.sampler] +drop_degenerate_groups = true [[orchestrator.train.source]] name = "wiki-search" diff --git a/configs/debug/algo/echo.toml b/configs/debug/algo/echo.toml index e9bdddccb3..722f961e59 100644 --- a/configs/debug/algo/echo.toml +++ b/configs/debug/algo/echo.toml @@ -44,12 +44,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 af1984c7ca..988f8d2931 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 detections 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) +- [Detections](#detections) - [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 `RolloutSource` (`p | `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 | 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: -- `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`. +- `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 at arrival, so it pays compute on rollouts a detection may then exclude. +- `score_group(group)` — the finalized cohort, synchronous: group-relative credit (GRPO/MaxRL baselines). `group` is a list of `Rollout`. 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. @@ -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` (they always ship; their signal is not advantage-based) 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` (they always ship) and ship no advantage stream. | | `sft` | `ce` | Cross-entropy on the sampled tokens. Assigns no advantage — trains on every sampled token. | ### Default Advantage @@ -382,7 +382,7 @@ id = "null" type = "subprocess" ``` -`group_size` controls how many problems are proposed from each source task. `env.n` controls how many solvers attempt each proposed problem. If a comparison contains only one trace—for example, a solver when `env.n = 1`—its advantage is zero and the zero-advantage filter removes it. +`group_size` controls how many problems are proposed from each source task. `env.n` controls how many solvers attempt each proposed problem. If a comparison contains only one trace—for example, a solver when `env.n = 1`—its advantage is zero and it ships no samples. This algorithm is accepted only for proposer-solver envs. Use the env's `train_proposer` and `train_solver` settings if you want to train only one role. @@ -439,7 +439,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 (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). -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. +Each per-token list must match the rollout's completion-token count exactly — validated loudly when the view writes it. Shipping and metrics derive from the streams (an all-zero stream ships no samples unless the algorithm trains on zero advantage; 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 @@ -454,30 +454,28 @@ 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 rollout arrival, before the group-time detections, so an excluded rollout still cost its reference compute — accepted for the simpler one-rollout-at-a-time shape. -## Filters +## Detections -Filters drop rollouts between scoring and training. Built-ins (composable): +Detections flag generation pathology — the policy melting down mid-sample — by reading a rollout's own token ids and logprobs. They are rollout-granularity predicates, evaluated once at group finalization (the pipeline's single decision point). Built-ins: -| Filter | Effect | +| Detection | Fires when | |---|---| -| `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. | +| `gibberish` | any sampled token is both rare (`token_id > token_id_threshold`) and high-entropy (`logprob < -log(vocab_size) - logprob_offset`) | +| `repetition` | `window` consecutive sampled tokens each exceed `prob_threshold` — a high-confidence repetition loop | -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: +Both run in monitor mode by default (results are recorded as metrics, nothing is dropped). With `enforce = true`, a detected rollout ships no training samples and never occupies a batch slot — the batch backfills from fresh groups — while its reward still counts toward the group baseline. Setting the list replaces the defaults wholesale: ```toml -[[orchestrator.post_batch_filters]] -type = "zero_advantage" - -[[orchestrator.post_batch_filters]] -type = "repetition" -threshold = 0.4 +[[orchestrator.detections]] +type = "gibberish" +enforce = true ``` -Filtered rollouts still appear in W&B distributions, just not in the trainer batch — useful for spotting whether filtering is doing its job. +Zero-advantage handling needs no configuration: an all-zero advantage stream ships no samples unless the source's algorithm declares `trains_on_zero_advantage` (echo does — its observation CE trains through collapsed advantages), and `[orchestrator.sampler] drop_degenerate_groups = true` additionally keeps whole zero-signal groups out of the batch so it backfills (size the extra inference with `oversampling_factor`). + +Detected and degenerate rollouts still appear in W&B distributions and the `detections/*` / `sampler/*` metrics, just not in the trainer batch — useful for spotting whether the hygiene is doing its job. ## Multi-Turn Trajectories diff --git a/docs/overview.md b/docs/overview.md index b33c8a0736..25375d6300 100644 --- a/docs/overview.md +++ b/docs/overview.md @@ -40,6 +40,6 @@ The `rl` entrypoint reads `examples/basic/reverse-text/rl.toml`, splits it into - **[Training](training.md)** — Launch and observe RL and SFT runs. - **[Inference](inference.md)** — vLLM-backed server (or fleet) holding the current policy. - **[Scaling](scaling.md)** — Single-GPU through multi-node clusters via FSDP / EP / CP and SLURM. -- **[Algorithms](algorithms.md)** — Async semantics, loss / advantage / filter plugins, trajectory merging. +- **[Algorithms](algorithms.md)** — Async semantics, loss / advantage plugins, detections, trajectory merging. - **[Advanced](advanced.md)** — Custom modeling, multimodal, LoRA, P/D inference. - **[Development](development.md)** — Test suite, pre-commit hooks, adding a new model. diff --git a/examples/advanced/glm-5.2/swe.toml b/examples/advanced/glm-5.2/swe.toml index 657a488c9b..a2e429af4f 100644 --- a/examples/advanced/glm-5.2/swe.toml +++ b/examples/advanced/glm-5.2/swe.toml @@ -88,7 +88,7 @@ id = "bash" type = "prime" labels = ["glm5-pd-disag", "swe-bench-verified"] -[[orchestrator.post_batch_filters]] +[[orchestrator.detections]] type = "gibberish" enforce = true diff --git a/examples/advanced/intellect-3.1/rl.toml b/examples/advanced/intellect-3.1/rl.toml index 43f7515568..ef7b64b098 100644 --- a/examples/advanced/intellect-3.1/rl.toml +++ b/examples/advanced/intellect-3.1/rl.toml @@ -119,9 +119,8 @@ id = "null" [orchestrator.train.source.env.agent.runtime] type = "subprocess" -[[orchestrator.pre_batch_filters]] -type = "zero_advantage" -enforce = true +[orchestrator.sampler] +drop_degenerate_groups = true [orchestrator.eval] interval = 25 diff --git a/examples/basic/wiki-search/rl.toml b/examples/basic/wiki-search/rl.toml index f9f256387b..6729dbfde0 100644 --- a/examples/basic/wiki-search/rl.toml +++ b/examples/basic/wiki-search/rl.toml @@ -52,9 +52,8 @@ id = "null" [orchestrator.train.source.env.agent.runtime] type = "subprocess" -[[orchestrator.pre_batch_filters]] -type = "zero_advantage" -enforce = true +[orchestrator.sampler] +drop_degenerate_groups = true [ckpt] # Checkpoint at the end of training 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 74af6c0de8..ae4f05a1cb 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py @@ -303,11 +303,12 @@ class CheckpointConfig(BaseConfig): # Flags rare tokens generated at high entropy (Section 5.2, https://arxiv.org/abs/2510.02387). -class GibberishFilterConfig(BaseConfig): +class GibberishDetectionConfig(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.""" + """When True, detected rollouts ship no training samples (their reward still counts toward + the group baseline). 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.""" @@ -319,11 +320,12 @@ class GibberishFilterConfig(BaseConfig): # 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): +class RepetitionDetectionConfig(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.""" + """When True, detected rollouts ship no training samples (their reward still counts toward + the group baseline). When False, only track detection metrics.""" window: int = Field(3_000, ge=1) """Consecutive high-probability steps required to flag the rollout.""" @@ -332,18 +334,20 @@ 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" +DetectionConfig: TypeAlias = Annotated[ + GibberishDetectionConfig | RepetitionDetectionConfig, + Field(discriminator="type"), +] - enforce: bool = True - """When True, skip detected rollouts entirely so they are not sent to the trainer. When False, only track detection metrics.""" +class SamplerConfig(BaseConfig): + """The task sampler's decision knobs (its stats collection is always on).""" -FilterConfig: TypeAlias = Annotated[ - GibberishFilterConfig | RepetitionFilterConfig | ZeroAdvantageFilterConfig, - Field(discriminator="type"), -] + drop_degenerate_groups: bool = False + """Drop a finalized group that produced no training signal (every advantage stream all-zero) + instead of letting its rollouts occupy batch slots — the batch then backfills from fresh + groups, at inference cost bounded by ``oversampling_factor``. Groups of sources whose + algorithm declares ``trains_on_zero_advantage`` (echo) are never dropped.""" class FileSystemWeightBroadcastConfig(BaseConfig): @@ -416,23 +420,16 @@ 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(), + detections: list[DetectionConfig] = [ + GibberishDetectionConfig(), + RepetitionDetectionConfig(), ] - """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.""" + """Generation-pathology detections, evaluated per rollout at group finalization. Both registered in + monitor mode by default; flip ``enforce=true`` per type to keep detected rollouts out of + training (they never enter the batch, so the batch backfills; their reward still counts + toward the group baseline). Setting this replaces the defaults wholesale.""" + + sampler: SamplerConfig = SamplerConfig() log: LogConfig = LogConfig() @@ -531,13 +528,10 @@ def auto_setup_prime_monitor_run_name(self): 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." - ) + def validate_unique_detection_types(self): + types = [d.type for d in self.detections] + if len(types) != len(set(types)): + raise ValueError(f"Duplicate detection types: {types}. Each detection type may only appear once.") return self @model_validator(mode="after") diff --git a/src/prime_rl/orchestrator/algo/base.py b/src/prime_rl/orchestrator/algo/base.py index 7b47210b6a..45c6d6a4f4 100644 --- a/src/prime_rl/orchestrator/algo/base.py +++ b/src/prime_rl/orchestrator/algo/base.py @@ -119,6 +119,13 @@ class Algorithm: action_loss_type: ClassVar[ActionLossType] = "rl" + trains_on_zero_advantage: ClassVar[bool] = False + """Whether an all-zero advantage stream still carries training signal for + this algorithm. False for credit-assignment algorithms (an all-zero stream + means no gradient — the sink ships no samples for it, and the task + sampler's degenerate-group gate may drop the whole group); echo sets True + because its observation CE trains even when the GRPO advantage collapses.""" + def __init__(self, config: AlgoConfig, policy_pool: InferencePool): self.policy_pool = policy_pool self.connected_pools: list[InferencePool] = [] # frozen pools connected in setup(); closed at shutdown diff --git a/src/prime_rl/orchestrator/algo/echo.py b/src/prime_rl/orchestrator/algo/echo.py index d4ecf74fa3..3cb365af75 100644 --- a/src/prime_rl/orchestrator/algo/echo.py +++ b/src/prime_rl/orchestrator/algo/echo.py @@ -23,6 +23,10 @@ class EchoAlgorithm(GRPOAlgorithm): mask and its denominator. An optional user filter narrows the selection per rollout (e.g. dropping tool-output warnings).""" + # Observation CE trains even when the GRPO advantage collapses to zero — + # zero-advantage rollouts of echo envs stay in the batch and ship. + trains_on_zero_advantage = True + def __init__(self, config: EchoAlgoConfig, policy_pool: InferencePool): super().__init__(config, policy_pool) self.role_weights = { diff --git a/src/prime_rl/orchestrator/algo/max_rl.py b/src/prime_rl/orchestrator/algo/max_rl.py index 9a3978108d..ff6ccfbf8b 100644 --- a/src/prime_rl/orchestrator/algo/max_rl.py +++ b/src/prime_rl/orchestrator/algo/max_rl.py @@ -20,8 +20,8 @@ class MaxRLAlgorithm(Algorithm): likelihood as it grows). Assumes non-negative (canonically binary) rewards; a group with mean reward - <= 0 carries no signal and gets zero advantages (the zero-advantage filter - drops it, matching the paper's no-success convention).""" + <= 0 carries no signal and gets zero advantages (the sink ships no samples + for all-zero streams, matching the paper's no-success convention).""" async def score_group(self, group: list[Rollout]) -> None: rewards = torch.tensor([rollout.reward for rollout in group], dtype=torch.float32) diff --git a/src/prime_rl/orchestrator/ckpt.py b/src/prime_rl/orchestrator/ckpt.py index 1e9b5173cf..958e1be47d 100644 --- a/src/prime_rl/orchestrator/ckpt.py +++ b/src/prime_rl/orchestrator/ckpt.py @@ -1,6 +1,7 @@ """Checkpoint manager for the orchestrator state (``Progress`` counters + -``TrainSource`` data position). Layout: -``/checkpoints/step_N/orchestrator/progress.pt``.""" +``TaskSampler`` data position and task stats). The sampler state is saved +under the ``"train_source"`` key so checkpoints stay loadable across the +rename. Layout: ``/checkpoints/step_N/orchestrator/progress.pt``.""" from __future__ import annotations @@ -14,7 +15,7 @@ import torch from prime_rl.configs.orchestrator import CheckpointConfig -from prime_rl.orchestrator.train_source import TrainSource +from prime_rl.orchestrator.task_sampler import TaskSampler from prime_rl.orchestrator.types import Progress from prime_rl.utils.logger import format_time, get_logger from prime_rl.utils.pathing import get_ckpt_dir, get_step_path @@ -28,7 +29,7 @@ def __init__(self, output_dir: Path, config: CheckpointConfig) -> None: def get_ckpt_path(self, step: int) -> Path: return get_step_path(self.ckpt_dir, step) / "orchestrator" - def save(self, progress: Progress, train_source: TrainSource, step: int) -> None: + def save(self, progress: Progress, task_sampler: TaskSampler, step: int) -> None: ckpt_path = self.get_ckpt_path(step) ckpt_path.mkdir(parents=True, exist_ok=True) start = time.perf_counter() @@ -37,7 +38,7 @@ def save(self, progress: Progress, train_source: TrainSource, step: int) -> None fd, tmp_name = tempfile.mkstemp(dir=ckpt_path, prefix="progress.pt.", suffix=".tmp") try: with os.fdopen(fd, "wb") as f: - torch.save({"progress": progress, "train_source": train_source.state_dict()}, f) + torch.save({"progress": progress, "train_source": task_sampler.state_dict()}, f) os.replace(tmp_name, ckpt_path / "progress.pt") except BaseException: with contextlib.suppress(OSError): @@ -47,7 +48,7 @@ def save(self, progress: Progress, train_source: TrainSource, step: int) -> None f"Orchestrator checkpoint saved to {ckpt_path} in {format_time(time.perf_counter() - start)}" ) - def load(self, progress: Progress, train_source: TrainSource, step: int) -> None: + def load(self, progress: Progress, task_sampler: TaskSampler, step: int) -> None: ckpt_path = self.get_ckpt_path(step) state_file = ckpt_path / "progress.pt" if not state_file.exists(): @@ -63,11 +64,11 @@ def load(self, progress: Progress, train_source: TrainSource, step: int) -> None for key, value in asdict(saved).items(): if hasattr(progress, key): setattr(progress, key, value) - train_source.load_state_dict(state["train_source"]) + task_sampler.load_state_dict(state["train_source"]) for name, position in state["train_source"]["envs"].items(): - if name not in train_source.base_rows: + if name not in task_sampler.base_rows: continue - rows = train_source.base_rows[name] + rows = task_sampler.base_rows[name] num_tasks = len(rows) if rows is not None else "infinite" get_logger().info( f"Resumed data position for env {name} - epoch={position['epoch']}, " diff --git a/src/prime_rl/orchestrator/detections.py b/src/prime_rl/orchestrator/detections.py new file mode 100644 index 0000000000..9c077fc6d8 --- /dev/null +++ b/src/prime_rl/orchestrator/detections.py @@ -0,0 +1,150 @@ +"""Generation-pathology detections — rollout-level sample hygiene. + +Each detection is a rollout-granularity predicate over the rollout's own +branches (token ids + logprobs), evaluated once at group finalization — the +pipeline's single decision point. Detection results are always recorded; when +``enforce=True`` a detected rollout ships no training samples and never enters +the batch (the batch backfills from fresh groups). Its reward is kept as-is so +it still counts toward the group baseline — the failure signal is real even +when the tokens are poisoned. + +Task-level admission (which tasks to roll out, which finalized groups count) +is the task sampler's job; detections guard the one thing selection cannot: +the policy melting down mid-sample. +""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import TYPE_CHECKING, Protocol + +from prime_rl.configs.orchestrator import DetectionConfig +from prime_rl.utils.logger import get_logger + +if TYPE_CHECKING: + from prime_rl.orchestrator.types import Rollout + + +@dataclass +class DetectionResult: + detected: bool + + +class RolloutDetection(Protocol): + name: str + enforce: bool + + def check(self, rollout: Rollout) -> DetectionResult: ... + + +@dataclass +class GibberishDetection: + """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) -> DetectionResult: + 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 DetectionResult(detected=True) + return DetectionResult(detected=False) + + +@dataclass +class RepetitionDetection: + """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) -> DetectionResult: + for branch in rollout.branches: + # Aligned branch streams (see GibberishDetection), 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 DetectionResult(detected=True) + return DetectionResult(detected=False) + + +def setup_detection(config: DetectionConfig, vocab_size: int) -> RolloutDetection: + """Create a RolloutDetection from a detection config.""" + if config.type == "gibberish": + return GibberishDetection( + 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 RepetitionDetection( + name="repetition", + window=config.window, + logprob_threshold=math.log(config.prob_threshold), + enforce=config.enforce, + ) + raise ValueError(f"Unknown detection type: {config.type}") + + +def setup_detections(configs: list[DetectionConfig], vocab_size: int) -> list[RolloutDetection]: + """Create RolloutDetections from a list of detection configs.""" + detections = [setup_detection(config, vocab_size) for config in configs] + if detections: + get_logger().info(f"Configured {len(detections)} rollout detection(s):") + for config, detection in zip(configs, detections): + mode = "Enforcing" if detection.enforce else "Monitoring" + params = ", ".join(f"{k}={v}" for k, v in config.model_dump().items()) + get_logger().info(f" {mode} {detection.name} detection ({params})") + return detections + + +def run_detections(detections: list[RolloutDetection], rollout: Rollout) -> None: + """Stamp one rollout in place with per-detection results + the exclusion + verdict: ``rollout.detections`` records per-name bools; ``is_excluded`` is + True iff an enforcing detection fired. First match wins (no double + counting). Reward and trajectory tokens are left untouched so the rollout + still contributes to baseline calculations and metric aggregation.""" + rollout.detections = {d.name: False for d in detections} + rollout.is_excluded = False + for detection in detections: + if detection.check(rollout).detected: + rollout.detections[detection.name] = True + if detection.enforce: + rollout.is_excluded = True + break diff --git a/src/prime_rl/orchestrator/dispatcher.py b/src/prime_rl/orchestrator/dispatcher.py index 80049cd461..ad32992345 100644 --- a/src/prime_rl/orchestrator/dispatcher.py +++ b/src/prime_rl/orchestrator/dispatcher.py @@ -36,7 +36,7 @@ from prime_rl.orchestrator.envs import EvalEnvs, TrainEnvs from prime_rl.orchestrator.eval_source import EvalSource -from prime_rl.orchestrator.train_source import TrainSource +from prime_rl.orchestrator.task_sampler import TaskSampler from prime_rl.orchestrator.types import ( GroupState, InflightRollout, @@ -115,7 +115,7 @@ def drain_keys(*, train_envs: set[str], eval_envs: set[str]) -> list[str]: class RolloutDispatcher: """``await dispatcher.start()`` runs the dispatch loop until ``stop()``. - Pulls examples from ``TrainSource`` / ``EvalSource``, schedules + Pulls examples from ``TaskSampler`` / ``EvalSource``, schedules rollouts under shared capacity, and emits ``Rollout``\\ s to ``out_q``. The watcher drives ``on_version_pending`` for off-policy cancellation; the orchestrator triggers eval epochs.""" @@ -125,7 +125,7 @@ def __init__( *, train_envs: TrainEnvs, eval_envs: EvalEnvs | None, - train_source: TrainSource, + task_sampler: TaskSampler, eval_source: EvalSource | None, policy_pool: InferencePool, policy: Policy, @@ -139,7 +139,7 @@ def __init__( # Train rollouts go to the env's rollout-source pool; eval always # evaluates the policy. self.policy_pool = policy_pool - self.train_source = train_source + self.task_sampler = task_sampler self.eval_source = eval_source self.max_off_policy_steps = max_off_policy_steps @@ -367,7 +367,7 @@ def next_fresh_group(self, kind: RolloutKind, envs) -> GroupState | None: """Pop the next example from the corresponding source and wrap it in a ``GroupState``. Returns ``None`` if the source is empty.""" if kind == "train": - source = self.train_source + source = self.task_sampler else: assert self.eval_source is not None source = self.eval_source diff --git a/src/prime_rl/orchestrator/envs.py b/src/prime_rl/orchestrator/envs.py index 193d93fa82..c8a675be1f 100644 --- a/src/prime_rl/orchestrator/envs.py +++ b/src/prime_rl/orchestrator/envs.py @@ -57,7 +57,7 @@ def __init__(self, config: EnvConfig, address: str): self.tasks: Iterator[vf.Task] | None = None """The env's tasks, client-side, set at ``start()``. A finite taskset is materialized (``num_tasks`` is its count) and iterated from there; an infinite - one streams off its generator. Consumed once — by ``TrainSource`` (train) or + one streams off its generator. Consumed once — by ``TaskSampler`` (train) or ``EvalEnv.start`` (eval).""" self._env_client: EnvClient | None = None diff --git a/src/prime_rl/orchestrator/filters.py b/src/prime_rl/orchestrator/filters.py deleted file mode 100644 index ad023fd928..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 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 190a53601f..f4e8cb6dfb 100644 --- a/src/prime_rl/orchestrator/metrics.py +++ b/src/prime_rl/orchestrator/metrics.py @@ -342,16 +342,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. + # The pipeline verdicts are per-trace (an untrainable seat is 0.0 throughout, and + # detections 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}/is_excluded/mean"] = sum(float(r.is_excluded) for r in rollouts) / len(rollouts) + 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}/detections/{name}/mean": sum(1 for r in rollouts if r.detections.get(name)) / len(rollouts) for name in names } return out @@ -414,7 +414,7 @@ def __iter__(self) -> Iterator[Rollout]: @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]) + return TrainRollouts([r for r in self.rollouts if not r.has_error and not r.is_excluded and r.agent.trainable]) def by_env(self) -> dict[str, TrainRollouts]: grouped: dict[str, list[Rollout]] = {} diff --git a/src/prime_rl/orchestrator/orchestrator.py b/src/prime_rl/orchestrator/orchestrator.py index d27eda3f98..360dea2396 100644 --- a/src/prime_rl/orchestrator/orchestrator.py +++ b/src/prime_rl/orchestrator/orchestrator.py @@ -5,7 +5,7 @@ - ``RolloutDispatcher`` schedules rollouts; emits ``Rollout`` (train/eval discriminated by ``kind``) on its queue. -- ``TrainSink`` ingests train rollouts (tokenize → advantages → filters) +- ``TrainSink`` ingests train rollouts (tokenize → detections → advantages) 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.detections import setup_detections 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.packing import BatchPacker from prime_rl.orchestrator.patches import ( @@ -54,8 +54,8 @@ monkey_patch_oai_iterable_types, ) from prime_rl.orchestrator.periodic_logger import PeriodicLogger +from prime_rl.orchestrator.task_sampler import TaskSampler from prime_rl.orchestrator.train_sink import TrainSink -from prime_rl.orchestrator.train_source import TrainSource from prime_rl.orchestrator.types import ( EvalBatch, Policy, @@ -95,9 +95,9 @@ # 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 ship nothing — usually a +# homogeneous-reward dataset (every advantage stream collapses) or an +# over-eager detection config; fail loudly instead of spinning MAX_CONSECUTIVE_EMPTY_BATCHES = 10 # Maximum batches the orchestrator may run ahead of the trainer. The @@ -131,7 +131,7 @@ class Orchestrator: sender: MicroBatchSender | None packer: BatchPacker train_envs: TrainEnvs - train_source: TrainSource + task_sampler: TaskSampler train_sink: TrainSink dispatcher: RolloutDispatcher watcher: WeightWatcher @@ -247,9 +247,8 @@ async def setup(self) -> None: if usage_base_url and usage_api_key: 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") + # Detections apply to train rollouts only + detections = setup_detections(config.detections, vocab_size=self.tokenizer.vocab_size) get_logger().info("Loading training environments") self.train_envs = TrainEnvs( @@ -328,10 +327,10 @@ async def setup(self) -> None: self.lora_name = config.model.lora.name if config.model.lora else None - self.train_source = TrainSource(self.train_envs) + self.task_sampler = TaskSampler(self.train_envs, drop_degenerate_groups=config.sampler.drop_degenerate_groups) if self.resume_step is not None and self.ckpt_manager is not None: - self.ckpt_manager.load(self.progress, self.train_source, step=self.resume_step) + self.ckpt_manager.load(self.progress, self.task_sampler, step=self.resume_step) # The checkpoint finished step ``resume_step``; resume at the next step. Derive the step # from ``resume_step`` (not the loaded progress.step) so it stays coordinated with the # trainer even when ``ckpt.skip_progress`` left the counter unrestored. @@ -402,7 +401,7 @@ async def setup(self) -> None: self.dispatcher = RolloutDispatcher( train_envs=self.train_envs, eval_envs=self.eval_envs, - train_source=self.train_source, + task_sampler=self.task_sampler, eval_source=self.eval_source, policy_pool=self.policy_inference, policy=self.policy, @@ -417,8 +416,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, - pre_filters=pre_filters, - post_filters=post_filters, + detections=detections, + on_group_finalized=self.task_sampler.observe, ) self.eval_sink = EvalSink(eval_envs=self.eval_envs) if self.eval_envs is not None else None self.watcher = WeightWatcher( @@ -505,7 +504,7 @@ async def start(self) -> None: if self.ckpt_manager is not None and self.progress.step > 1: self.progress.step -= 1 get_logger().info("Writing final checkpoint") - self.ckpt_manager.save(self.progress, self.train_source, step=self.progress.step) + self.ckpt_manager.save(self.progress, self.task_sampler, step=self.progress.step) await self.stop() if clean_exit: get_logger().success("Orchestrator finished.") @@ -602,7 +601,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 detection config ([[orchestrator.detections]]) or task difficulty." ) return self.consecutive_empty_batches = 0 @@ -611,7 +610,7 @@ async def finalize_train_batch(self, batch: TrainBatch) -> None: if effective and n_trainable / len(effective) <= 0.1: get_logger().warning( f"Only {n_trainable}/{len(effective)} effective rollouts are trainable " - f"({n_trainable / len(effective):.1%}) — consider reviewing task difficulty / filter config" + f"({n_trainable / len(effective):.1%}) — consider reviewing task difficulty / detection config" ) # Ship batch ``step`` only once the trainer has published v{step-1-TARGET_LAG}. @@ -695,12 +694,15 @@ 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 + if self.train_sink.detection_seen > 0: + metrics["detections/all/excluded_rate"] = ( + self.train_sink.detection_excluded / self.train_sink.detection_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 + for name, count in self.train_sink.detection_by_name.items(): + metrics[f"detections/all/{name}/rate"] = count / self.train_sink.detection_seen + for env_name, count in self.train_sink.degenerate_groups_dropped.items(): + metrics[f"sampler/{env_name}/dropped_degenerate_groups"] = float(count) + metrics |= self.task_sampler.metrics() self.monitor.log(metrics, step=step) self.wait_for_policy_time = 0.0 self.monitor.log_samples(effective.rollouts, step=step) @@ -729,7 +731,7 @@ async def finalize_train_batch(self, batch: TrainBatch) -> None: self.log_train_batch(batch, step=step, step_time=step_time) - self.train_sink.reset_pre_filter_stats() + self.train_sink.reset_detection_stats() self.maybe_trigger_eval(self.progress.step) trim_process_memory() @@ -915,8 +917,8 @@ async def maybe_save_ckpt(self, step: int) -> float: get_logger().info(f"Saving checkpoint at step {step}") t = time.perf_counter() # Synchronous on purpose: the payload is tiny, and snapshotting on the - # event loop keeps the dispatcher from mutating TrainSource mid-save - self.ckpt_manager.save(self.progress, self.train_source, step) + # event loop keeps the dispatcher from mutating TaskSampler mid-save + self.ckpt_manager.save(self.progress, self.task_sampler, step) return time.perf_counter() - t def update_dispatch_gate(self) -> None: diff --git a/src/prime_rl/orchestrator/train_source.py b/src/prime_rl/orchestrator/task_sampler.py similarity index 59% rename from src/prime_rl/orchestrator/train_source.py rename to src/prime_rl/orchestrator/task_sampler.py index 9c401821cf..26874c64c0 100644 --- a/src/prime_rl/orchestrator/train_source.py +++ b/src/prime_rl/orchestrator/task_sampler.py @@ -1,29 +1,42 @@ -"""TrainSource: weighted round-robin across train envs, infinite pull. +"""TaskSampler: picks what to roll out next, remembers how it went. -Weights are each env's configured ``ratio`` (default 1, i.e. equal weight -per env). An env serves the tasks the orchestrator loaded client-side: a -finite one as a shuffled table (reshuffled with ``seed=epoch`` on cursor -exhaustion), an infinite one (``num_tasks is None``) straight off its -generator — every pull is a fresh task and there are no epochs to shuffle.""" +The pick side is weighted round-robin across train envs by their configured +``ratio`` (default 1), then the env's next task: a finite taskset as a +shuffled table (reshuffled with ``seed=epoch`` on cursor exhaustion), an +infinite one (``num_tasks is None``) straight off its generator — every pull +is a fresh task and there are no epochs to shuffle. + +The observe side is the sink's group-finalization hook: every finalized train +group updates the per-task outcome stats (:mod:`.task_stats`), always on, so +difficulty estimates are warm before anything reads them. Nothing consults the +stats when sampling yet — picks are identical to the pre-stats behavior — +and the sampler stays advisory either way: it returns tasks and consumes +outcomes, never tokens or training samples.""" from __future__ import annotations import random from collections.abc import Iterator +from typing import TYPE_CHECKING import verifiers.v1 as vf from prime_rl.orchestrator.envs import TrainEnvs +from prime_rl.orchestrator.task_stats import TaskStats + +if TYPE_CHECKING: + from prime_rl.orchestrator.types import Rollout -class TrainSource: +class TaskSampler: """``next_example()`` picks a weighted-RR env and returns its next example. Returned dicts carry ``env_name`` + ``task``, whose data is shipped to the env server at dispatch. - The data position round-trips through a checkpoint via ``state_dict()`` - / ``load_state_dict()``: per-env ``{epoch, cursor}`` plus the env-choice - RNG state, making the dispatch sequence reproducible across resumes. For + The data position and the task stats round-trip through a checkpoint via + ``state_dict()`` / ``load_state_dict()``: per-env ``{epoch, cursor}`` plus + the env-choice RNG state, making the dispatch sequence reproducible across + resumes. For a finite env, epochs are 1-indexed and seed that epoch's shuffle; for an infinite env the epoch stays 1 and the cursor counts generator pulls, replayed on restore by fast-forwarding the generator (exact iff it's @@ -31,11 +44,13 @@ class TrainSource: batches), so a resume skips the tasks that were in flight at checkpoint time.""" - def __init__(self, train_envs: TrainEnvs) -> None: + def __init__(self, train_envs: TrainEnvs, *, drop_degenerate_groups: bool = False) -> None: self.rng = random.Random(42) self.envs = list(train_envs) if not self.envs: - raise ValueError("TrainSource needs at least one train env") + raise ValueError("TaskSampler needs at least one train env") + self.task_stats = TaskStats() + self.drop_degenerate_groups = drop_degenerate_groups # A finite env's example table in canonical order (each epoch's shuffle # starts from this); ``None`` for an infinite env, whose generator @@ -71,15 +86,33 @@ def _shuffle(self, env_name: str) -> list[dict] | None: random.Random(self.epochs[env_name]).shuffle(rows) return rows + def observe(self, group: list[Rollout]) -> bool: + """The sink's group-finalization hook: fold one finalized train group + into the per-task stats. Returns the sampler's drop verdict — True + asks the sink to keep a degenerate (zero-signal) group out of the + batch so it backfills from fresh groups; the sink still exempts + sources whose algorithm trains on zero advantage.""" + signal = self.task_stats.observe(group) + return self.drop_degenerate_groups and not signal + + def metrics(self) -> dict[str, float]: + """Drain the sampler metric family (pool occupancy, coverage, + realized signal rate, wasted tokens) for the step log.""" + return self.task_stats.metrics({env.name: env.num_tasks for env in self.envs}) + def state_dict(self) -> dict: - """Env-choice RNG state + per-env ``{epoch, cursor}``.""" + """Env-choice RNG state + per-env ``{epoch, cursor}`` + task stats.""" return { "rng": self.rng.getstate(), "envs": {name: {"epoch": self.epochs[name], "cursor": self.cursors[name]} for name in self.epochs}, + "stats": self.task_stats.state_dict(), } def load_state_dict(self, state_dict: dict) -> None: self.rng.setstate(state_dict["rng"]) + # Checkpoints written before task stats existed load with empty stats + # and warm back up from live groups. + self.task_stats.load_state_dict(state_dict.get("stats", {})) for name, position in state_dict["envs"].items(): if name not in self.base_rows: continue diff --git a/src/prime_rl/orchestrator/task_stats.py b/src/prime_rl/orchestrator/task_stats.py new file mode 100644 index 0000000000..04309bc0bc --- /dev/null +++ b/src/prime_rl/orchestrator/task_stats.py @@ -0,0 +1,187 @@ +"""Per-task outcome statistics — the task sampler's memory. + +Every finalized train group is a free experiment: ``group_size`` episodes of +the current policy against one task. ``TaskStats`` keeps that evidence as +discounted success/failure pseudo-counts (a Beta posterior with forgetting) +plus reward EMAs, keyed by ``(env_name, task_key, agent role)``. The discount +is the staleness answer — the policy moves, so old outcomes must fade — and +nothing is ever evicted: estimates drift back toward the prior when a task +goes unobserved, so no task is written off permanently. + +Keys are content hashes of the task's data, recomputable from the trace echo +(``rollout.task.data``), so stats survive dataset reordering and resumes and +can be joined against saved trace records offline. Roles are kept separate +because multi-agent groups mix reward scales (a proposer's reward says nothing +about solver difficulty). + +Nothing here decides anything: this module only remembers and reports. +Sampling weights that *read* the posterior arrive with the weighted-sampling +config surface. +""" + +from __future__ import annotations + +import hashlib +import json +import math +from collections import defaultdict +from dataclasses import asdict, dataclass +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from prime_rl.orchestrator.types import Rollout + +# Evidence discount per observed group: old outcomes fade as the policy moves. +DECAY = 0.9 +# Beta prior pseudo-counts (alpha, beta): an unseen task sits at p_hat = 0.5. +PRIOR = (1.0, 1.0) +# A trace counts as a success when its reward reaches this threshold. +SUCCESS_THRESHOLD = 0.5 +# p_hat bands for the pool occupancy metrics. +HOPELESS_BELOW = 0.05 +SATURATED_ABOVE = 0.95 + + +def task_key(task_data: dict) -> str: + """Stable content key for one task: blake2b over the canonical JSON of its + dumped data. Identical content hashes identically wherever it round-trips + (dispatch request, trace echo, saved records).""" + canonical = json.dumps(task_data, sort_keys=True, separators=(",", ":"), default=str) + return hashlib.blake2b(canonical.encode(), digest_size=8).hexdigest() + + +@dataclass +class TaskStat: + """Discounted evidence for one ``(env, task, role)``.""" + + s: float = 0.0 + """Discounted success count.""" + f: float = 0.0 + """Discounted failure count.""" + reward_mean: float = 0.0 + reward_std: float = 0.0 + """EMA of the per-group reward std — the signal proxy for non-binary rewards.""" + draws_per_group: float = 0.0 + """EMA of role traces observed per group (1 for single-agent envs).""" + visits: int = 0 + last_seen_version: int = 0 + + @property + def p_hat(self) -> float: + """Posterior mean success rate under the Beta prior.""" + alpha, beta = PRIOR + return (alpha + self.s) / (alpha + beta + self.s + self.f) + + def update( + self, *, successes: int, failures: int, reward_mean: float, reward_std: float, draws: int, version: int + ) -> None: + self.s = DECAY * self.s + successes + self.f = DECAY * self.f + failures + if self.visits == 0: + self.reward_mean = reward_mean + self.reward_std = reward_std + self.draws_per_group = float(draws) + else: + w = 1.0 - DECAY + self.reward_mean += w * (reward_mean - self.reward_mean) + self.reward_std += w * (reward_std - self.reward_std) + self.draws_per_group += w * (draws - self.draws_per_group) + self.visits += 1 + self.last_seen_version = max(self.last_seen_version, version) + + +class TaskStats: + """The store plus per-tick counters for the sampler metric family. + + ``observe`` consumes one finalized train group; ``metrics`` drains the + tick counters and snapshots the pool occupancy. Only clean, trainable + traces update evidence — errored episodes and off-policy cancellations + (which arrive as error markers) say nothing about task difficulty — but a + degenerate zero-signal group is itself the strongest difficulty datum and + updates counts like any other outcome. + """ + + def __init__(self) -> None: + # env -> task_key -> role -> TaskStat + self.stats: dict[str, dict[str, dict[str, TaskStat]]] = {} + # Per-tick counters, drained by ``metrics()``. + self._groups: dict[str, int] = defaultdict(int) + self._signal_groups: dict[str, int] = defaultdict(int) + self._tokens: dict[str, int] = defaultdict(int) + self._wasted_tokens: dict[str, int] = defaultdict(int) + + def observe(self, group: list[Rollout]) -> bool: + """Fold one finalized group into the stats; returns whether the group + bought any gradient (a nonzero advantage on some rollout).""" + env_name = group[0].env_name + tokens = sum(r.num_total_tokens for r in group) + signal = any(r.is_trainable for r in group) + self._groups[env_name] += 1 + self._signal_groups[env_name] += int(signal) + self._tokens[env_name] += tokens + if not signal: + self._wasted_tokens[env_name] += tokens + + clean = [r for r in group if not r.has_error and r.agent.trainable] + if not clean: + return signal + key = task_key(clean[0].task.data.model_dump(mode="json")) + by_role: dict[str, list[Rollout]] = defaultdict(list) + for rollout in clean: + by_role[rollout.agent.name].append(rollout) + for role, rollouts in by_role.items(): + rewards = [r.reward for r in rollouts] + mean = sum(rewards) / len(rewards) + std = math.sqrt(sum((x - mean) ** 2 for x in rewards) / len(rewards)) + successes = sum(1 for x in rewards if x >= SUCCESS_THRESHOLD) + stat = self.stats.setdefault(env_name, {}).setdefault(key, {}).setdefault(role, TaskStat()) + stat.update( + successes=successes, + failures=len(rewards) - successes, + reward_mean=mean, + reward_std=std, + draws=len(rewards), + version=max(r.policy_version for r in rollouts), + ) + return signal + + def metrics(self, num_tasks: dict[str, int | None]) -> dict[str, float]: + """Sampler metric family, per env. Pool occupancy counts one unit per + tracked ``(task, role)`` stat; ``unseen``/``coverage`` need the finite + table size and are skipped for infinite tasksets. Tick counters + (realized signal rate, wasted tokens) are drained on read.""" + out: dict[str, float] = {} + for env_name, total in num_tasks.items(): + tracked = self.stats.get(env_name, {}) + units = [stat for roles in tracked.values() for stat in roles.values()] + if units: + hopeless = sum(1 for u in units if u.p_hat < HOPELESS_BELOW) + saturated = sum(1 for u in units if u.p_hat > SATURATED_ABOVE) + out[f"sampler/{env_name}/pool/hopeless"] = float(hopeless) + out[f"sampler/{env_name}/pool/saturated"] = float(saturated) + out[f"sampler/{env_name}/pool/learnable"] = float(len(units) - hopeless - saturated) + out[f"sampler/{env_name}/p_hat/mean"] = sum(u.p_hat for u in units) / len(units) + if total is not None: + seen = len(tracked) + out[f"sampler/{env_name}/pool/unseen"] = float(max(0, total - seen)) + out[f"sampler/{env_name}/coverage"] = seen / total if total else 0.0 + groups = self._groups.pop(env_name, 0) + if groups: + out[f"sampler/{env_name}/groups_observed"] = float(groups) + out[f"sampler/{env_name}/realized_signal_rate"] = self._signal_groups.pop(env_name, 0) / groups + tokens = self._tokens.pop(env_name, 0) + if tokens: + out[f"sampler/{env_name}/wasted_token_frac"] = self._wasted_tokens.pop(env_name, 0) / tokens + return out + + def state_dict(self) -> dict: + return { + env: {key: {role: asdict(stat) for role, stat in roles.items()} for key, roles in tasks.items()} + for env, tasks in self.stats.items() + } + + def load_state_dict(self, state: dict) -> None: + self.stats = { + env: {key: {role: TaskStat(**stat) for role, stat in roles.items()} for key, roles in tasks.items()} + for env, tasks in state.items() + } diff --git a/src/prime_rl/orchestrator/train_sink.py b/src/prime_rl/orchestrator/train_sink.py index 5f052c14ed..eafac306d3 100644 --- a/src/prime_rl/orchestrator/train_sink.py +++ b/src/prime_rl/orchestrator/train_sink.py @@ -4,11 +4,15 @@ 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. -3. ``process_batch`` — applies post-batch filter annotations and assembles - the trainer-bound ``TrainingSample`` list. Returns a ``TrainBatch``. +2. ``process_group`` — the single decision point: filters errored rollouts, + hands the trainable survivors to the env algorithm's ``finalize_group`` + (advantages + per-sample wire stamping), runs the generation-pathology + detections, reports the finalized group to the task sampler (whose + verdict may drop a degenerate group so the batch backfills), and queues + the admitted survivors. +3. ``process_batch`` — assembles the trainer-bound ``TrainingSample`` list; + all-zero advantage streams ship nothing unless the source's algorithm + trains on zero advantage. Returns a ``TrainBatch``. ``add()`` takes one episode (``list[Rollout]``) and returns ``TrainBatch | None``; group accounting counts episodes, never loose traces. @@ -21,10 +25,11 @@ import asyncio import uuid from collections import defaultdict +from collections.abc import Callable from prime_rl.configs.orchestrator import OrchestratorConfig +from prime_rl.orchestrator.detections import RolloutDetection, run_detections 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 Rollout, TrainBatch @@ -57,8 +62,8 @@ 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], + detections: list[RolloutDetection], + on_group_finalized: Callable[[list[Rollout]], bool] | None = None, ) -> None: assert (batch_size is None) != (token_batch_size is None), ( "Exactly one of batch_size / token_batch_size must be set" @@ -69,11 +74,15 @@ 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.detections = detections + # Fired once per finalized group with the full cohort (errored + + # excluded + survivors) — the task sampler's outcome feedback. Its + # return value is the sampler's degenerate-group drop verdict; the + # sink stays sampler-blind, the callback stays token-blind. + self.on_group_finalized = on_group_finalized # Observation window for the next shipped batch: rollouts of groups - # finalized since the last ship (errored + filtered + survivors). + # finalized since the last ship (errored + excluded + survivors). # In-progress groups stay out until they finalize. self.pending_rollouts: TrainRollouts = TrainRollouts() # Keyed by the dispatcher's group UUID. ``(env_name, task_idx)`` @@ -89,10 +98,11 @@ def __init__( # re-sums per arrival. self.pending_tokens: int = 0 - # Reset by the orchestrator after each ship via ``reset_pre_filter_stats`` - self.pre_filter_seen = 0 - self.pre_filter_dropped = 0 - self.pre_filter_dropped_by_name: dict[str, int] = {} + # Reset by the orchestrator after each ship via ``reset_detection_stats`` + self.detection_seen = 0 + self.detection_excluded = 0 + self.detection_by_name: dict[str, int] = {} + self.degenerate_groups_dropped: dict[str, int] = defaultdict(int) def group_size_for(self, env_name: str) -> int: return self.train_envs.get(env_name).config.group_size @@ -167,7 +177,8 @@ async def process_rollout(self, rollout: Rollout) -> None: async def process_group(self, group_id: uuid.UUID) -> None: """Finalize one GRPO group: drop errored rollouts, assign advantages, - run pre-batch filters, append survivors to ``pending_batch``.""" + report the group to the task sampler, append admitted survivors to + ``pending_batch``.""" group = self.pending_groups.pop(group_id, []) self.pending_group_episodes.pop(group_id, None) if not group: @@ -191,6 +202,10 @@ async def process_group(self, group_id: uuid.UUID) -> None: f"Finished group | env={env_name} task_idx={task_idx} | " f"rollouts={len(group)} (errored={num_errored}) | dropped: no trainable survivors" ) + # Still an outcome: an all-errored/untrainable group is pure waste + # and the sampler should know. + if self.on_group_finalized is not None: + self.on_group_finalized(group) return # Advantages + per-sample wire stamping (advantage stream, loss @@ -205,44 +220,51 @@ 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 + # Advantages are stamped — the group is now evidence. The sampler's + # verdict may drop a degenerate group so the batch backfills; sources + # whose algorithm trains on zero advantage are exempt. + drop_group = False + if self.on_group_finalized is not None: + drop_group = self.on_group_finalized(group) and not env.algorithm.trains_on_zero_advantage + if drop_group: + self.degenerate_groups_dropped[env_name] += 1 + + excluded_by_name: dict[str, int] = {} + num_excluded = 0 for r in survivors: - self.pre_filter_seen += 1 - if r.is_filtered: - self.pre_filter_dropped += 1 - num_filtered += 1 - for name, hit in r.filter_results.items(): + # Rollout-granularity predicates, evaluated here with everything + # else — group finalization is the pipeline's one decision point. + run_detections(self.detections, r) + self.detection_seen += 1 + if r.is_excluded: + self.detection_excluded += 1 + num_excluded += 1 + for name, hit in r.detections.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 + self.detection_by_name[name] = self.detection_by_name.get(name, 0) + 1 + excluded_by_name[name] = excluded_by_name.get(name, 0) + 1 + continue + if drop_group: 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 "—" + detection_str = ", ".join(f"{n}={c}" for n, c in excluded_by_name.items()) if excluded_by_name 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}, excluded={num_excluded}" + f"{', dropped: degenerate' if drop_group else ''}) | " + f"reward={avg_reward:.4f} | detections: {detection_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.""" + set) and assemble 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 :] @@ -259,15 +281,14 @@ 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. An all-zero advantage stream is + # dead weight for the trainer unless the source's algorithm declares otherwise + # (echo's observation CE trains through collapsed advantages). + samples: list[TrainingSample] = [sample for r in cohort if self._ships(r) 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 + # last ship (errored + excluded + survivors) — while ``samples`` is the shipped cohort's # trainable payload. ``rollouts.effective`` / ``rollouts.metrics`` derive the clean subset + # metric views on demand. Reset the window only when the batch actually ships (non-empty # samples) — an empty batch is dropped unlogged by the orchestrator, so keep accumulating its @@ -277,7 +298,18 @@ def process_batch(self) -> TrainBatch: self.pending_rollouts = TrainRollouts() return TrainBatch(rollouts=rollouts, samples=samples) - def reset_pre_filter_stats(self) -> None: - self.pre_filter_seen = 0 - self.pre_filter_dropped = 0 - self.pre_filter_dropped_by_name.clear() + def _ships(self, rollout: Rollout) -> bool: + """Whether a queued rollout's samples ship to the trainer: an all-zero + advantage stream carries no gradient (``None`` streams — opd/opsd/sft — + always ship; their signal is not advantage-based).""" + if rollout.advantages is None: + return True + if any(a != 0.0 for a in rollout.advantages): + return True + return self.train_envs.get(rollout.env_name).algorithm.trains_on_zero_advantage + + def reset_detection_stats(self) -> None: + self.detection_seen = 0 + self.detection_excluded = 0 + self.detection_by_name.clear() + self.degenerate_groups_dropped.clear() diff --git a/src/prime_rl/orchestrator/types.py b/src/prime_rl/orchestrator/types.py index b77de8ba31..d2f0fd83ca 100644 --- a/src/prime_rl/orchestrator/types.py +++ b/src/prime_rl/orchestrator/types.py @@ -99,8 +99,11 @@ class Rollout(vf.Trace[DataT], Generic[DataT]): # 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) + # Arrival-time detection results (per-name) + the exclusion verdict: an + # enforcing detection keeps the rollout's samples out of training while its + # reward still counts toward the group baseline. + detections: dict[str, bool] = Field(default_factory=dict, exclude=True) + is_excluded: bool = Field(default=False, exclude=True) eval_step: int | None = Field(default=None, exclude=True) def assign_advantages(self, values: float | list[float]) -> None: diff --git a/tests/unit/orchestrator/test_filters.py b/tests/unit/orchestrator/test_detections.py similarity index 51% rename from tests/unit/orchestrator/test_filters.py rename to tests/unit/orchestrator/test_detections.py index 69ce76d029..cafea0b83e 100644 --- a/tests/unit/orchestrator/test_filters.py +++ b/tests/unit/orchestrator/test_detections.py @@ -3,19 +3,19 @@ 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.configs.orchestrator import GibberishDetectionConfig, RepetitionDetectionConfig +from prime_rl.orchestrator.detections import ( + GibberishDetection, + RepetitionDetection, + run_detections, + setup_detection, + setup_detections, ) from prime_rl.orchestrator.types import Rollout 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 + """An assistant node whose tokens are all model-sampled (the detections read each node's masked-True tokens + logprobs).""" return vf.MessageNode( message=vf.AssistantMessage(content="x"), @@ -48,7 +48,7 @@ def _make_rollout( multi_step: bool = False, ) -> Rollout: """Build a ``Rollout`` (a message-graph trace) carrying the completion tokens — enough for - the filters to inspect each node's sampled tokens / logprobs.""" + the detections to inspect each node's sampled tokens / logprobs.""" if multi_step: mid = len(completion_ids) // 2 nodes = [ @@ -68,38 +68,38 @@ def _make_rollout( return rollout -def _make_gibberish_filter(vocab_size=128_000, token_id_threshold=100_000, logprob_offset=2.0, enforce=False): +def _make_gibberish_detection(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( + return GibberishDetection( 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( +def _make_repetition_detection(window=5, prob_threshold=0.99, enforce=False): + return RepetitionDetection( name="repetition", window=window, logprob_threshold=math.log(prob_threshold), enforce=enforce ) -# --- GibberishFilter tests --- +# --- GibberishDetection tests --- def test_gibberish_detects_rare_low_prob_token(): - gibberish_filter = _make_gibberish_filter() + gibberish_detection = _make_gibberish_detection() - result = gibberish_filter.check( + result = gibberish_detection.check( _make_rollout( completion_ids=[50, 120_000, 80], - completion_logprobs=[-1.0, gibberish_filter.logprob_threshold - 1.0, -0.5], + completion_logprobs=[-1.0, gibberish_detection.logprob_threshold - 1.0, -0.5], ) ) assert result.detected is True def test_gibberish_ignores_normal_tokens(): - gibberish_filter = _make_gibberish_filter() + gibberish_detection = _make_gibberish_detection() - result = gibberish_filter.check( + result = gibberish_detection.check( _make_rollout( completion_ids=[10, 200, 5000], completion_logprobs=[-1.0, -2.0, -3.0], @@ -109,9 +109,9 @@ def test_gibberish_ignores_normal_tokens(): def test_gibberish_ignores_high_prob_rare_token(): - gibberish_filter = _make_gibberish_filter() + gibberish_detection = _make_gibberish_detection() - result = gibberish_filter.check( + result = gibberish_detection.check( _make_rollout( completion_ids=[120_000], completion_logprobs=[-0.5], @@ -121,12 +121,12 @@ def test_gibberish_ignores_high_prob_rare_token(): def test_gibberish_works_across_trajectory_steps(): - gibberish_filter = _make_gibberish_filter() + gibberish_detection = _make_gibberish_detection() - result = gibberish_filter.check( + result = gibberish_detection.check( _make_rollout( completion_ids=[50, 60, 120_000, 80], - completion_logprobs=[-1.0, -0.5, gibberish_filter.logprob_threshold - 1.0, -0.5], + completion_logprobs=[-1.0, -0.5, gibberish_detection.logprob_threshold - 1.0, -0.5], multi_step=True, ) ) @@ -138,26 +138,26 @@ def test_gibberish_aligns_logprobs_under_generation_prompt_scaffold(): 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() + gibberish_detection = _make_gibberish_detection() rollout = Rollout[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])], + nodes=[_scaffold_assistant_node([50, 80, 120_000], [-1.0, -0.5, gibberish_detection.logprob_threshold - 1.0])], rewards={"reward": vf.Reward(score=1.0)}, ) - result = gibberish_filter.check(rollout) + result = gibberish_detection.check(rollout) assert result.detected is True -# --- RepetitionFilter tests --- +# --- RepetitionDetection tests --- def test_repetition_triggers_after_window(): - repetition_filter = _make_repetition_filter(window=5) + repetition_detection = _make_repetition_detection(window=5) - result = repetition_filter.check( + result = repetition_detection.check( _make_rollout( completion_ids=list(range(5)), completion_logprobs=[-0.001] * 5, @@ -167,9 +167,9 @@ def test_repetition_triggers_after_window(): def test_repetition_no_trigger_below_window(): - repetition_filter = _make_repetition_filter(window=5) + repetition_detection = _make_repetition_detection(window=5) - result = repetition_filter.check( + result = repetition_detection.check( _make_rollout( completion_ids=list(range(4)), completion_logprobs=[-0.001] * 4, @@ -179,10 +179,10 @@ def test_repetition_no_trigger_below_window(): def test_repetition_resets_on_low_prob(): - repetition_filter = _make_repetition_filter(window=5) + repetition_detection = _make_repetition_detection(window=5) logprobs = [-0.001] * 3 + [-2.0] + [-0.001] * 3 - result = repetition_filter.check( + result = repetition_detection.check( _make_rollout( completion_ids=list(range(7)), completion_logprobs=logprobs, @@ -192,9 +192,9 @@ def test_repetition_resets_on_low_prob(): def test_repetition_varied_probs_no_trigger(): - repetition_filter = _make_repetition_filter(window=3) + repetition_detection = _make_repetition_detection(window=3) - result = repetition_filter.check( + result = repetition_detection.check( _make_rollout( completion_ids=list(range(6)), completion_logprobs=[-0.001, -3.0, -0.001, -3.0, -0.001, -3.0], @@ -203,76 +203,76 @@ def test_repetition_varied_probs_no_trigger(): assert result.detected is False -# --- setup_filter / setup_filters tests --- +# --- setup_detection / setup_detections 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_detection_gibberish(): + config = GibberishDetectionConfig(token_id_threshold=100_000, logprob_offset=2.0) + gibberish_detection = setup_detection(config, vocab_size=128_000) + assert isinstance(gibberish_detection, GibberishDetection) + assert gibberish_detection.name == "gibberish" + assert gibberish_detection.token_id_threshold == 100_000 + assert abs(gibberish_detection.logprob_threshold - (-math.log(128_000) - 2.0)) < 1e-10 + assert gibberish_detection.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_detection_gibberish_enforce(): + config = GibberishDetectionConfig(enforce=True) + gibberish_detection = setup_detection(config, vocab_size=128_000) + assert gibberish_detection.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_detection_repetition(): + config = RepetitionDetectionConfig(window=3_000, prob_threshold=0.99) + repetition_detection = setup_detection(config, vocab_size=128_000) + assert isinstance(repetition_detection, RepetitionDetection) + assert repetition_detection.name == "repetition" + assert repetition_detection.window == 3_000 + assert abs(repetition_detection.logprob_threshold - math.log(0.99)) < 1e-10 + assert repetition_detection.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_detection_repetition_enforce(): + config = RepetitionDetectionConfig(enforce=True) + repetition_detection = setup_detection(config, vocab_size=128_000) + assert repetition_detection.enforce is True -def test_setup_filters_multiple(): +def test_setup_detections_multiple(): configs = [ - GibberishFilterConfig(), - RepetitionFilterConfig(), + GibberishDetectionConfig(), + RepetitionDetectionConfig(), ] - 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" + detections = setup_detections(configs, vocab_size=128_000) + assert len(detections) == 2 + assert detections[0].name == "gibberish" + assert detections[1].name == "repetition" -# --- apply_filters tests (enforce=True) --- +# --- run_detections tests (enforce=True) --- -def test_apply_filters_enforced_flags_rollout(): - gibberish_filter = _make_gibberish_filter(enforce=True) +def test_run_detections_enforced_flags_rollout(): + gibberish_detection = _make_gibberish_detection(enforce=True) rollout = _make_rollout( completion_ids=[120_000], - completion_logprobs=[gibberish_filter.logprob_threshold - 1.0], + completion_logprobs=[gibberish_detection.logprob_threshold - 1.0], reward=1.0, ) - apply_filters([gibberish_filter], [rollout]) + run_detections([gibberish_detection], 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 + assert rollout.detections == {"gibberish": True} + assert rollout.is_excluded is True -def test_apply_filters_preserves_clean_rollouts(): - gibberish_filter = _make_gibberish_filter(enforce=True) +def test_run_detections_preserves_clean_rollouts(): + gibberish_detection = _make_gibberish_detection(enforce=True) rollout = _make_rollout( completion_ids=[50, 60, 70], @@ -280,129 +280,131 @@ def test_apply_filters_preserves_clean_rollouts(): reward=1.0, ) - apply_filters([gibberish_filter], [rollout]) + run_detections([gibberish_detection], 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 + assert rollout.detections == {"gibberish": False} + assert rollout.is_excluded 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) +def test_run_detections_first_detection_wins(): + gibberish_detection = _make_gibberish_detection(enforce=True) + repetition_detection = _make_repetition_detection(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], + completion_logprobs=[gibberish_detection.logprob_threshold - 1.0, -0.001, -0.001], reward=1.0, ) - apply_filters([gibberish_filter, repetition_filter], [rollout]) + run_detections([gibberish_detection, repetition_detection], rollout) assert rollout.stop_condition is None - assert rollout.filter_results == {"gibberish": True, "repetition": False} - assert rollout.is_filtered is True + assert rollout.detections == {"gibberish": True, "repetition": False} + assert rollout.is_excluded is True -def test_apply_filters_empty_list(): +def test_run_detections_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 + run_detections([], rollout) + assert rollout.detections == {} + assert rollout.is_excluded is False assert rollout.reward == 1.0 -def test_apply_filters_mixed_batch(): - gibberish_filter = _make_gibberish_filter(enforce=True) +def test_run_detections_mixed_batch(): + gibberish_detection = _make_gibberish_detection(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 + completion_ids=[120_000], completion_logprobs=[gibberish_detection.logprob_threshold - 1.0], reward=1.0 ) - apply_filters([gibberish_filter], [clean, dirty]) + for r in (clean, dirty): + run_detections([gibberish_detection], r) assert clean.reward == 1.0 assert dirty.reward == 1.0 - assert clean.is_filtered is False - assert dirty.is_filtered is True + assert clean.is_excluded is False + assert dirty.is_excluded is True -def test_apply_filters_enforced_preserves_rollout_tokens(): - gibberish_filter = _make_gibberish_filter(enforce=True) +def test_run_detections_enforced_preserves_rollout_tokens(): + gibberish_detection = _make_gibberish_detection(enforce=True) rollout = _make_rollout( completion_ids=[10, 120_000, 30], - completion_logprobs=[-1.0, gibberish_filter.logprob_threshold - 1.0, -0.5], + completion_logprobs=[-1.0, gibberish_detection.logprob_threshold - 1.0, -0.5], reward=1.0, ) - apply_filters([gibberish_filter], [rollout]) + run_detections([gibberish_detection], rollout) assert rollout.nodes[0].token_ids == [10, 120_000, 30] assert rollout.nodes[0].logprobs == [ -1.0, - gibberish_filter.logprob_threshold - 1.0, + gibberish_detection.logprob_threshold - 1.0, -0.5, ] assert rollout.nodes[0].mask == [True, True, True] - assert rollout.is_filtered is True + assert rollout.is_excluded is True -def test_apply_filters_preserves_existing_stop_condition(): - gibberish_filter = _make_gibberish_filter(enforce=True) +def test_run_detections_preserves_existing_stop_condition(): + gibberish_detection = _make_gibberish_detection(enforce=True) rollout = _make_rollout( completion_ids=[120_000], - completion_logprobs=[gibberish_filter.logprob_threshold - 1.0], + completion_logprobs=[gibberish_detection.logprob_threshold - 1.0], reward=1.0, ) rollout.stop_condition = "generation_truncated" - apply_filters([gibberish_filter], [rollout]) + run_detections([gibberish_detection], rollout) assert rollout.stop_condition == "generation_truncated" - assert rollout.is_filtered is True + assert rollout.is_excluded is True -# --- apply_filters tests (monitor-only, enforce=False) --- +# --- run_detections tests (monitor-only, enforce=False) --- -def test_apply_filters_monitor_only_tracks_detection(): - gibberish_filter = _make_gibberish_filter(enforce=False) +def test_run_detections_monitor_only_tracks_detection(): + gibberish_detection = _make_gibberish_detection(enforce=False) rollout = _make_rollout( completion_ids=[120_000], - completion_logprobs=[gibberish_filter.logprob_threshold - 1.0], + completion_logprobs=[gibberish_detection.logprob_threshold - 1.0], reward=1.0, ) - apply_filters([gibberish_filter], [rollout]) + run_detections([gibberish_detection], 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 + assert rollout.detections == {"gibberish": True} + assert rollout.is_excluded is False -def test_apply_filters_monitor_only_mixed_batch(): - gibberish_filter = _make_gibberish_filter(enforce=False) +def test_run_detections_monitor_only_mixed_batch(): + gibberish_detection = _make_gibberish_detection(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 + completion_ids=[120_000], completion_logprobs=[gibberish_detection.logprob_threshold - 1.0], reward=1.0 ) - apply_filters([gibberish_filter], [clean, dirty]) + for r in (clean, dirty): + run_detections([gibberish_detection], r) assert clean.reward == 1.0 assert dirty.reward == 1.0 - assert clean.is_filtered is False - assert dirty.is_filtered is False + assert clean.is_excluded is False + assert dirty.is_excluded is False diff --git a/tests/unit/orchestrator/test_metrics.py b/tests/unit/orchestrator/test_metrics.py index 546e594d9a..566ac99b90 100644 --- a/tests/unit/orchestrator/test_metrics.py +++ b/tests/unit/orchestrator/test_metrics.py @@ -32,8 +32,8 @@ def mk( group_id: str = "g0", trainable: bool = True, is_trainable: bool = True, - is_filtered: bool = False, - filter_results: dict | None = None, + is_excluded: bool = False, + detections: dict | None = None, setup: float = 0.0, agent: float = 0.0, agent_model: float = 0.0, @@ -64,8 +64,8 @@ def mk( group_id=group_id, agent=SimpleNamespace(trainable=trainable, name=agent_name), is_trainable=is_trainable, - is_filtered=is_filtered, - filter_results=filter_results or {}, + is_excluded=is_excluded, + detections=detections or {}, timing=SimpleNamespace( setup=SimpleNamespace(duration=setup), agent=SimpleNamespace( @@ -93,12 +93,12 @@ 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")] + [mk(env_name="a"), mk(env_name="a", has_error=True), mk(env_name="b", is_excluded=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 assert isinstance(eff, TrainRollouts) and len(eff) == 2 # same type, errored + filtered dropped - assert all(not r.has_error and not r.is_filtered and r in rc.rollouts for r in eff) # view of references + assert all(not r.has_error and not r.is_excluded 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()) @@ -162,7 +162,7 @@ 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([mk(is_truncated=True), mk(has_error=True, error_type="ProviderError"), mk(is_excluded=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 @@ -234,16 +234,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_excluded=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/is_excluded/mean"] == 0.5 + assert out["train/agg/all/agent/detections/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) + assert not any("is_trainable" in k or "is_excluded" in k or "/detections/" in k for k in eval_out) def test_eval_avg_at_k_and_pass_k(): diff --git a/tests/unit/orchestrator/test_task_stats.py b/tests/unit/orchestrator/test_task_stats.py new file mode 100644 index 0000000000..e81609a7f2 --- /dev/null +++ b/tests/unit/orchestrator/test_task_stats.py @@ -0,0 +1,78 @@ +import verifiers.v1 as vf + +from prime_rl.orchestrator.task_stats import DECAY, PRIOR, TaskStats, task_key +from prime_rl.orchestrator.types import Rollout + + +def make_rollout( + *, reward: float | None = None, env_name: str = "env", idx: int = 0, ok: bool = True, trainable: bool = True +) -> Rollout: + rollout = Rollout( + task=vf.TraceTask(type="Task", data=vf.TaskData(idx=idx, prompt=f"task {idx}")), + agent=vf.AgentInfo(config=vf.AgentConfig(), trainable=trainable), + ok=ok, + ) + if not ok: + rollout.errors = [vf.Error(type="TestError", message="boom")] + if reward is not None: + rollout.rewards = {"main": vf.Reward(score=reward, weight=1.0)} + rollout.env_name = env_name + return rollout + + +def test_task_key_is_canonical_and_content_sensitive(): + assert task_key({"a": 1, "b": 2}) == task_key({"b": 2, "a": 1}) + assert task_key({"a": 1}) != task_key({"a": 2}) + + +def test_observe_accumulates_discounted_evidence_per_role(): + stats = TaskStats() + group = [make_rollout(reward=1.0), make_rollout(reward=0.0)] + stats.observe(group) + + key = task_key(group[0].task.data.model_dump(mode="json")) + stat = stats.stats["env"][key]["agent"] + assert stat.s == 1.0 and stat.f == 1.0 and stat.visits == 1 + alpha, beta = PRIOR + assert stat.p_hat == (alpha + 1.0) / (alpha + beta + 2.0) + + # A second all-success group: prior counts decay, new evidence lands whole. + stats.observe([make_rollout(reward=1.0), make_rollout(reward=1.0)]) + assert stat.s == DECAY * 1.0 + 2.0 and stat.f == DECAY * 1.0 + assert stat.visits == 2 + + +def test_errored_rollouts_update_tick_counters_but_not_evidence(): + stats = TaskStats() + stats.observe([make_rollout(ok=False), make_rollout(ok=False)]) + assert stats.stats == {} + metrics = stats.metrics({"env": 10}) + assert metrics["sampler/env/groups_observed"] == 1.0 + assert metrics["sampler/env/realized_signal_rate"] == 0.0 + assert metrics["sampler/env/pool/unseen"] == 10.0 + # Tick counters drain on read; snapshot metrics persist. + assert "sampler/env/groups_observed" not in stats.metrics({"env": 10}) + + +def test_signal_rate_reads_nonzero_advantages(): + stats = TaskStats() + signal = [make_rollout(reward=1.0), make_rollout(reward=0.0)] + signal[0].advantages = [0.5] + signal[1].advantages = [-0.5] + degenerate = [make_rollout(reward=1.0, idx=1), make_rollout(reward=1.0, idx=1)] + for rollout in degenerate: + rollout.advantages = [0.0] + stats.observe(signal) + stats.observe(degenerate) + metrics = stats.metrics({"env": None}) + assert metrics["sampler/env/realized_signal_rate"] == 0.5 + assert "sampler/env/coverage" not in metrics # infinite taskset + + +def test_state_dict_roundtrip(): + stats = TaskStats() + stats.observe([make_rollout(reward=1.0), make_rollout(reward=0.0)]) + restored = TaskStats() + restored.load_state_dict(stats.state_dict()) + key = task_key(make_rollout().task.data.model_dump(mode="json")) + assert restored.stats["env"][key]["agent"] == stats.stats["env"][key]["agent"]