diff --git a/configs/ci/nightly-fft/wiki-search.toml b/configs/ci/nightly-fft/wiki-search.toml index ffdc1b1394..49c534e9a1 100644 --- a/configs/ci/nightly-fft/wiki-search.toml +++ b/configs/ci/nightly-fft/wiki-search.toml @@ -19,10 +19,6 @@ batch_size = 512 group_size = 16 oversampling_factor = 2.0 -[[orchestrator.pre_batch_filters]] -type = "zero_advantage" -enforce = true - [[orchestrator.train.source]] name = "wiki-search" diff --git a/configs/debug/algo/echo.toml b/configs/debug/algo/echo.toml index 62d562b735..a9bd418436 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 92958ddccd..efb77a03b3 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, curricula, 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) +- [Curricula](#curricula) - [Multi-Turn Trajectories](#multi-turn-trajectories) - [Extension Property](#extension-property) - [Best-Effort Interleaving](#best-effort-interleaving) @@ -164,12 +164,12 @@ At runtime, each env's resolved config builds two objects: a `Sampler` (`prime_r | `hierarchical_grpo` | `HierarchicalGRPOAlgorithm` | `score_group`: GRPO baseline per episode for solvers, per group for the proposer | | `opd` | `OPDAlgorithm` | `score_rollout`: own-context prefill under the teacher | | `opsd` | `OPSDAlgorithm` | `score_rollout`: demo-conditioned prefill under the live policy | -| `sft` | `SFTDistillAlgorithm` | `score_group`: group-norm credit (feeds filters) | +| `sft` | `SFTDistillAlgorithm` | no credit assignment; CE on sampled tokens | 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. +- `score_group(group)` — the completed 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` 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` 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. 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. Curriculum admission and metrics can inspect these streams after group scoring. 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,43 @@ type = "opsd" demo_key = "demonstration" ``` -Scoring runs at arrival, *before* the pre-batch filters, so a rollout that is later filtered still cost its reference compute — accepted for the simpler one-rollout-at-a-time shape (advantage-based filters never fire for opd/opsd anyway, since neither assigns an advantage). +Scoring runs at arrival, before curriculum admission, so a rollout that is later rejected still costs its reference compute. -## Filters +By default, zero-advantage RL tokens are removed after a complete batch cohort has been collected and before its payload is packed for the trainer. This does not backfill the batch. Samples that still carry CE or reference-KL components are retained, while pure zero-advantage RL samples are not shipped. Removing an RL token also removes its trainer/inference mismatch-KL contribution. Set `orchestrator.train.filter_zero_advantages = false` to retain them. -Filters drop rollouts between scoring and training. Built-ins (composable): +## Curricula -| Filter | Effect | -|---|---| -| `gibberish` | Drops rollouts whose mean log-prob fall below a threshold — usually a sign of degenerate output. | -| `repetition` | Drops rollouts with high n-gram repetition. | -| `zero_advantage` | Drops rollouts whose advantage is zero, so the trainer doesn't waste tokens on them. | +Each training source has a `Curriculum` composed from one `TaskSampler` and any number of named `AdmissionGate`s. The sampler chooses tasks and observes every finalized result. Every gate evaluates every result; the group trains only if every gate admits it. Rejected groups remain observable while the orchestrator samples again to fill the batch. -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: +Samplers and gates can be stateful. Their `state_dict`, `load_state_dict`, and `metrics` methods are included in orchestrator checkpoints and logged under `curriculum//`. + +Three small implementations are included: + +- `StandardSampler` is the default: it advances the task iterator and cycles finite tasksets in source order. +- `DifficultyPoolSampler` samples finite tasksets with replacement and tracks each task's latest valid mean group reward. Each named pool has an inclusive reward threshold and a relative per-task sampling weight; weight `0` disables sampling from that pool. Unseen tasks use neutral weight `1.0`, so pool observations affect sampling immediately without waiting for a full taskset pass. +- `AdvRangeGate` rejects a group when every trainable-token advantage falls inside `reject_min` through `reject_max`. Unlike the built-in post-batch zero-advantage filtering, rejection requests replacement work. Groups without an advantage stream are admitted. ```toml -[[orchestrator.post_batch_filters]] -type = "zero_advantage" +[orchestrator.train.source.curriculum.sampler] +type = "difficulty_pool" -[[orchestrator.post_batch_filters]] -type = "repetition" -threshold = 0.4 -``` +[orchestrator.train.source.curriculum.sampler.pools.hard] +threshold = 0.25 +weight = 0.2 -Filtered rollouts still appear in the W&B metrics, just not in the trainer batch — useful for spotting whether filtering is doing its job. +[orchestrator.train.source.curriculum.sampler.pools.normal] +threshold = 0.75 +weight = 1.0 + +[orchestrator.train.source.curriculum.sampler.pools.easy] +threshold = 1.0 +weight = 0.2 + +[orchestrator.train.source.curriculum.gates.low_signal] +type = "advantage_range" +reject_min = -0.05 +reject_max = 0.05 +``` ## Multi-Turn Trajectories diff --git a/examples/advanced/glm-5.2/swe.toml b/examples/advanced/glm-5.2/swe.toml index e02257486f..492c9384d8 100644 --- a/examples/advanced/glm-5.2/swe.toml +++ b/examples/advanced/glm-5.2/swe.toml @@ -86,11 +86,6 @@ id = "bash" [orchestrator.eval.source.env.agent.runtime] labels = ["glm5-pd-disag", "swe-bench-verified"] -[[orchestrator.post_batch_filters]] -type = "gibberish" -enforce = true - [inference] # we need <0.85 bc glm5 layers are too large for 0.85 use_deep_gemm = true - diff --git a/examples/advanced/intellect-3.1/rl.toml b/examples/advanced/intellect-3.1/rl.toml index d6c2ef2b5e..b7ba405f88 100644 --- a/examples/advanced/intellect-3.1/rl.toml +++ b/examples/advanced/intellect-3.1/rl.toml @@ -50,6 +50,7 @@ oversampling_factor = 2 [[orchestrator.train.source]] name = "swe" ratio = 0.3 +curriculum = { gates = { zero_advantage = { type = "advantage_range" } } } [orchestrator.train.source.env.taskset] id = "r2e-gym" @@ -63,6 +64,7 @@ labels = ["intellect-3.1", "swe"] [[orchestrator.train.source]] name = "deepdive" ratio = 0.2 +curriculum = { gates = { zero_advantage = { type = "advantage_range" } } } [orchestrator.train.source.env.taskset] id = "deepdive" @@ -78,6 +80,7 @@ labels = ["intellect-3.1", "deepdive"] [[orchestrator.train.source]] name = "math" ratio = 0.3 +curriculum = { gates = { zero_advantage = { type = "advantage_range" } } } [orchestrator.train.source.env.taskset] id = "i3_math" @@ -91,6 +94,7 @@ type = "subprocess" [[orchestrator.train.source]] name = "logic" ratio = 0.2 +curriculum = { gates = { zero_advantage = { type = "advantage_range" } } } [orchestrator.train.source.env.taskset] id = "i3_logic" @@ -107,6 +111,7 @@ type = "subprocess" [[orchestrator.train.source]] name = "code" ratio = 0.2 +curriculum = { gates = { zero_advantage = { type = "advantage_range" } } } [orchestrator.train.source.env.taskset] id = "i3_code" @@ -117,10 +122,6 @@ id = "null" [orchestrator.train.source.env.agent.runtime] type = "subprocess" -[[orchestrator.pre_batch_filters]] -type = "zero_advantage" -enforce = true - [orchestrator.eval] interval = 25 diff --git a/examples/basic/wiki-search/rl.toml b/examples/basic/wiki-search/rl.toml index f704fa6892..b6085e6078 100644 --- a/examples/basic/wiki-search/rl.toml +++ b/examples/basic/wiki-search/rl.toml @@ -42,6 +42,7 @@ max_completion_tokens = 512 [[orchestrator.train.source]] name = "wiki-search" +curriculum = { gates = { zero_advantage = { type = "advantage_range" } } } [orchestrator.train.source.env.taskset] id = "wiki-search" @@ -52,10 +53,6 @@ id = "null" [orchestrator.train.source.env.agent.runtime] type = "subprocess" -[[orchestrator.pre_batch_filters]] -type = "zero_advantage" -enforce = true - [ckpt] # Checkpoint at the end of training [inference.vllm] diff --git a/packages/prime-rl-configs/src/prime_rl/configs/algorithm.py b/packages/prime-rl-configs/src/prime_rl/configs/algorithm.py index 18f13ac82a..739cd6e4bb 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/algorithm.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/algorithm.py @@ -234,8 +234,7 @@ class MaxRLAlgoConfig(BaseAlgoConfig): objective: low-pass-rate examples get ~1/p weight, and ``group_size`` is the truncation order interpolating REINFORCE (1) → exact maximum likelihood (∞). Designed for non-negative (canonically binary) rewards; - a group with mean reward 0 carries zero advantages everywhere (the - zero-advantage filter drops it, matching the paper's K=0 convention).""" + a group with mean reward 0 carries zero advantages everywhere.""" action_loss_type: ClassVar[ActionLossType] = "rl" @@ -300,8 +299,8 @@ class OPDAlgoConfig(BaseAlgoConfig): a reference model, evaluated in the trainer from reference prefill logprobs scored over each sample's own context (``ref_logprobs`` on the wire, ``ref_kl`` loss component). No scalar advantage is assigned — - rollouts keep ``advantages=None`` (advantage-based filters never fire) and - samples ship no advantage stream. ``group_size`` only fans out sampling.""" + rollouts keep ``advantages=None`` and samples ship no advantage stream. + ``group_size`` only fans out sampling.""" action_loss_type: ClassVar[ActionLossType] = "ref_kl" @@ -322,8 +321,7 @@ class OPSDAlgoConfig(BaseAlgoConfig): prepended as a leading system message. The sample is scored verbatim (no re-rendering), so it's robust to tool/multimodal prompts and works for any number of turns. No scalar advantage is assigned — rollouts keep - ``advantages=None`` (advantage-based filters never fire) and samples ship no - advantage stream.""" + ``advantages=None`` and samples ship no advantage stream.""" action_loss_type: ClassVar[ActionLossType] = "ref_kl" @@ -347,8 +345,8 @@ class SFTAlgoConfig(BaseAlgoConfig): type: Literal["sft"] = "sft" """SFT distillation: cross-entropy on the sampled tokens. The ``ce`` loss ignores advantages and SFT assigns none — it trains on every sampled token. - Reward-based filtering, if wanted, is an explicit filter, not smuggled - through an unused advantage stream.""" + A curriculum can reject results using reward or any other finalized + rollout data.""" action_loss_type: ClassVar[ActionLossType] = "ce" 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 7ef9eeb9f6..e316a2746b 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py @@ -157,6 +157,77 @@ def validate_env(self): return self +class StandardSamplerConfig(BaseConfig): + type: Literal["standard"] = "standard" + + +class DifficultyPoolConfig(BaseConfig): + threshold: float + """Inclusive maximum reward assigned to this pool.""" + + weight: float = Field(ge=0) + """Relative per-task sampling weight.""" + + +def default_difficulty_pools() -> dict[str, DifficultyPoolConfig]: + return { + "hard": DifficultyPoolConfig(threshold=0.25, weight=0.2), + "normal": DifficultyPoolConfig(threshold=0.75, weight=1.0), + "easy": DifficultyPoolConfig(threshold=1.0, weight=0.2), + } + + +class DifficultyPoolSamplerConfig(BaseConfig): + type: Literal["difficulty_pool"] = "difficulty_pool" + + pools: dict[str, DifficultyPoolConfig] = Field(default_factory=default_difficulty_pools) + """Named pools ordered by their reward thresholds.""" + + seed: int = 42 + + @model_validator(mode="after") + def validate_pools(self): + if not self.pools: + raise ValueError("DifficultyPoolSampler requires at least one pool") + thresholds = [pool.threshold for pool in self.pools.values()] + if len(set(thresholds)) != len(thresholds): + raise ValueError("Difficulty pool thresholds must be unique") + if not any(pool.weight > 0 for pool in self.pools.values()): + raise ValueError("At least one difficulty pool must have a positive weight") + return self + + +TaskSamplerConfig: TypeAlias = Annotated[ + StandardSamplerConfig | DifficultyPoolSamplerConfig, + Field(discriminator="type"), +] + + +class AdvRangeGateConfig(BaseConfig): + type: Literal["advantage_range"] = "advantage_range" + + reject_min: float = 0.0 + reject_max: float = 0.0 + + @model_validator(mode="after") + def validate_range(self): + if self.reject_min > self.reject_max: + raise ValueError("reject_min must be less than or equal to reject_max") + return self + + +AdmissionGateConfig: TypeAlias = AdvRangeGateConfig + + +class CurriculumConfig(BaseConfig): + sampler: TaskSamplerConfig = Field(default_factory=StandardSamplerConfig) + """Task selection policy. The default cycles through the task iterator in source order.""" + + gates: dict[str, AdmissionGateConfig] = Field(default_factory=dict) + """Named admission policies. Every gate observes every finalized group, + and a group trains only when every gate admits it.""" + + class TrainSourceConfig(EnvConfig): sampling: TrainSamplingConfig = TrainSamplingConfig() """Per-env sampling overrides. Unset fields inherit from the group-level train sampling config.""" @@ -170,6 +241,10 @@ class TrainSourceConfig(EnvConfig): ``orchestrator.algo`` when unset; set ``type`` (and its params) to give this env its own algorithm.""" + curriculum: CurriculumConfig | None = None + """User-authored task sampler and admission gates. The default cycles + through the taskset and admits every finalized group.""" + class EvalSourceConfig(EnvConfig): sampling: EvalSamplingConfig = EvalSamplingConfig() @@ -192,6 +267,9 @@ class TrainConfig(BaseConfig): sampling: TrainSamplingConfig = TrainSamplingConfig() """Shared training sampling configuration.""" + filter_zero_advantages: bool = True + """Remove zero-advantage RL tokens after collecting a batch, before shipping it.""" + @model_validator(mode="after") def resolve_env_defaults(self): """Resolve per-env overrides: inherit group-level sampling (the worker ``pool`` @@ -298,50 +376,6 @@ class CheckpointConfig(BaseConfig): """Skip loading the progress from checkpoint.""" -# Flags rare tokens generated at high entropy (Section 5.2, https://arxiv.org/abs/2510.02387). -class GibberishFilterConfig(BaseConfig): - type: Literal["gibberish"] = "gibberish" - - enforce: bool = False - """When True, skip detected rollouts entirely so they are not sent to the trainer. When False, only track detection metrics.""" - - token_id_threshold: int = 100_000 - """Token IDs above this are candidates for gibberish. BPE tokens are sorted by merge order.""" - - logprob_offset: float = 2.0 - """Offset from uniform-distribution logprob. Threshold = ``-log(vocab_size) - logprob_offset``.""" - - -# Flags rollouts stuck in a repetition loop: emits high-confidence tokens for an extended stretch. -# Flagged when `window` consecutive tokens are each sampled with probability above `prob_threshold`. -# (Section 3.2, https://arxiv.org/abs/2506.13585) -class RepetitionFilterConfig(BaseConfig): - type: Literal["repetition"] = "repetition" - - enforce: bool = False - """When True, skip detected rollouts entirely so they are not sent to the trainer. When False, only track detection metrics.""" - - window: int = Field(3_000, ge=1) - """Consecutive high-probability steps required to flag the rollout.""" - - prob_threshold: float = Field(0.99, gt=0, le=1) - """Tokens sampled with probability above this are considered repetitive. Consecutive such tokens count toward the window.""" - - -# Flags rollouts with zero advantage. -class ZeroAdvantageFilterConfig(BaseConfig): - type: Literal["zero_advantage"] = "zero_advantage" - - enforce: bool = True - """When True, skip detected rollouts entirely so they are not sent to the trainer. When False, only track detection metrics.""" - - -FilterConfig: TypeAlias = Annotated[ - GibberishFilterConfig | RepetitionFilterConfig | ZeroAdvantageFilterConfig, - Field(discriminator="type"), -] - - class FileSystemWeightBroadcastConfig(BaseConfig): type: Literal["filesystem"] = "filesystem" @@ -412,24 +446,6 @@ class OrchestratorConfig(BaseConfig): eval: EvalConfig | None = None """Evaluation configuration.""" - pre_batch_filters: list[FilterConfig] = [ - GibberishFilterConfig(enforce=False), - RepetitionFilterConfig(enforce=False), - ZeroAdvantageFilterConfig(enforce=False), - ] - """Filters applied *before* a rollout enters the training batch buffer. - All three filter types are registered in monitor mode by default; flip ``enforce=true`` per type - to drop matching rollouts before they consume a slot in the batch (e.g. a zero-advantage group - never makes it into a training batch).""" - - post_batch_filters: list[FilterConfig] = [ - GibberishFilterConfig(), - RepetitionFilterConfig(), - ZeroAdvantageFilterConfig(), - ] - """Filters applied *after* a batch has been assembled. Each filter annotates each rollout; - rollouts flagged by an enforcing filter are still recorded but not shipped to the trainer.""" - log: LogConfig = LogConfig() env_vars: EnvVars = {} @@ -522,16 +538,6 @@ def auto_setup_prime_monitor_name(self): self.monitors.prime.name = self.monitors.wandb.name return self - @model_validator(mode="after") - def validate_unique_filter_types(self): - for slot_name in ("pre_batch_filters", "post_batch_filters"): - types = [f.type for f in getattr(self, slot_name)] - if len(types) != len(set(types)): - raise ValueError( - f"Duplicate filter types in {slot_name}: {types}. Each filter type may only appear once per slot." - ) - return self - @model_validator(mode="after") def inherit_env_algorithms(self): """Envs without their own algorithm inherit the top-level one. diff --git a/pyproject.toml b/pyproject.toml index 1731d3a5e3..682efe01e0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -33,7 +33,7 @@ dependencies = [ "setproctitle>=1.3.0", "uvloop>=0.21.0", "torchtitan", - "verifiers[harbor]>=0.3.1.dev14", + "verifiers[harbor]>=0.3.1.dev33", "renderers", "dion", "tilelang>=0.1.8", diff --git a/skills/training/monitor-run/SKILL.md b/skills/training/monitor-run/SKILL.md index 5d31d0c19d..6a1e9e4898 100644 --- a/skills/training/monitor-run/SKILL.md +++ b/skills/training/monitor-run/SKILL.md @@ -92,9 +92,9 @@ All metrics print to the console log (and W&B when configured). **Progress** — orchestrator log. Rollout metrics mirror the episode/trace hierarchy, at two levels: - `{scope}/{subset}//` — episode-level facts only: the token/turn/branch counts, summed over an episode's traces. -- `{scope}/{subset}///` — every trace-level metric (reward, truncation, errors, timing, env metrics, filter verdicts, eval scores), keyed by agent name so seats never mix. Flat over that agent's traces: one sample is one trace, so an in-episode fan-out like n solvers contributes n samples. +- `{scope}/{subset}///` — every trace-level metric (reward, truncation, errors, timing, env metrics, curriculum admission, eval scores), keyed by agent name so seats never mix. Flat over that agent's traces: one sample is one trace, so an in-episode fan-out like n solvers contributes n samples. -`scope` is `train/agg` (all train envs) or `train/` (`eval/` for eval); `subset` is `all` (every rollout) or `effective` (post-filter). Single-agent envs have one agent — usually `agent` — and one trace per episode, so both levels agree; multi-agent envs name each seat (`proposer`, `solver`, `judge`, …). +`scope` is `train/agg` (all train envs) or `train/` (`eval/` for eval); `subset` is `all` (every rollout) or `effective` (admitted, clean, and trainable). Single-agent envs have one agent — usually `agent` — and one trace per episode, so both levels agree; multi-agent envs name each seat (`proposer`, `solver`, `judge`, …). | Metric | Description | |--------|-------------| @@ -104,7 +104,8 @@ All metrics print to the console log (and W&B when configured). | `train//effective//num_turns/mean` | avg turns for that agent alone (also token counts, `num_branches`) | | `train/agg/effective//is_truncated/mean` | fraction of that agent's rollouts truncated | | `train/agg/all//has_error/mean` | fraction of that agent's rollouts errored (per-type under `train/agg/all//error/`; also `dispatcher/errored/{train,eval}`) | -| `train/agg/all//is_trainable/mean` | fraction carrying a training signal — 0.0 for a frozen seat like a judge (also `is_filtered`, `filters/`) | +| `train/agg/all//is_trainable/mean` | fraction carrying a training signal — 0.0 for a frozen seat like a judge | +| `train/agg/all//is_admitted/mean` | fraction accepted by the source curriculum; per-source counters and custom policy metrics live under `curriculum//` | | `train//effective//metrics//mean` | env-specific metrics for that agent (e.g. pass rate) | | `train//effective//timing/agent/model/mean` | model vs harness share of that agent's phase | | `eval//effective//{avg@k,pass@k}` | eval scores for that agent, when configured | @@ -147,7 +148,7 @@ curl -s http://localhost:8100/metrics | grep -E "num_requests|gpu_cache_usage" JSONL files of `vf.Trace` records (training tensors excluded), one line per trace — a multi-agent env's episode contributes several lines sharing one `info.episode_id`. `all` -gets every completed rollout the moment it arrives — errored, filtered, and never-batched +gets every completed rollout the moment it arrives — errored, curriculum-rejected, and never-batched ones included — so it's crash-durable; `effective` gets the clean trainable subset that went into the step's train batch (eval: the non-errored trainable epoch cohort; multiple eval envs share the step file) — untrainable traces (a frozen judge's) appear only in `all`. Each record carries `run` (`{type, id, step}`; for eval, `step` is the trigger step), diff --git a/src/prime_rl/orchestrator/algo/base.py b/src/prime_rl/orchestrator/algo/base.py index 3540864277..b7161f5125 100644 --- a/src/prime_rl/orchestrator/algo/base.py +++ b/src/prime_rl/orchestrator/algo/base.py @@ -19,8 +19,8 @@ I/O against another model — an inference pool the algorithm connected in ``setup()`` (a frozen teacher) or the live policy (opsd's self-distillation), queried with bounded concurrency. No siblings. -- ``score_group(group)`` — the cohort, on group completion, *before* filtering - (filters read the streams): group-relative credit (GRPO/MaxRL baselines). +- ``score_group(group)`` — the cohort, on group completion: group-relative + credit (GRPO/MaxRL baselines). How rollouts are *produced* is not the algorithm's concern: that is the env's :class:`~prime_rl.orchestrator.sampler.Sampler`, and sample construction @@ -95,21 +95,18 @@ class Algorithm: directly — read the trace, write credit via :meth:`Rollout.assign_advantages`. They are async so either stage may do I/O — e.g. a process-reward model or a - teacher at arrival, or a judge at group time whose signal a pre-batch - filter then reads; a hook that only does advantage math simply never - awaits. + teacher at arrival, or a judge at group time; a hook that only does + advantage math simply never awaits. - :meth:`score_rollout` — one rollout, on arrival: rollout-local credit, observation ce weights, or per-token results from a model the algorithm connected in :meth:`setup` (e.g. teacher reference logprobs). Default: nothing. - - :meth:`score_group` — the cohort, *before* filtering (filters read the - streams): group-relative credit. Default: nothing — rollouts keep - ``advantages=None``, so advantage-based filters skip them. + - :meth:`score_group` — the cohort: group-relative credit. Default: + nothing — rollouts keep ``advantages=None``. - Model I/O lives in :meth:`score_rollout`: it runs at arrival, *before* the - pre-batch filters, so it pays compute on rollouts that may then be filtered - out — accepted for the simpler one-rollout-at-a-time shape. + Model I/O lives in :meth:`score_rollout`: it runs at arrival, so it pays + compute on rollouts that the curriculum may later reject. Constructed with the algorithm config it interprets plus the live policy pool (``self.policy_pool`` — always available, never closed by the @@ -145,8 +142,7 @@ async def score_rollout(self, rollout: Rollout) -> None: group stats.""" async def score_group(self, group: list[Rollout]) -> None: - """Group phase, the finalized cohort, before filtering: write - group-relative credit.""" + """Group phase over the finalized cohort: write group-relative credit.""" async def finalize_rollout(self, rollout: Rollout) -> None: """Arrival phase (non-virtual): rollout-local scoring as each rollout is diff --git a/src/prime_rl/orchestrator/algo/max_rl.py b/src/prime_rl/orchestrator/algo/max_rl.py index 9a3978108d..279a39e468 100644 --- a/src/prime_rl/orchestrator/algo/max_rl.py +++ b/src/prime_rl/orchestrator/algo/max_rl.py @@ -20,8 +20,7 @@ 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.""" 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/algo/opd.py b/src/prime_rl/orchestrator/algo/opd.py index a4b3a06f72..8dc8b513e4 100644 --- a/src/prime_rl/orchestrator/algo/opd.py +++ b/src/prime_rl/orchestrator/algo/opd.py @@ -19,9 +19,8 @@ class OPDAlgorithm(Algorithm): The policy samples its own rollouts; at ship time each sample's full context is prefill-scored under the teacher (``ref_logprobs`` on the wire), and the trainer evaluates the KL against the live policy. No - credit is assigned — rollouts keep ``advantages=None`` (advantage-based - filters never fire) and samples ship no advantage stream; ``group_size`` - only fans out sampling.""" + credit is assigned — rollouts keep ``advantages=None`` and samples ship no + advantage stream; ``group_size`` only fans out sampling.""" action_loss_type = "ref_kl" diff --git a/src/prime_rl/orchestrator/algo/sft.py b/src/prime_rl/orchestrator/algo/sft.py index c8c51221f3..9f1aa91d33 100644 --- a/src/prime_rl/orchestrator/algo/sft.py +++ b/src/prime_rl/orchestrator/algo/sft.py @@ -8,7 +8,7 @@ class SFTDistillAlgorithm(Algorithm): rollouts (``sampling.source``); the policy trains with CE on its tokens. Assigns no advantage — the ``ce`` loss ignores credit, and SFT trains on - every sampled token. Reward-based filtering, if wanted, is an explicit - filter, not smuggled through an unused advantage stream.""" + every sampled token. A curriculum can reject results using reward or any + other finalized rollout data.""" action_loss_type = "ce" diff --git a/src/prime_rl/orchestrator/ckpt.py b/src/prime_rl/orchestrator/ckpt.py index ef85dcdc9e..1dd4a4191a 100644 --- a/src/prime_rl/orchestrator/ckpt.py +++ b/src/prime_rl/orchestrator/ckpt.py @@ -1,5 +1,4 @@ -"""Checkpoint manager for the orchestrator state (``Progress`` counters + -``TrainSource`` data position). Layout: +"""Checkpoint manager for orchestrator progress and train-source state. Layout: ``/checkpoints/step_N/orchestrator/progress.pt``.""" from __future__ import annotations @@ -57,7 +56,7 @@ def load(self, progress: Progress, train_source: TrainSource, step: int, path: P get_logger().debug(f"Loading checkpoint from {state_file}") start = time.perf_counter() if self.config.skip_progress: - get_logger().info("Skipping progress and data position loading from checkpoint") + get_logger().info("Skipping progress and train source loading from checkpoint") else: with open(state_file, "rb") as f: state = torch.load(f, weights_only=False) @@ -66,15 +65,9 @@ def load(self, progress: Progress, train_source: TrainSource, step: int, path: P if hasattr(progress, key): setattr(progress, key, value) train_source.load_state_dict(state["train_source"]) - for name, position in state["train_source"]["envs"].items(): - if name not in train_source.base_rows: - continue - rows = train_source.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']}, " - f"cursor={position['cursor']}/{num_tasks}" - ) + for name in state["train_source"]["envs"]: + if name in train_source.curricula: + get_logger().info(f"Resumed curriculum state for env {name}") get_logger().debug(f"Orchestrator checkpoint loaded in {format_time(time.perf_counter() - start)}") diff --git a/src/prime_rl/orchestrator/curriculum/__init__.py b/src/prime_rl/orchestrator/curriculum/__init__.py new file mode 100644 index 0000000000..40c7e829db --- /dev/null +++ b/src/prime_rl/orchestrator/curriculum/__init__.py @@ -0,0 +1,14 @@ +"""Task selection and finalized-sample admission for training environments.""" + +from prime_rl.orchestrator.curriculum.base import Curriculum +from prime_rl.orchestrator.curriculum.gates import AdmissionGate, AdvRangeGate +from prime_rl.orchestrator.curriculum.samplers import DifficultyPoolSampler, StandardSampler, TaskSampler + +__all__ = [ + "AdmissionGate", + "AdvRangeGate", + "Curriculum", + "DifficultyPoolSampler", + "StandardSampler", + "TaskSampler", +] diff --git a/src/prime_rl/orchestrator/curriculum/base.py b/src/prime_rl/orchestrator/curriculum/base.py new file mode 100644 index 0000000000..bcebec3b65 --- /dev/null +++ b/src/prime_rl/orchestrator/curriculum/base.py @@ -0,0 +1,89 @@ +"""Curriculum composition and lifecycle.""" + +from __future__ import annotations + +from collections.abc import Iterator, Sequence +from typing import TYPE_CHECKING, Any + +import verifiers.v1 as vf + +from prime_rl.orchestrator.curriculum.gates import AdmissionGate +from prime_rl.orchestrator.curriculum.samplers import TaskSampler + +if TYPE_CHECKING: + from prime_rl.configs.orchestrator import CurriculumConfig + from prime_rl.orchestrator.types import Rollout + + +class Curriculum: + """One task sampler composed with zero or more admission gates.""" + + def __init__( + self, + config: CurriculumConfig | None, + tasks: Sequence[vf.Task] | Iterator[vf.Task], + ) -> None: + from prime_rl.configs.orchestrator import ( + AdvRangeGateConfig, + CurriculumConfig, + DifficultyPoolSamplerConfig, + StandardSamplerConfig, + ) + from prime_rl.orchestrator.curriculum.gates import AdvRangeGate + from prime_rl.orchestrator.curriculum.samplers import DifficultyPoolSampler, StandardSampler + + config = CurriculumConfig() if config is None else config + if isinstance(config.sampler, StandardSamplerConfig): + self.sampler: TaskSampler = StandardSampler(tasks) + elif isinstance(config.sampler, DifficultyPoolSamplerConfig): + self.sampler = DifficultyPoolSampler(config.sampler, tasks) + else: + raise TypeError(f"Unsupported task sampler config: {type(config.sampler).__name__}") + + self.gates: dict[str, AdmissionGate] = {} + for name, gate_config in config.gates.items(): + if isinstance(gate_config, AdvRangeGateConfig): + gate: AdmissionGate = AdvRangeGate(gate_config) + else: + raise TypeError(f"Unsupported admission gate config: {type(gate_config).__name__}") + self.gates[name] = gate + + def on_result(self, group: list[Rollout]) -> bool: + """Observe every result, evaluate every gate, and combine with AND.""" + if not group: + raise ValueError("Cannot report an empty rollout group") + task_keys = {rollout.task.key for rollout in group} + if None in task_keys: + raise ValueError("A finalized group is missing Task.key") + if len(task_keys) != 1: + raise ValueError(f"A finalized group contains multiple task keys: {task_keys}") + self.sampler.observe(group) + decisions: list[bool] = [] + for name, gate in self.gates.items(): + decision = gate.admit(group) + if not isinstance(decision, bool): + raise TypeError(f"AdmissionGate {name!r}.admit() must return bool, got {type(decision).__name__}") + decisions.append(decision) + return all(decisions) + + def state_dict(self) -> dict[str, Any]: + return { + "sampler": self.sampler.state_dict(), + "gates": {name: gate.state_dict() for name, gate in self.gates.items()}, + } + + def load_state_dict(self, state_dict: dict[str, Any]) -> None: + if "sampler" not in state_dict: + self.sampler.load_state_dict(state_dict) + return + self.sampler.load_state_dict(state_dict["sampler"]) + for name, gate_state in state_dict["gates"].items(): + gate = self.gates.get(name) + if gate is not None: + gate.load_state_dict(gate_state) + + def metrics(self) -> dict[str, float]: + metrics = {f"sampler/{name}": float(value) for name, value in self.sampler.metrics().items()} + for gate_name, gate in self.gates.items(): + metrics |= {f"gate/{gate_name}/{name}": float(value) for name, value in gate.metrics().items()} + return metrics diff --git a/src/prime_rl/orchestrator/curriculum/gates/__init__.py b/src/prime_rl/orchestrator/curriculum/gates/__init__.py new file mode 100644 index 0000000000..3c5176dcad --- /dev/null +++ b/src/prime_rl/orchestrator/curriculum/gates/__init__.py @@ -0,0 +1,6 @@ +"""Admission gate interfaces and implementations.""" + +from prime_rl.orchestrator.curriculum.gates.adv import AdvRangeGate +from prime_rl.orchestrator.curriculum.gates.base import AdmissionGate + +__all__ = ["AdmissionGate", "AdvRangeGate"] diff --git a/src/prime_rl/orchestrator/curriculum/gates/adv.py b/src/prime_rl/orchestrator/curriculum/gates/adv.py new file mode 100644 index 0000000000..2074a44565 --- /dev/null +++ b/src/prime_rl/orchestrator/curriculum/gates/adv.py @@ -0,0 +1,33 @@ +"""Advantage-based training-sample admission.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from prime_rl.orchestrator.curriculum.gates.base import AdmissionGate + +if TYPE_CHECKING: + from prime_rl.configs.orchestrator import AdvRangeGateConfig + from prime_rl.orchestrator.types import Rollout + + +class AdvRangeGate(AdmissionGate): + """Reject groups whose trainable-token advantages all fall inside a range. + + The default ``[0, 0]`` interval filters groups with no online learning + signal. Groups without an advantage stream are admitted. + """ + + def __init__(self, config: AdvRangeGateConfig) -> None: + self.config = config + + def admit(self, group: list[Rollout]) -> bool: + advantages: list[float] = [] + for rollout in group: + if rollout.advantages is None: + continue + trainable = [value for sample in rollout.samples for value in sample.mask] + advantages.extend(advantage for advantage, keep in zip(rollout.advantages, trainable, strict=True) if keep) + if not advantages: + return True + return not all(self.config.reject_min <= advantage <= self.config.reject_max for advantage in advantages) diff --git a/src/prime_rl/orchestrator/curriculum/gates/base.py b/src/prime_rl/orchestrator/curriculum/gates/base.py new file mode 100644 index 0000000000..25d1c9037c --- /dev/null +++ b/src/prime_rl/orchestrator/curriculum/gates/base.py @@ -0,0 +1,27 @@ +"""Admission gate interface.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from prime_rl.orchestrator.types import Rollout + + +class AdmissionGate: + """Base class for user-authored training-sample admission policies.""" + + def admit(self, group: list[Rollout]) -> bool: + """Return whether a finalized group should enter the training batch.""" + return True + + def state_dict(self) -> dict[str, Any]: + """Return checkpoint state owned by this gate.""" + return {} + + def load_state_dict(self, state_dict: dict[str, Any]) -> None: + """Restore checkpoint state before results resume.""" + + def metrics(self) -> dict[str, float]: + """Return metrics relative to this gate's namespace.""" + return {} diff --git a/src/prime_rl/orchestrator/curriculum/samplers/__init__.py b/src/prime_rl/orchestrator/curriculum/samplers/__init__.py new file mode 100644 index 0000000000..7c1879671e --- /dev/null +++ b/src/prime_rl/orchestrator/curriculum/samplers/__init__.py @@ -0,0 +1,7 @@ +"""Task sampler interfaces and implementations.""" + +from prime_rl.orchestrator.curriculum.samplers.base import TaskSampler +from prime_rl.orchestrator.curriculum.samplers.pool import DifficultyPoolSampler +from prime_rl.orchestrator.curriculum.samplers.standard import StandardSampler + +__all__ = ["DifficultyPoolSampler", "StandardSampler", "TaskSampler"] diff --git a/src/prime_rl/orchestrator/curriculum/samplers/base.py b/src/prime_rl/orchestrator/curriculum/samplers/base.py new file mode 100644 index 0000000000..188271a300 --- /dev/null +++ b/src/prime_rl/orchestrator/curriculum/samplers/base.py @@ -0,0 +1,35 @@ +"""Task sampler interface.""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from collections.abc import Iterator +from typing import TYPE_CHECKING, Any + +import verifiers.v1 as vf + +if TYPE_CHECKING: + from prime_rl.orchestrator.types import Rollout + + +class TaskSampler(Iterator[vf.Task], ABC): + """Base class for user-authored task selection policies.""" + + @abstractmethod + def __next__(self) -> vf.Task: + """Choose the next task.""" + raise NotImplementedError + + def observe(self, group: list[Rollout]) -> None: + """Update sampling state from a finalized group.""" + + def state_dict(self) -> dict[str, Any]: + """Return checkpoint state owned by this sampler.""" + return {} + + def load_state_dict(self, state_dict: dict[str, Any]) -> None: + """Restore checkpoint state before sampling resumes.""" + + def metrics(self) -> dict[str, float]: + """Return metrics relative to this sampler's namespace.""" + return {} diff --git a/src/prime_rl/orchestrator/curriculum/samplers/pool.py b/src/prime_rl/orchestrator/curriculum/samplers/pool.py new file mode 100644 index 0000000000..5492e9002a --- /dev/null +++ b/src/prime_rl/orchestrator/curriculum/samplers/pool.py @@ -0,0 +1,95 @@ +"""Difficulty-pool task sampling.""" + +from __future__ import annotations + +import random +from collections import Counter +from collections.abc import Iterator, Sequence +from typing import TYPE_CHECKING, Any + +import verifiers.v1 as vf + +from prime_rl.orchestrator.curriculum.samplers.base import TaskSampler + +if TYPE_CHECKING: + from prime_rl.configs.orchestrator import DifficultyPoolSamplerConfig + from prime_rl.orchestrator.types import Rollout + + +class DifficultyPoolSampler(TaskSampler): + """Weight finite tasks by a pool derived from their latest group mean. + + Each pool's threshold is its inclusive maximum reward; the final pool is + the catch-all. Unseen tasks have neutral weight, so observations affect + sampling immediately without requiring a full taskset pass. + """ + + def __init__( + self, + config: DifficultyPoolSamplerConfig, + tasks: Sequence[vf.Task] | Iterator[vf.Task], + ) -> None: + if not isinstance(tasks, Sequence): + raise ValueError("DifficultyPoolSampler requires a finite taskset") + self.tasks = tuple(tasks) + if not self.tasks: + raise ValueError("DifficultyPoolSampler requires at least one task") + keys = [task.key for task in self.tasks] + duplicates = {key for key, count in Counter(keys).items() if count > 1} + if duplicates: + raise ValueError(f"Task keys must be unique within a taskset: {sorted(duplicates)}") + self.tasks_by_key = dict(zip(keys, self.tasks)) + self.rng = random.Random(config.seed) + self.pools = config.pools + ordered = sorted(self.pools.items(), key=lambda item: item[1].threshold) + self._ordered_pools = tuple(ordered) + self.task_rewards: dict[str, float] = {} + + def task_pool(self, task_key: str) -> str | None: + """Return the task's current pool, or ``None`` until it has a score.""" + score = self.task_rewards.get(task_key) + if score is None: + return None + for name, pool in self._ordered_pools: + if score <= pool.threshold: + return name + return self._ordered_pools[-1][0] + + def __next__(self) -> vf.Task: + weights = [] + for task in self.tasks: + pool = self.task_pool(task.key) + weights.append(1.0 if pool is None else self.pools[pool].weight) + if not any(weights): + raise RuntimeError("DifficultyPoolSampler has no tasks in a pool with positive weight") + return self.rng.choices(self.tasks, weights=weights, k=1)[0] + + def observe(self, group: list[Rollout]) -> None: + rewards = [rollout.reward for rollout in group if not rollout.has_error and rollout.agent.trainable] + if not rewards: + return + task_key = group[0].task.key + if task_key is None: + raise ValueError("A finalized group is missing Task.key") + self.task_rewards[task_key] = sum(rewards) / len(rewards) + + def state_dict(self) -> dict[str, Any]: + return { + "rng": self.rng.getstate(), + "task_rewards": dict(self.task_rewards), + } + + def load_state_dict(self, state_dict: dict[str, Any]) -> None: + self.rng.setstate(state_dict["rng"]) + self.task_rewards = dict(state_dict["task_rewards"]) + + def metrics(self) -> dict[str, float]: + occupancy = dict.fromkeys(self.pools, 0) + for task_key in self.task_rewards: + pool = self.task_pool(task_key) + if pool is not None: + occupancy[pool] += 1 + return { + "pool/unseen": float(len(self.tasks_by_key) - len(self.task_rewards)), + **{f"pool/{name}": float(count) for name, count in occupancy.items()}, + } diff --git a/src/prime_rl/orchestrator/curriculum/samplers/standard.py b/src/prime_rl/orchestrator/curriculum/samplers/standard.py new file mode 100644 index 0000000000..b5d0f9e33a --- /dev/null +++ b/src/prime_rl/orchestrator/curriculum/samplers/standard.py @@ -0,0 +1,49 @@ +"""Standard task iteration.""" + +from __future__ import annotations + +import itertools +from collections import Counter +from collections.abc import Iterator, Sequence +from typing import Any + +import verifiers.v1 as vf + +from prime_rl.orchestrator.curriculum.samplers.base import TaskSampler + + +class StandardSampler(TaskSampler): + """Advance the task iterator, cycling finite tasksets in source order.""" + + def __init__(self, tasks: Sequence[vf.Task] | Iterator[vf.Task]) -> None: + self.tasks = tuple(tasks) if isinstance(tasks, Sequence) else None + if self.tasks is not None: + if not self.tasks: + raise ValueError("A finite curriculum needs at least one task") + keys = [task.key for task in self.tasks] + duplicates = {key for key, count in Counter(keys).items() if count > 1} + if duplicates: + raise ValueError(f"Task keys must be unique within a taskset: {sorted(duplicates)}") + self.task_iterator = itertools.cycle(self.tasks) + else: + self.task_iterator = tasks + self.cursor = 0 + + def __next__(self) -> vf.Task: + task = next(self.task_iterator) + self.cursor += 1 + return task + + def state_dict(self) -> dict[str, Any]: + return {"cursor": self.cursor} + + def load_state_dict(self, state_dict: dict[str, Any]) -> None: + cursor = state_dict["cursor"] + if self.tasks is not None: + self.task_iterator = itertools.cycle(self.tasks) + offset = cursor % len(self.tasks) + else: + offset = cursor + for _ in range(offset): + next(self.task_iterator) + self.cursor = cursor diff --git a/src/prime_rl/orchestrator/dispatcher.py b/src/prime_rl/orchestrator/dispatcher.py index c6091ede59..76368bc5fb 100644 --- a/src/prime_rl/orchestrator/dispatcher.py +++ b/src/prime_rl/orchestrator/dispatcher.py @@ -113,6 +113,18 @@ def drain_keys(*, train_envs: set[str], eval_envs: set[str]) -> list[str]: return keys +def error_trace_task(group: GroupState | None) -> vf.TraceTask: + """Preserve task identity on a synthetic error trace.""" + if group is None: + return vf.TraceTask(type="Task", data=vf.TaskData(idx=-1, prompt=None)) + return vf.TraceTask( + type=type(group.task).__name__, + data=vf.TaskData(idx=group.task.data.idx, prompt=None), + key=group.task.key, + hash=group.task.hash, + ) + + class RolloutDispatcher: """``await dispatcher.start()`` runs the dispatch loop until ``stop()``. Pulls examples from ``TrainSource`` / ``EvalSource``, schedules @@ -476,10 +488,9 @@ async def handle_completed_rollout(self, task: asyncio.Task) -> None: return except Exception as exc: get_logger().warning(f"Rollout task failed in group {meta.group_id} ({meta.env_name}): {exc!r}") - task_idx = group.task.data.idx if group is not None else -1 rollouts = [ Rollout( - task=vf.TraceTask(type="Task", data=vf.TaskData(idx=task_idx, prompt=None)), + task=error_trace_task(group), agent=vf.AgentInfo(config=vf.AgentConfig()), ) ] @@ -532,8 +543,6 @@ async def drop_group(self, group_id: uuid.UUID) -> int: (both in-flight and not-yet-scheduled). Returns the count for off-policy metrics.""" group = self.groups.pop(group_id, None) - task_idx = group.task.data.idx if group is not None else -1 - # Sync claim phase: pop matching tasks from ``self.inflight`` and # release their permits in one non-yielding sweep. After this loop # the dropped tasks are no longer reachable from ``self.inflight``, @@ -552,7 +561,7 @@ async def drop_group(self, group_id: uuid.UUID) -> int: last_meta: InflightRollout | None = claimed[-1][1] if claimed else None for _, meta in claimed: trace = Rollout( - task=vf.TraceTask(type="Task", data=vf.TaskData(idx=task_idx, prompt=None)), + task=error_trace_task(group), agent=vf.AgentInfo(config=vf.AgentConfig()), ok=False, errors=[vf.Error(type="Cancelled", message="Off-policy cancel")], @@ -579,7 +588,7 @@ async def drop_group(self, group_id: uuid.UUID) -> int: unscheduled_cancelled = group.rollouts_to_schedule for _ in range(unscheduled_cancelled): trace = Rollout( - task=vf.TraceTask(type="Task", data=vf.TaskData(idx=task_idx, prompt=None)), + task=error_trace_task(group), agent=vf.AgentInfo(config=vf.AgentConfig()), ok=False, errors=[vf.Error(type="Cancelled", message="Off-policy cancel")], diff --git a/src/prime_rl/orchestrator/envs.py b/src/prime_rl/orchestrator/envs.py index 21c4445086..f6b392fe76 100644 --- a/src/prime_rl/orchestrator/envs.py +++ b/src/prime_rl/orchestrator/envs.py @@ -10,7 +10,7 @@ (``task_data``); the server pydantic-validates it into the taskset's declared ``TaskData`` type and runs it. That keeps the server (and every worker in its pool) stateless about data — no per-worker dataset loads, no idx-addressed task -cache — and gives the orchestrator real tasks to cycle, shuffle, and filter. +cache — and gives the orchestrator real tasks to sample. The server answers one ``Episode`` per run request, whose traces we validate into ``Trace[WireTaskData]`` — real ``vf.Trace``\\ s (never loose dicts) whose task diff --git a/src/prime_rl/orchestrator/eval_sink.py b/src/prime_rl/orchestrator/eval_sink.py index 51e7e0506b..7f733e11b5 100644 --- a/src/prime_rl/orchestrator/eval_sink.py +++ b/src/prime_rl/orchestrator/eval_sink.py @@ -1,6 +1,6 @@ """EvalSink: three-level rollout sink for eval epochs. -Same shape as ``TrainSink``, but no tokenization / advantages / filters: +Same shape as ``TrainSink``, but no tokenization, advantages, or admission: 1. ``process_rollout`` — no-op. 2. ``process_group`` — at ``group_size`` episodes, move the rollouts 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..c231911eae 100644 --- a/src/prime_rl/orchestrator/metrics.py +++ b/src/prime_rl/orchestrator/metrics.py @@ -333,7 +333,7 @@ def to_wandb(self, *, prefix: str, subset: Subset) -> dict[str, float]: class TrainMetrics(EpisodeMetrics): - """Common metrics plus the per-agent filter-pipeline rates. ``reward`` (flat over all traces) + """Common metrics plus per-agent pipeline rates. ``reward`` (flat over all traces) serves the console log lines; the wandb reward stats are per-agent.""" @property @@ -342,18 +342,12 @@ 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. + # Pipeline verdicts are per-trace, so they read per agent like the rest. for agent, traces in self.by_agent().items(): p = f"{prefix}/{subset}/{agent}" rollouts = traces.rollouts out[f"{p}/is_trainable/mean"] = sum(float(r.is_trainable) for r in rollouts) / len(rollouts) - out[f"{p}/is_filtered/mean"] = sum(float(r.is_filtered) for r in rollouts) / len(rollouts) - names = sorted({name for r in rollouts for name in r.filter_results}) - out |= { - f"{p}/filters/{name}/mean": sum(1 for r in rollouts if r.filter_results.get(name)) / len(rollouts) - for name in names - } + out[f"{p}/is_admitted/mean"] = sum(float(r.is_admitted) for r in rollouts) / len(rollouts) return out @@ -396,8 +390,8 @@ def to_wandb(self, *, prefix: str, subset: Subset) -> dict[str, float]: class TrainRollouts: - """A list of train rollouts (everything that came back, errored + filtered + untrainable - included). ``effective`` is the clean trainable subset (a view of the same traces); + """A list of train rollouts (everything that came back, including rejected, + errored and untrainable traces). ``effective`` is the clean trainable subset; ``metrics`` builds ``TrainMetrics`` over them.""" def __init__(self, rollouts: list[Rollout] | None = None) -> None: @@ -414,7 +408,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 r.is_admitted and not r.has_error 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 8f9c641ef8..08fb1f0189 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 → advantages → admission) and returns a ``TrainBatch`` when the threshold is met. - ``EvalSink`` ingests eval rollouts and returns an ``EvalBatch`` (the full returned cohort) on epoch completion. @@ -47,7 +47,6 @@ 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 ( @@ -93,9 +92,7 @@ # 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 contain no samples. MAX_CONSECUTIVE_EMPTY_BATCHES = 10 # Maximum batches the orchestrator may run ahead of the trainer. The @@ -236,10 +233,6 @@ async def setup(self) -> None: if config.heartbeat is not None: self.heart = Heartbeat(config.heartbeat.url) - # Filters apply to train rollouts only - pre_filters = setup_filters(config.pre_batch_filters, vocab_size=self.tokenizer.vocab_size, kind="pre-batch") - post_filters = setup_filters(config.post_batch_filters, vocab_size=self.tokenizer.vocab_size, kind="post-batch") - get_logger().info("Loading training environments") self.train_envs = TrainEnvs( config.train.source, @@ -407,8 +400,7 @@ 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, + on_result=self.train_source.on_result, ) self.eval_sink = EvalSink(eval_envs=self.eval_envs) if self.eval_envs is not None else None self.watcher = WeightWatcher( @@ -520,7 +512,7 @@ async def main_loop(self) -> None: except asyncio.TimeoutError: continue - # Every completed rollout — errored, filtered, or never batched — lands in the + # Every completed rollout — errored, rejected, or never batched — lands in the # ``all`` trace file the moment it arrives, so it survives crashes and drains. # Train rollouts belong to the batch window currently collecting (``progress.step``), # eval rollouts to the step whose eval triggered them. @@ -583,14 +575,14 @@ async def finalize_train_batch(self, batch: TrainBatch) -> None: if not batch.samples: self.consecutive_empty_batches += 1 get_logger().warning( - f"Step {step}: empty train batch (0 of {len(batch.rollouts)} generated rollouts shipped — " - f"all errored or filtered out) " - f"(consecutive empty batches: {self.consecutive_empty_batches}/{MAX_CONSECUTIVE_EMPTY_BATCHES})" + f"Step {step}: empty train batch after {len(batch.rollouts)} finalized rollouts " + f"(consecutive empty batches: " + f"{self.consecutive_empty_batches}/{MAX_CONSECUTIVE_EMPTY_BATCHES})" ) 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 algorithm credit and task difficulty." ) return self.consecutive_empty_batches = 0 @@ -599,7 +591,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" ) # Ship batch ``step`` only once the trainer has published v{step-1-TARGET_LAG}. @@ -646,14 +638,14 @@ async def finalize_train_batch(self, batch: TrainBatch) -> None: trim_process_memory() # Rollout metrics over the {agg,} × {all,effective} matrix. ``batch.rollouts`` is the - # full arrival window (errored + filtered included); ``.effective`` is the clean subset. + # full arrival window (errored + rejected included); ``.effective`` is the clean subset. metrics: dict[str, float] = {} for subset, pool in (("all", batch.rollouts), ("effective", effective)): metrics |= pool.metrics.to_wandb(prefix="train/agg", subset=subset) for env_name, env_pool in pool.by_env().items(): metrics |= env_pool.metrics.to_wandb(prefix=f"train/{env_name}", subset=subset) - # Progress / timing / env-share / pre-filter accounting (assembled here, not in the metrics + # Progress / timing / env-share accounting (assembled here, not in the metrics # objects). ``num_tokens`` is over the full arrival window; the input/output breakdown is over # the effective (shipped) subset, summing the same ``vf.Trace`` token properties the metric # matrix reports. @@ -679,12 +671,7 @@ 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 - ) - for name, count in self.train_sink.pre_filter_dropped_by_name.items(): - metrics[f"pre_filters/all/{name}/rate"] = count / self.train_sink.pre_filter_seen + metrics |= self.train_source.metrics() await monitors.log(metrics, step=step) self.wait_for_policy_time = 0.0 @@ -697,7 +684,6 @@ 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.maybe_trigger_eval(self.progress.step) trim_process_memory() diff --git a/src/prime_rl/orchestrator/train_sink.py b/src/prime_rl/orchestrator/train_sink.py index f6cad74ee8..a235c40a56 100644 --- a/src/prime_rl/orchestrator/train_sink.py +++ b/src/prime_rl/orchestrator/train_sink.py @@ -4,11 +4,12 @@ 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 +2. ``process_group`` — removes 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``. + per-sample wire stamping), then asks the source curriculum whether the + result should train. +3. ``process_batch`` — assembles the trainer-bound ``TrainingSample`` list + and optionally removes zero-advantage RL payload. ``add()`` takes one episode (``list[Rollout]``) and returns ``TrainBatch | None``; group accounting counts episodes, never loose traces. @@ -21,16 +22,18 @@ import asyncio import uuid from collections import defaultdict +from collections.abc import Callable from prime_rl.configs.orchestrator import OrchestratorConfig from prime_rl.orchestrator.envs import TrainEnvs -from prime_rl.orchestrator.filters import RolloutFilter, apply_filters from prime_rl.orchestrator.metrics import TrainRollouts from prime_rl.orchestrator.trajectories import trace_to_samples from prime_rl.orchestrator.types import Rollout, TrainBatch from prime_rl.transport import TrainingSample from prime_rl.utils.logger import get_logger +MAX_CONSECUTIVE_ZERO_OUTPUT_BATCH_EQUIVALENTS = 10 + def payload_tokens(rollout: Rollout) -> int: """Token cost of the rollout's trainer-bound payload — the samples built by @@ -45,6 +48,38 @@ def payload_tokens(rollout: Rollout) -> int: return sum(len(sample.token_ids) for sample in rollout.samples) or rollout.num_total_tokens +def _prune_zero_advantages(sample: TrainingSample) -> bool: + """Remove zero-advantage tokens from the RL component. + + Return whether the sample still carries any RL, CE, or reference-KL + component and therefore needs to be shipped. + """ + if sample.advantages is None: + return True + + if sample.rl_weights is None: + rl_weights = [1.0 if trainable else 0.0 for trainable in sample.mask] + else: + rl_weights = list(sample.rl_weights) + + changed = False + for index, (trainable, advantage, weight) in enumerate( + zip(sample.mask, sample.advantages, rl_weights, strict=True) + ): + if trainable and advantage == 0.0 and weight != 0.0: + rl_weights[index] = 0.0 + changed = True + + if not changed: + return True + + sample.rl_weights = rl_weights + has_rl = any(trainable and weight != 0.0 for trainable, weight in zip(sample.mask, rl_weights, strict=True)) + has_ce = sample.ce_weights is not None and any(weight != 0.0 for weight in sample.ce_weights) + has_ref_kl = sample.ref_kl_weights is not None and any(weight != 0.0 for weight in sample.ref_kl_weights) + return has_rl or has_ce or has_ref_kl + + class TrainSink: """Three-level train sink. Constructed once, fed via ``add(rollout)``.""" @@ -57,8 +92,7 @@ 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], + on_result: 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 +103,10 @@ 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.on_result = on_result # Observation window for the next shipped batch: rollouts of groups - # finalized since the last ship (errored + filtered + survivors). + # finalized since the last ship (errored + rejected + admitted). # In-progress groups stay out until they finalize. self.pending_rollouts: TrainRollouts = TrainRollouts() # Keyed by the dispatcher's group UUID. ``(env_name, task_idx)`` @@ -88,11 +121,10 @@ def __init__( # runs), kept in sync on append/pop so the readiness check never # 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] = {} + # Finalized work since the most recent positive contribution, measured + # in the active batch unit. + self.zero_output_units: int = 0 + self.reported_zero_output_windows: int = 0 def group_size_for(self, env_name: str) -> int: return self.train_envs.get(env_name).config.group_size @@ -166,8 +198,7 @@ async def process_rollout(self, rollout: Rollout) -> None: await self.train_envs.get(rollout.env_name).algorithm.finalize_rollout(rollout) 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``.""" + """Finalize one group, ask its curriculum for admission, and queue it.""" group = self.pending_groups.pop(group_id, []) self.pending_group_episodes.pop(group_id, None) if not group: @@ -187,6 +218,8 @@ async def process_group(self, group_id: uuid.UUID) -> None: # Untrainable traces carry no samples and must not skew the group baseline. survivors = [r for r in survivors if r.agent.trainable] if not survivors: + self._admit(group) + self._record_zero_output(group) get_logger().debug( f"Finished group | env={env_name} task_idx={task_idx} | " f"rollouts={len(group)} (errored={num_errored}) | dropped: no trainable survivors" @@ -205,44 +238,65 @@ 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 + if not self._admit(group): + self._record_zero_output(group) + get_logger().debug( + f"Finished group | env={env_name} task_idx={task_idx} | " + f"rollouts={len(group)} (errored={num_errored}) | rejected by curriculum" + ) + return + 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(): - if hit: - self.pre_filter_dropped_by_name[name] = self.pre_filter_dropped_by_name.get(name, 0) + 1 - filtered_by_name[name] = filtered_by_name.get(name, 0) + 1 - continue - # Reset annotations so the post-batch filter pass starts clean - r.filter_results = {} - r.is_filtered = False self.pending_batch.append(r) if self.token_batch_size is not None: self.pending_tokens += payload_tokens(r) + self.zero_output_units = 0 + self.reported_zero_output_windows = 0 - # 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 "—" 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}) | reward={avg_reward:.4f}" + ) + + def _admit(self, group: list[Rollout]) -> bool: + admitted = self.on_result(group) if self.on_result is not None else True + for rollout in group: + rollout.is_admitted = admitted + return admitted + + def _record_zero_output(self, group: list[Rollout]) -> None: + if self.batch_size is not None: + self.zero_output_units += len(group) + else: + payload = sum(payload_tokens(rollout) for rollout in group) + self.zero_output_units += payload or self.config.seq_len * len(group) + self._check_zero_output_budget() + + def _check_zero_output_budget(self) -> None: + target = self.batch_size if self.batch_size is not None else self.token_batch_size + assert target is not None + windows = self.zero_output_units // target + if windows <= self.reported_zero_output_windows: + return + self.reported_zero_output_windows = windows + get_logger().warning( + f"No admitted train payload after {self.zero_output_units} finalized units " + f"(consecutive zero-output batch equivalents: " + f"{windows}/{MAX_CONSECUTIVE_ZERO_OUTPUT_BATCH_EQUIVALENTS})" ) + if windows >= MAX_CONSECUTIVE_ZERO_OUTPUT_BATCH_EQUIVALENTS: + raise RuntimeError( + f"{windows} consecutive zero-output batch equivalents — " + "check the curriculum admission policy and task difficulty." + ) 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 +313,13 @@ 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] + if self.config.train.filter_zero_advantages: + for rollout in cohort: + rollout.samples = [sample for sample in rollout.samples if _prune_zero_advantages(sample)] + samples: list[TrainingSample] = [sample for rollout in cohort for sample in rollout.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 + rejected + admitted) — 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 @@ -276,8 +328,3 @@ def process_batch(self) -> TrainBatch: if samples: 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() diff --git a/src/prime_rl/orchestrator/train_source.py b/src/prime_rl/orchestrator/train_source.py index 9c401821cf..36d6b1f970 100644 --- a/src/prime_rl/orchestrator/train_source.py +++ b/src/prime_rl/orchestrator/train_source.py @@ -1,35 +1,20 @@ -"""TrainSource: weighted round-robin across train envs, infinite pull. - -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.""" +"""Training source selection and curriculum lifecycle.""" from __future__ import annotations import random -from collections.abc import Iterator - -import verifiers.v1 as vf +from collections import defaultdict +from typing import TYPE_CHECKING, Any +from prime_rl.orchestrator.curriculum import Curriculum from prime_rl.orchestrator.envs import TrainEnvs +if TYPE_CHECKING: + from prime_rl.orchestrator.types import Rollout -class TrainSource: - """``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 - 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 - deterministic). Cursors advance at dispatch time (ahead of shipped - batches), so a resume skips the tasks that were in flight at checkpoint - time.""" +class TrainSource: + """Mix train envs and host one user-authored curriculum per env.""" def __init__(self, train_envs: TrainEnvs) -> None: self.rng = random.Random(42) @@ -37,73 +22,59 @@ def __init__(self, train_envs: TrainEnvs) -> None: if not self.envs: raise ValueError("TrainSource needs at least one train env") - # A finite env's example table in canonical order (each epoch's shuffle - # starts from this); ``None`` for an infinite env, whose generator - # (``self.iters``) is pulled per example. - self.base_rows: dict[str, list[dict] | None] = {} - self.examples: dict[str, list[dict] | None] = {} - self.iters: dict[str, Iterator[vf.Task]] = {} - self.epochs: dict[str, int] = {} - self.cursors: dict[str, int] = {} + self.curricula: dict[str, Curriculum] = {} for env in self.envs: - assert env.tasks is not None, f"env {env.name} not started" - if env.num_tasks is None: # infinite: pull the generator per example - rows: list[dict] | None = None - self.iters[env.name] = env.tasks - else: - rows = [{"task": task, "env_name": env.name} for task in env.tasks] - self.base_rows[env.name] = rows - self.epochs[env.name] = 1 - self.cursors[env.name] = 0 - self.examples[env.name] = self._shuffle(env.name) + if env.tasks is None: + raise RuntimeError(f"env {env.name} not started") + tasks = env.tasks if env.num_tasks is None else list(env.tasks) + self.curricula[env.name] = Curriculum(env.config.curriculum, tasks) - self.env_names = [e.name for e in self.envs] - self.weights: list[float] = [float(e.config.ratio) for e in self.envs] + self.env_names = [env.name for env in self.envs] + self.weights = [float(env.config.ratio) for env in self.envs] + self._admitted: dict[str, int] = defaultdict(int) + self._rejected: dict[str, int] = defaultdict(int) - def _shuffle(self, env_name: str) -> list[dict] | None: - """The env's example table shuffled for its current epoch — a pure - function of (canonical order, epoch), so a restored position replays - the exact epoch permutation.""" - rows = self.base_rows[env_name] - if rows is None: - return None - rows = rows.copy() - random.Random(self.epochs[env_name]).shuffle(rows) - return rows + def next_example(self) -> dict[str, Any]: + env_name = self.rng.choices(self.env_names, weights=self.weights, k=1)[0] + return { + "env_name": env_name, + "task": next(self.curricula[env_name].sampler), + } + + def on_result(self, group: list[Rollout]) -> bool: + """Report a finalized group and return whether it should train.""" + if not group: + raise ValueError("Cannot report an empty rollout group") + env_name = group[0].env_name + admitted = self.curricula[env_name].on_result(group) + if not isinstance(admitted, bool): + raise TypeError(f"Curriculum.on_result() must return bool, got {type(admitted).__name__}") + if admitted: + self._admitted[env_name] += 1 + else: + self._rejected[env_name] += 1 + return admitted + + def metrics(self) -> dict[str, float]: + metrics: dict[str, float] = {} + for env_name, curriculum in self.curricula.items(): + admitted = self._admitted.pop(env_name, 0) + rejected = self._rejected.pop(env_name, 0) + total = admitted + rejected + if total: + metrics[f"curriculum/{env_name}/admission_rate"] = admitted / total + metrics |= {f"curriculum/{env_name}/{name}": float(value) for name, value in curriculum.metrics().items()} + return metrics - def state_dict(self) -> dict: - """Env-choice RNG state + per-env ``{epoch, cursor}``.""" + def state_dict(self) -> dict[str, Any]: return { "rng": self.rng.getstate(), - "envs": {name: {"epoch": self.epochs[name], "cursor": self.cursors[name]} for name in self.epochs}, + "envs": {name: curriculum.state_dict() for name, curriculum in self.curricula.items()}, } - def load_state_dict(self, state_dict: dict) -> None: + def load_state_dict(self, state_dict: dict[str, Any]) -> None: self.rng.setstate(state_dict["rng"]) - for name, position in state_dict["envs"].items(): - if name not in self.base_rows: - continue - self.epochs[name] = position["epoch"] - self.cursors[name] = position["cursor"] - if self.base_rows[name] is None: - for _ in range(position["cursor"]): - next(self.iters[name]) - else: - self.examples[name] = self._shuffle(name) - - def next_example(self) -> dict | None: - env_name = self.rng.choices(self.env_names, weights=self.weights, k=1)[0] - rows = self.examples[env_name] - cursor = self.cursors[env_name] - if rows is None: # infinite env: pull the next generated task - task = next(self.iters[env_name]) - self.cursors[env_name] = cursor + 1 - return {"task": task, "env_name": env_name} - if cursor >= len(rows): - self.epochs[env_name] += 1 - rows = self._shuffle(env_name) - self.examples[env_name] = rows - cursor = 0 - example = rows[cursor] - self.cursors[env_name] = cursor + 1 - return example + for name, curriculum_state in state_dict["envs"].items(): + curriculum = self.curricula.get(name) + if curriculum is not None: + curriculum.load_state_dict(curriculum_state) diff --git a/src/prime_rl/orchestrator/types.py b/src/prime_rl/orchestrator/types.py index b77de8ba31..c597cc110a 100644 --- a/src/prime_rl/orchestrator/types.py +++ b/src/prime_rl/orchestrator/types.py @@ -96,11 +96,9 @@ class Rollout(vf.Trace[DataT], Generic[DataT]): samples: list[TrainingSample] = Field(default_factory=list, exclude=True) # Per-token rl advantage stream, full-length-N (= len(token_ids)) per # sample, concatenated across the rollout's samples in order; 0.0 on - # non-trainable positions. None = no credit assigned (advantage-based - # filters skip it; the wire ships no advantage stream). + # non-trainable positions. None means no credit was assigned. 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) + is_admitted: bool = Field(default=True, exclude=True) eval_step: int | None = Field(default=None, exclude=True) def assign_advantages(self, values: float | list[float]) -> None: @@ -140,9 +138,9 @@ def is_trainable(self) -> bool: @dataclass class TrainBatch: """``rollouts`` is the observation window since the last ship — every rollout of every group - finalized in that span (errored + filtered included; rollouts of still-incomplete groups wait + finalized in that span (errored + rejected included; rollouts of still-incomplete groups wait for a later window). Its ``.effective`` / ``.metrics`` views drive logging. ``samples`` is the - trainer-bound payload (the shipped cohort's post-filter survivors) — an empty list means nothing + trainer-bound payload from the admitted cohort — an empty list means nothing ships, which would stall the trainer. Trainable counts derive from ``rollouts.effective`` (``r.is_trainable``) and token totals from ``samples``, so neither is carried as a field.""" diff --git a/tests/unit/orchestrator/test_curriculum.py b/tests/unit/orchestrator/test_curriculum.py new file mode 100644 index 0000000000..5856d69d45 --- /dev/null +++ b/tests/unit/orchestrator/test_curriculum.py @@ -0,0 +1,173 @@ +from collections.abc import Iterator +from types import SimpleNamespace + +import pytest +import verifiers.v1 as vf + +from prime_rl.configs.orchestrator import ( + AdvRangeGateConfig, + CurriculumConfig, + DifficultyPoolConfig, + DifficultyPoolSamplerConfig, +) +from prime_rl.orchestrator.curriculum import ( + AdvRangeGate, + Curriculum, + StandardSampler, +) +from prime_rl.orchestrator.train_source import TrainSource +from prime_rl.orchestrator.types import Rollout +from prime_rl.transport import TrainingSample + + +def make_task(idx: int) -> vf.Task: + return vf.Task(vf.TaskData(idx=idx, prompt=f"task {idx}")) + + +def make_rollout( + task: vf.Task, + *, + env_name: str = "test", + reward: float = 0.0, + advantages: list[float] | None = None, +) -> Rollout: + samples = [] + if advantages is not None: + samples = [ + TrainingSample( + token_ids=list(range(len(advantages))), + mask=[True] * len(advantages), + logprobs=[0.0] * len(advantages), + temperatures=[1.0] * len(advantages), + advantages=advantages, + env_name=env_name, + ) + ] + return Rollout( + task=vf.TraceTask( + type=type(task).__name__, + data=task.data, + key=task.key, + hash=task.hash, + ), + agent=vf.AgentInfo(config=vf.AgentConfig()), + env_name=env_name, + rewards={"reward": vf.Reward(score=reward)}, + advantages=advantages, + samples=samples, + ok=True, + ) + + +def test_default_curriculum_resumes_finite_and_infinite_tasksets() -> None: + tasks = [make_task(i) for i in range(5)] + finite = StandardSampler(tasks) + for _ in range(3): + next(finite) + state = finite.state_dict() + expected = [next(finite).key for _ in range(4)] + + restored = StandardSampler(tasks) + restored.load_state_dict(state) + assert [next(restored).key for _ in range(4)] == expected + + def task_stream() -> Iterator[vf.Task]: + yield from (make_task(i) for i in range(10)) + + infinite = StandardSampler(task_stream()) + next(infinite) + next(infinite) + restored_infinite = StandardSampler(task_stream()) + restored_infinite.load_state_dict(infinite.state_dict()) + assert next(restored_infinite).key == make_task(2).key + + +def test_curriculum_requires_unique_finite_task_keys() -> None: + task = make_task(0) + with pytest.raises(ValueError, match="Task keys must be unique"): + StandardSampler([task, task]) + + +def test_train_source_composes_sampler_and_all_gates_with_state_and_metrics() -> None: + tasks = [make_task(i) for i in range(3)] + pools = {"all": DifficultyPoolConfig(threshold=1.0, weight=1.0)} + config = SimpleNamespace( + ratio=1.0, + curriculum=CurriculumConfig( + sampler=DifficultyPoolSamplerConfig(pools=pools), + gates={ + "reject": AdvRangeGateConfig(), + "admit": AdvRangeGateConfig(reject_min=0.5, reject_max=0.5), + }, + ), + ) + env = SimpleNamespace(name="test", tasks=iter(tasks), num_tasks=len(tasks), config=config) + source = TrainSource([env]) + + sampled = source.next_example()["task"] + assert source.on_result([make_rollout(sampled, reward=0.25, advantages=[0.0])]) is False + assert source.metrics() == { + "curriculum/test/admission_rate": 0.0, + "curriculum/test/sampler/pool/unseen": 2.0, + "curriculum/test/sampler/pool/all": 1.0, + } + + state = source.state_dict() + restored = TrainSource([SimpleNamespace(name="test", tasks=iter(tasks), num_tasks=len(tasks), config=config)]) + restored.load_state_dict(state) + assert restored.curricula["test"].state_dict()["sampler"]["task_rewards"] == {sampled.key: 0.25} + assert restored.curricula["test"].state_dict()["gates"] == { + "reject": {}, + "admit": {}, + } + + +def test_difficulty_pools_stack_with_advantage_gate_and_resume_sampling() -> None: + tasks = [make_task(i) for i in range(3)] + pools = { + "hard": DifficultyPoolConfig(threshold=0.25, weight=0.0), + "normal": DifficultyPoolConfig(threshold=0.75, weight=1.0), + "easy": DifficultyPoolConfig(threshold=1.0, weight=0.0), + } + sampler_config = DifficultyPoolSamplerConfig(pools=pools, seed=7) + gate_config = AdvRangeGateConfig() + config = CurriculumConfig( + sampler=sampler_config, + gates={"zero_advantage": gate_config}, + ) + curriculum = Curriculum(config, tasks) + rewards = {0: 0.1, 1: 0.5, 2: 0.9} + decisions = [] + for index, task in enumerate(tasks): + rollout = make_rollout(task, reward=rewards[task.data.idx], advantages=[float(index > 0)]) + decisions.append(curriculum.on_result([rollout])) + + assert decisions == [False, True, True] + assert curriculum.metrics() == { + "sampler/pool/unseen": 0.0, + "sampler/pool/hard": 1.0, + "sampler/pool/normal": 1.0, + "sampler/pool/easy": 1.0, + } + state = curriculum.state_dict() + expected = [next(curriculum.sampler).key for _ in range(10)] + assert set(expected) == {tasks[1].key} + restored = Curriculum(config, tasks) + restored.load_state_dict(state) + assert [next(restored.sampler).key for _ in range(10)] == expected + + +def test_advantage_range_gate_generalizes_zero_advantage_rejection() -> None: + task = make_task(0) + zero_gate = AdvRangeGate(AdvRangeGateConfig()) + assert zero_gate.admit([make_rollout(task, advantages=[0.0, 0.0])]) is False + assert zero_gate.admit([make_rollout(task, advantages=[0.0, 0.2])]) is True + assert zero_gate.admit([make_rollout(task)]) is True + + tolerance_gate = AdvRangeGate(AdvRangeGateConfig(reject_min=-0.1, reject_max=0.1)) + assert tolerance_gate.admit([make_rollout(task, advantages=[-0.05, 0.0, 0.05])]) is False + + masked = make_rollout(task, advantages=[0.0, 0.5]) + masked.samples[0].mask = [False, True] + positive_gate = AdvRangeGate(AdvRangeGateConfig(reject_min=0.5, reject_max=0.5)) + assert positive_gate.admit([masked]) is False diff --git a/tests/unit/orchestrator/test_filters.py b/tests/unit/orchestrator/test_filters.py deleted file mode 100644 index 69ce76d029..0000000000 --- a/tests/unit/orchestrator/test_filters.py +++ /dev/null @@ -1,408 +0,0 @@ -import math -import uuid - -import verifiers.v1 as vf - -from prime_rl.configs.orchestrator import GibberishFilterConfig, RepetitionFilterConfig -from prime_rl.orchestrator.filters import ( - GibberishFilter, - RepetitionFilter, - apply_filters, - setup_filter, - setup_filters, -) -from prime_rl.orchestrator.types import 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 - masked-True tokens + logprobs).""" - return vf.MessageNode( - message=vf.AssistantMessage(content="x"), - token_ids=token_ids, - mask=[True] * len(token_ids), - logprobs=logprobs, - ) - - -def _scaffold_assistant_node( - completion_ids: list[int], completion_logprobs: list[float], *, scaffold: int = 2 -) -> vf.MessageNode: - """A realistic v1 assistant node: a leading generation-prompt scaffold (mask=False, not - model-sampled) then the sampled completion. ``logprobs`` cover only the completion suffix - (vLLM returns logprobs for generated tokens only) — the exact layout where per-node - ``zip(token_ids, logprobs, mask)`` mispairs and the branch streams normalize.""" - return vf.MessageNode( - message=vf.AssistantMessage(content="x"), - token_ids=[1] * scaffold + completion_ids, - mask=[False] * scaffold + [True] * len(completion_ids), - logprobs=completion_logprobs, - ) - - -def _make_rollout( - completion_ids: list[int], - completion_logprobs: list[float], - *, - reward: float = 1.0, - multi_step: bool = False, -) -> Rollout: - """Build a ``Rollout`` (a message-graph trace) carrying the completion tokens — enough for - the filters to inspect each node's sampled tokens / logprobs.""" - if multi_step: - mid = len(completion_ids) // 2 - nodes = [ - _assistant_node(completion_ids[:mid], completion_logprobs[:mid]), - _assistant_node(completion_ids[mid:], completion_logprobs[mid:]), - ] - else: - nodes = [_assistant_node(completion_ids, completion_logprobs)] - rollout = Rollout[vf.TaskData]( - task=vf.TraceTask(type="Task", data=vf.TaskData(idx=0, prompt="")), - agent=vf.AgentInfo(config=vf.AgentConfig()), - nodes=nodes, - rewards={"reward": vf.Reward(score=reward)}, - ) - rollout.env_name = "test" - rollout.group_id = uuid.uuid4() - return rollout - - -def _make_gibberish_filter(vocab_size=128_000, token_id_threshold=100_000, logprob_offset=2.0, enforce=False): - logprob_threshold = -math.log(vocab_size) - logprob_offset - return GibberishFilter( - name="gibberish", token_id_threshold=token_id_threshold, logprob_threshold=logprob_threshold, enforce=enforce - ) - - -def _make_repetition_filter(window=5, prob_threshold=0.99, enforce=False): - return RepetitionFilter( - name="repetition", window=window, logprob_threshold=math.log(prob_threshold), enforce=enforce - ) - - -# --- GibberishFilter tests --- - - -def test_gibberish_detects_rare_low_prob_token(): - gibberish_filter = _make_gibberish_filter() - - result = gibberish_filter.check( - _make_rollout( - completion_ids=[50, 120_000, 80], - completion_logprobs=[-1.0, gibberish_filter.logprob_threshold - 1.0, -0.5], - ) - ) - assert result.detected is True - - -def test_gibberish_ignores_normal_tokens(): - gibberish_filter = _make_gibberish_filter() - - result = gibberish_filter.check( - _make_rollout( - completion_ids=[10, 200, 5000], - completion_logprobs=[-1.0, -2.0, -3.0], - ) - ) - assert result.detected is False - - -def test_gibberish_ignores_high_prob_rare_token(): - gibberish_filter = _make_gibberish_filter() - - result = gibberish_filter.check( - _make_rollout( - completion_ids=[120_000], - completion_logprobs=[-0.5], - ) - ) - assert result.detected is False - - -def test_gibberish_works_across_trajectory_steps(): - gibberish_filter = _make_gibberish_filter() - - result = gibberish_filter.check( - _make_rollout( - completion_ids=[50, 60, 120_000, 80], - completion_logprobs=[-1.0, -0.5, gibberish_filter.logprob_threshold - 1.0, -0.5], - multi_step=True, - ) - ) - assert result.detected is True - - -def test_gibberish_aligns_logprobs_under_generation_prompt_scaffold(): - """Regression: the assistant node carries a generation-prompt scaffold (mask=False) and - suffix-only logprobs, and the gibberish token is the LAST completion token. The old - per-node ``zip(token_ids, logprobs, mask)`` truncated at len(logprobs) and never examined - it; reading the aligned branch streams detects it.""" - gibberish_filter = _make_gibberish_filter() - - rollout = 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])], - rewards={"reward": vf.Reward(score=1.0)}, - ) - - result = gibberish_filter.check(rollout) - assert result.detected is True - - -# --- RepetitionFilter tests --- - - -def test_repetition_triggers_after_window(): - repetition_filter = _make_repetition_filter(window=5) - - result = repetition_filter.check( - _make_rollout( - completion_ids=list(range(5)), - completion_logprobs=[-0.001] * 5, - ) - ) - assert result.detected is True - - -def test_repetition_no_trigger_below_window(): - repetition_filter = _make_repetition_filter(window=5) - - result = repetition_filter.check( - _make_rollout( - completion_ids=list(range(4)), - completion_logprobs=[-0.001] * 4, - ) - ) - assert result.detected is False - - -def test_repetition_resets_on_low_prob(): - repetition_filter = _make_repetition_filter(window=5) - - logprobs = [-0.001] * 3 + [-2.0] + [-0.001] * 3 - result = repetition_filter.check( - _make_rollout( - completion_ids=list(range(7)), - completion_logprobs=logprobs, - ) - ) - assert result.detected is False - - -def test_repetition_varied_probs_no_trigger(): - repetition_filter = _make_repetition_filter(window=3) - - result = repetition_filter.check( - _make_rollout( - completion_ids=list(range(6)), - completion_logprobs=[-0.001, -3.0, -0.001, -3.0, -0.001, -3.0], - ) - ) - assert result.detected is False - - -# --- setup_filter / setup_filters tests --- - - -def test_setup_filter_gibberish(): - config = GibberishFilterConfig(token_id_threshold=100_000, logprob_offset=2.0) - gibberish_filter = setup_filter(config, vocab_size=128_000) - assert isinstance(gibberish_filter, GibberishFilter) - assert gibberish_filter.name == "gibberish" - assert gibberish_filter.token_id_threshold == 100_000 - assert abs(gibberish_filter.logprob_threshold - (-math.log(128_000) - 2.0)) < 1e-10 - assert gibberish_filter.enforce is False - - -def test_setup_filter_gibberish_enforce(): - config = GibberishFilterConfig(enforce=True) - gibberish_filter = setup_filter(config, vocab_size=128_000) - assert gibberish_filter.enforce is True - - -def test_setup_filter_repetition(): - config = RepetitionFilterConfig(window=3_000, prob_threshold=0.99) - repetition_filter = setup_filter(config, vocab_size=128_000) - assert isinstance(repetition_filter, RepetitionFilter) - assert repetition_filter.name == "repetition" - assert repetition_filter.window == 3_000 - assert abs(repetition_filter.logprob_threshold - math.log(0.99)) < 1e-10 - assert repetition_filter.enforce is False - - -def test_setup_filter_repetition_enforce(): - config = RepetitionFilterConfig(enforce=True) - repetition_filter = setup_filter(config, vocab_size=128_000) - assert repetition_filter.enforce is True - - -def test_setup_filters_multiple(): - configs = [ - GibberishFilterConfig(), - RepetitionFilterConfig(), - ] - filters = setup_filters(configs, vocab_size=128_000, kind="post-batch") - assert len(filters) == 2 - assert filters[0].name == "gibberish" - assert filters[1].name == "repetition" - - -# --- apply_filters tests (enforce=True) --- - - -def test_apply_filters_enforced_flags_rollout(): - gibberish_filter = _make_gibberish_filter(enforce=True) - - rollout = _make_rollout( - completion_ids=[120_000], - completion_logprobs=[gibberish_filter.logprob_threshold - 1.0], - reward=1.0, - ) - - apply_filters([gibberish_filter], [rollout]) - - assert rollout.reward == 1.0 - assert rollout.nodes[0].token_ids == [120_000] - assert rollout.nodes[0].mask == [True] - assert rollout.stop_condition is None - assert rollout.filter_results == {"gibberish": True} - assert rollout.is_filtered is True - - -def test_apply_filters_preserves_clean_rollouts(): - gibberish_filter = _make_gibberish_filter(enforce=True) - - rollout = _make_rollout( - completion_ids=[50, 60, 70], - completion_logprobs=[-1.0, -2.0, -1.5], - reward=1.0, - ) - - apply_filters([gibberish_filter], [rollout]) - - assert rollout.reward == 1.0 - assert rollout.nodes[0].token_ids == [50, 60, 70] - assert all(rollout.nodes[0].mask) - assert rollout.stop_condition is None - assert rollout.filter_results == {"gibberish": False} - assert rollout.is_filtered is False - - -def test_apply_filters_first_filter_wins(): - gibberish_filter = _make_gibberish_filter(enforce=True) - repetition_filter = _make_repetition_filter(window=2, enforce=True) - - rollout = _make_rollout( - completion_ids=[120_000, 1, 2], - completion_logprobs=[gibberish_filter.logprob_threshold - 1.0, -0.001, -0.001], - reward=1.0, - ) - - apply_filters([gibberish_filter, repetition_filter], [rollout]) - - assert rollout.stop_condition is None - assert rollout.filter_results == {"gibberish": True, "repetition": False} - assert rollout.is_filtered is True - - -def test_apply_filters_empty_list(): - rollout = _make_rollout( - completion_ids=[1, 2, 3], - completion_logprobs=[-1.0, -1.0, -1.0], - ) - apply_filters([], [rollout]) - assert rollout.filter_results == {} - assert rollout.is_filtered is False - assert rollout.reward == 1.0 - - -def test_apply_filters_mixed_batch(): - gibberish_filter = _make_gibberish_filter(enforce=True) - - clean = _make_rollout(completion_ids=[50], completion_logprobs=[-1.0], reward=1.0) - dirty = _make_rollout( - completion_ids=[120_000], completion_logprobs=[gibberish_filter.logprob_threshold - 1.0], reward=1.0 - ) - - apply_filters([gibberish_filter], [clean, dirty]) - - assert clean.reward == 1.0 - assert dirty.reward == 1.0 - assert clean.is_filtered is False - assert dirty.is_filtered is True - - -def test_apply_filters_enforced_preserves_rollout_tokens(): - gibberish_filter = _make_gibberish_filter(enforce=True) - - rollout = _make_rollout( - completion_ids=[10, 120_000, 30], - completion_logprobs=[-1.0, gibberish_filter.logprob_threshold - 1.0, -0.5], - reward=1.0, - ) - - apply_filters([gibberish_filter], [rollout]) - - assert rollout.nodes[0].token_ids == [10, 120_000, 30] - assert rollout.nodes[0].logprobs == [ - -1.0, - gibberish_filter.logprob_threshold - 1.0, - -0.5, - ] - assert rollout.nodes[0].mask == [True, True, True] - assert rollout.is_filtered is True - - -def test_apply_filters_preserves_existing_stop_condition(): - gibberish_filter = _make_gibberish_filter(enforce=True) - - rollout = _make_rollout( - completion_ids=[120_000], - completion_logprobs=[gibberish_filter.logprob_threshold - 1.0], - reward=1.0, - ) - rollout.stop_condition = "generation_truncated" - - apply_filters([gibberish_filter], [rollout]) - - assert rollout.stop_condition == "generation_truncated" - assert rollout.is_filtered is True - - -# --- apply_filters tests (monitor-only, enforce=False) --- - - -def test_apply_filters_monitor_only_tracks_detection(): - gibberish_filter = _make_gibberish_filter(enforce=False) - - rollout = _make_rollout( - completion_ids=[120_000], - completion_logprobs=[gibberish_filter.logprob_threshold - 1.0], - reward=1.0, - ) - - apply_filters([gibberish_filter], [rollout]) - - assert rollout.reward == 1.0 - assert all(rollout.nodes[0].mask) - assert rollout.stop_condition is None - assert rollout.filter_results == {"gibberish": True} - assert rollout.is_filtered is False - - -def test_apply_filters_monitor_only_mixed_batch(): - gibberish_filter = _make_gibberish_filter(enforce=False) - - clean = _make_rollout(completion_ids=[50], completion_logprobs=[-1.0], reward=1.0) - dirty = _make_rollout( - completion_ids=[120_000], completion_logprobs=[gibberish_filter.logprob_threshold - 1.0], reward=1.0 - ) - - apply_filters([gibberish_filter], [clean, dirty]) - - assert clean.reward == 1.0 - assert dirty.reward == 1.0 - assert clean.is_filtered is False - assert dirty.is_filtered is False diff --git a/tests/unit/orchestrator/test_metrics.py b/tests/unit/orchestrator/test_metrics.py index 546e594d9a..22b63594d6 100644 --- a/tests/unit/orchestrator/test_metrics.py +++ b/tests/unit/orchestrator/test_metrics.py @@ -32,8 +32,7 @@ def mk( group_id: str = "g0", trainable: bool = True, is_trainable: bool = True, - is_filtered: bool = False, - filter_results: dict | None = None, + is_admitted: bool = True, setup: float = 0.0, agent: float = 0.0, agent_model: float = 0.0, @@ -64,8 +63,7 @@ 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_admitted=is_admitted, timing=SimpleNamespace( setup=SimpleNamespace(duration=setup), agent=SimpleNamespace( @@ -93,12 +91,17 @@ 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_admitted=False), + 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 isinstance(eff, TrainRollouts) and len(eff) == 2 + assert all(r.is_admitted and not r.has_error and r in rc.rollouts for r in eff) 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 +165,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_admitted=False)]) 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 +237,15 @@ 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_admitted=False), + mk(is_trainable=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_admitted/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 key or "is_admitted" in key for key in eval_out) def test_eval_avg_at_k_and_pass_k():