diff --git a/configs/debug/concurrency.toml b/configs/debug/concurrency.toml index a6f3db076e..3215eacc88 100644 --- a/configs/debug/concurrency.toml +++ b/configs/debug/concurrency.toml @@ -14,8 +14,8 @@ # skip-optimizer checkpointing. Router replay is enabled — GLM-4.5-Air is a # glm4_moe model, which supports routed_experts. # -# GRPO uses the linear length penalty with the input-token and turns -# coefficients set to 0.1. +# Reward shaping: the linear length penalty with the output-token term off +# and the input-token and turns coefficients set to 0.1. max_steps = 1000 seq_len = 131072 @@ -94,8 +94,8 @@ max_off_policy_steps = 32 [orchestrator.algo] type = "grpo" -[orchestrator.algo.length_penalty] -type = "linear" +[[orchestrator.reward_shaping]] +type = "length_penalty" num_output_tokens_weight = 0.0 num_input_tokens_weight = 0.1 num_turns_weight = 0.1 diff --git a/deps/verifiers b/deps/verifiers index d4f3d5a8c6..fd50b582b6 160000 --- a/deps/verifiers +++ b/deps/verifiers @@ -1 +1 @@ -Subproject commit d4f3d5a8c6bae93c874f78d00eb89fd97c4b0275 +Subproject commit fd50b582b60f25893262012093106c931137750f diff --git a/docs/algorithms.md b/docs/algorithms.md index a0910f060d..62bcac7d23 100644 --- a/docs/algorithms.md +++ b/docs/algorithms.md @@ -143,7 +143,7 @@ At runtime, each env's resolved config builds two objects: a `GenerationSource` | `algo.type` | Class | hook(s) — stage | |---|---|---| -| `grpo` | `GRPOAlgorithm` | `score_group`: group-norm credit (optional length penalty) | +| `grpo` | `GRPOAlgorithm` | `score_group`: group-norm credit | | `echo` | `EchoAlgorithm` | `score_episode`: weighted ce on observation tokens; `score_group`: group-norm credit (inherited) | | `max_rl` | `MaxRLAlgorithm` | `score_group`: mean-normalized group credit | | `rae` | `RAEAlgorithm` | `score_group`: per-agent EMA-baseline credit | @@ -159,7 +159,7 @@ Algorithms operate on native verifier artifacts and annotate their message graph The pipeline drives these through `finalize_episode` and `finalize_group`. Advantages, reference logprobs, and named loss weights stay on verifier nodes through admission. Only admitted traces are flattened into `TrainingSample`s. -Class-level declarations state what the algorithm needs: which loss component its action tokens feed (`action_loss_type`). Every class is constructed with its algorithm config plus the one host-owned resource it can't rebuild — the live policy clients (`self.clients`). Everything else an algorithm needs it builds from its own config in `setup()`: `opd` connects its frozen `teacher`; `opsd` builds the renderer for its demonstration hint (tokenizer is always the live policy's — self-distillation has no separate model). The pipeline only ever calls the two `finalize_*` methods — writing your own algorithm is subclassing `Algorithm` and overriding the hooks its signal needs (see [Authoring an Algorithm](#authoring-an-algorithm)). Shared math (efficiency shaping, prefill alignment) lives as plain functions in `prime_rl.orchestrator.algo.advantage`. +Class-level declarations state what the algorithm needs: which loss component its action tokens feed (`action_loss_type`). Every class is constructed with its algorithm config plus the one host-owned resource it can't rebuild — the live policy clients (`self.clients`). Everything else an algorithm needs it builds from its own config in `setup()`: `opd` connects its frozen `teacher`; `opsd` builds the renderer for its demonstration hint (tokenizer is always the live policy's — self-distillation has no separate model). The pipeline only ever calls the two `finalize_*` methods — writing your own algorithm is subclassing `Algorithm` and overriding the hooks its signal needs (see [Authoring an Algorithm](#authoring-an-algorithm)). Reward shaping (e.g. the length penalty) is not an algorithm concern — it runs before the algorithm, see [Reward Shaping](#reward-shaping). ## Async / Off-Policy Training @@ -286,7 +286,7 @@ The per-token training signal is set by `algo.type` and the [algorithm](#the-alg | Type | Component | Effect | |---|---|---| -| `grpo` | `rl` | Group-norm: reward minus per-group baseline, optional length penalty. | +| `grpo` | `rl` | Group-norm: reward minus per-group baseline. | | `max_rl` | `rl` | Mean-normalized group credit (maximum-likelihood RL). | | `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. | @@ -301,16 +301,20 @@ The default advantage is per-group reward minus per-group baseline (DR-GRPO with This is intentionally simple — it does the right thing for most envs. Write a named algorithm class when you need group-aware shaping that depends on trajectory metadata (sub-agent rollouts, relative-rank shaping, …) — see [Authoring an Algorithm](#authoring-an-algorithm). -A **length penalty** (`length_penalty` on the `grpo`-family algorithms) can be layered on top to discourage rambling. The `linear` penalty subtracts a single `pass_rate`-scaled penalty from each reward before the GRPO baseline, combining output tokens (`num_output_tokens_weight`), input / context tokens (`num_input_tokens_weight`), and turns (`num_turns_weight`) — each normalized by the group's own max for that quantity, with `num_input_tokens_weight` and `num_turns_weight` defaulting to `0.1`. +### Reward Shaping -```toml -[orchestrator.algo] -type = "grpo" +Reward shaping changes *what* is rewarded; the algorithm decides how that reward becomes per-token credit. Shapers are configured as a list under `[[orchestrator.reward_shaping]]` (per-env override: `reward_shaping` on the source, `[]` to disable) and run on every finalized group before the env's algorithm scores it. Each shaper records an additive term on every trainable trace in `trace.reward_shaping[]`; algorithms score against `training_reward(trace)` — the env reward plus those terms — while `trace.reward` stays the env's verdict, so metrics, curricula and evals keep reading task success. Every algorithm that reads rewards (`grpo`, `echo`, `max_rl`, `rae`, `hierarchical_grpo`) sees shaped rewards. Each shaper's term is logged under `reward_shaping/`. -[orchestrator.algo.length_penalty] -type = "linear" +The **length penalty** (`type = "length_penalty"`) discourages rambling: it subtracts a single `pass_rate`-scaled penalty (where `pass_rate` is the group's mean reward) combining output tokens (`num_output_tokens_weight`, default `0.25`), input / context tokens (`num_input_tokens_weight`, default `0.1`), and turns (`num_turns_weight`, default `0.1`) — each normalized by the group's own max for that quantity. Only rollouts with `reward >= success_threshold` (default `1.0`) pay it: a wrong answer is not improved by being shorter, and penalizing failed long rollouts would teach the policy to give up early. Lower the threshold for envs with partial-credit rewards. + +```toml +[[orchestrator.reward_shaping]] +type = "length_penalty" +num_output_tokens_weight = 0.0 # penalize context and turns only ``` +A new shaper is a named class in `prime_rl.orchestrator.reward_shaping` (subclass `RewardShaper`, implement `shape_group`, register it in `SHAPER_CLASSES`) with a `type`-keyed config in `prime_rl.configs.reward_shaping`. + ### Hierarchical GRPO GRPO gives each rollout its reward minus the average reward of comparable rollouts. In an ordinary single-agent group, every rollout answers the same task, so one group average is enough. diff --git a/docs/training.md b/docs/training.md index 8bb2c656f6..49e924723b 100644 --- a/docs/training.md +++ b/docs/training.md @@ -61,6 +61,7 @@ A condensed view of the knobs you'll most often tune. For trainer-side paralleli | `orchestrator.group_size` | Rollouts generated per task. | | `orchestrator.max_off_policy_steps` | Maximum staleness of a trained rollout (default 8): the version a batch trains on minus the oldest version that generated the rollout, queue time included. Episodes past the bound are dropped; a group shares one dispatch version, so its episodes age out together. The main off-policy dial on long agentic rollouts — bump for throughput, lower for tighter on-policyness. Watch `off_policy/*` and `mismatch_kl/all/mean` when tuning. | | `[orchestrator.algo]` | Training algorithm — its `type` names it (`grpo` default, `max_rl`, `rae`, `hierarchical_grpo`, `opd`, `opsd`, `sft`, `echo`). See [Algorithms](#algorithms). | +| `[[orchestrator.reward_shaping]]` | Reward shapers applied before credit assignment (`length_penalty`). See [Reward Shaping](algorithms.md#reward-shaping). | | `[[orchestrator.train.source]]` | Training sources. List multiple tables for multi-env training; weight them via `ratio`. See [Configuration § Training sources](configuration.md#training-sources-orchestratortrainsource). | | `[[orchestrator.eval.source]]` + `orchestrator.eval.interval` | Eval environments and cadence (default every 100 steps). | diff --git a/examples/advanced/glm-4.5-air/search.toml b/examples/advanced/glm-4.5-air/search.toml index 9610214d76..edca289a18 100644 --- a/examples/advanced/glm-4.5-air/search.toml +++ b/examples/advanced/glm-4.5-air/search.toml @@ -84,8 +84,8 @@ thinking_retention = "all" [orchestrator.algo] type = "grpo" -[orchestrator.algo.length_penalty] -type = "linear" +[[orchestrator.reward_shaping]] +type = "length_penalty" num_output_tokens_weight = 0.0 num_input_tokens_weight = 0.1 num_turns_weight = 0.0 diff --git a/examples/advanced/glm-4.5-air/swe.toml b/examples/advanced/glm-4.5-air/swe.toml index 05f69f1e8c..6d36c5835a 100644 --- a/examples/advanced/glm-4.5-air/swe.toml +++ b/examples/advanced/glm-4.5-air/swe.toml @@ -92,8 +92,8 @@ max_off_policy_steps = 32 [orchestrator.algo] type = "grpo" -[orchestrator.algo.length_penalty] -type = "linear" +[[orchestrator.reward_shaping]] +type = "length_penalty" num_output_tokens_weight = 0.0 num_input_tokens_weight = 0.1 num_turns_weight = 0.1 diff --git a/examples/advanced/glm-4.5-air/terminal.toml b/examples/advanced/glm-4.5-air/terminal.toml index 1afece9609..55f0af605b 100644 --- a/examples/advanced/glm-4.5-air/terminal.toml +++ b/examples/advanced/glm-4.5-air/terminal.toml @@ -85,8 +85,8 @@ thinking_retention = "all" [orchestrator.algo] type = "grpo" -[orchestrator.algo.length_penalty] -type = "linear" +[[orchestrator.reward_shaping]] +type = "length_penalty" num_output_tokens_weight = 0.0 num_input_tokens_weight = 0.1 num_turns_weight = 0.1 diff --git a/examples/advanced/nemotron-3-super/swe.toml b/examples/advanced/nemotron-3-super/swe.toml index a09b1654fa..909647c9b5 100644 --- a/examples/advanced/nemotron-3-super/swe.toml +++ b/examples/advanced/nemotron-3-super/swe.toml @@ -75,8 +75,8 @@ max_inflight = 1536 [orchestrator.algo] type = "grpo" -[orchestrator.algo.length_penalty] -type = "linear" +[[orchestrator.reward_shaping]] +type = "length_penalty" num_output_tokens_weight = 0.0 num_input_tokens_weight = 0.1 num_turns_weight = 0.1 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 739cd6e4bb..2af4343c47 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/algorithm.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/algorithm.py @@ -88,28 +88,10 @@ class SamplingConfig(BaseConfig): # --------------------------------------------------------------------------- -# Shared sub-configs (length penalty, echo roles) +# Shared sub-configs (echo roles) # --------------------------------------------------------------------------- -class LinearLengthPenaltyConfig(BaseConfig): - """Linear ``pass_rate``-scaled penalty subtracted from each reward before the GRPO baseline — the sum of three terms (completion tokens, input tokens, turns), each normalized by the group's own max for that quantity and disabled by setting its coefficient to 0.""" - - type: Literal["linear"] = "linear" - - num_output_tokens_weight: float = Field(0.25, ge=0, allow_inf_nan=False) - """Scale on the output-token term. Each reward is reduced by ``num_output_tokens_weight * pass_rate * (rollout num_output_tokens / group's max num_output_tokens)`` — where ``pass_rate`` is the group's mean reward — before the GRPO baseline subtraction. Finite and non-negative; 0 disables the term.""" - - num_input_tokens_weight: float = Field(0.1, ge=0, allow_inf_nan=False) - """Scale on the input-token term — tokens the model conditioned on but did not generate (``num_total_tokens - num_output_tokens``: prompts, tool responses), as a fraction of the group's max input tokens. 0 disables the term.""" - - num_turns_weight: float = Field(0.1, ge=0, allow_inf_nan=False) - """Scale on the turns term (``pass_rate * (rollout num_turns / group's max num_turns)``). 0 disables the term.""" - - -LengthPenaltyConfig: TypeAlias = LinearLengthPenaltyConfig - - class EchoRoleConfig(BaseConfig): """Echo CE supervision for one message role.""" @@ -204,9 +186,6 @@ class GRPOAlgoConfig(BaseAlgoConfig): action_loss_type: ClassVar[ActionLossType] = "rl" - length_penalty: LengthPenaltyConfig | None = None - """Linear length penalty subtracted from each reward before the GRPO baseline (see ``LinearLengthPenaltyConfig``): a ``pass_rate``-scaled sum of output-token, input-token, and turns terms, each normalized by the group's own max for that quantity. None disables it.""" - class EchoAlgoConfig(GRPOAlgoConfig): type: Literal["echo"] = "echo" # type: ignore[assignment] 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 a6071edc0d..3d342144db 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py @@ -10,6 +10,7 @@ GRPOAlgoConfig, ) from prime_rl.configs.monitors import OrchestratorMonitorsConfig +from prime_rl.configs.reward_shaping import RewardShaperConfig from prime_rl.configs.shared import ( BaseModelConfig, BaseWeightBroadcastConfig, @@ -239,6 +240,12 @@ class TrainSourceConfig(EnvConfig): ``orchestrator.algo`` when unset; set ``type`` (and its params) to give this env its own algorithm.""" + reward_shaping: list[RewardShaperConfig] | None = None + """Reward shapers for this env, applied to the env reward before the + algorithm assigns credit. Inherits from the top-level + ``orchestrator.reward_shaping`` when unset; ``[]`` disables shaping for + this env.""" + curriculum: CurriculumConfig | None = None """User-authored task sampler and admission gates. The default cycles through the taskset and admits every finalized group.""" @@ -446,6 +453,12 @@ class OrchestratorConfig(BaseConfig): Defaults to ``grpo``. Override per source via ``[[orchestrator.train.source]]``'s ``algo``.""" + reward_shaping: list[RewardShaperConfig] = Field(default_factory=list) + """Reward shapers applied to every finalized group before the algorithm + assigns credit — additive terms on the env reward (``type`` names the + shaper, e.g. ``length_penalty``). Empty by default. Override per source via + ``[[orchestrator.train.source]]``'s ``reward_shaping``.""" + model: ModelConfig = ModelConfig() """The model being trained: its model fields plus the client of the live vLLM deployment (``[orchestrator.model] name = ...`` with @@ -557,6 +570,14 @@ def inherit_env_algorithms(self): env_cfg.algo = self.algo.model_copy(deep=True) return self + @model_validator(mode="after") + def inherit_env_reward_shaping(self): + """Envs without their own reward shaping inherit the top-level list.""" + for env_cfg in self.train.source: + if env_cfg.reward_shaping is None: + env_cfg.reward_shaping = [shaper.model_copy(deep=True) for shaper in self.reward_shaping] + return self + @model_validator(mode="after") def validate_env_algorithms(self): """Let each algorithm reject environments it cannot score correctly.""" diff --git a/skills/configs/SKILL.md b/skills/configs/SKILL.md index a7f0abfd18..f62dd83e24 100644 --- a/skills/configs/SKILL.md +++ b/skills/configs/SKILL.md @@ -71,6 +71,8 @@ The `sft` entrypoint takes the same eval shape at the top level for online evals **Algorithms** — `[orchestrator.algo] type = "grpo" | "max_rl" | "rae" | "hierarchical_grpo" | "opd" | "opsd" | "sft" | "echo"` — the type names the algorithm (credit assignment + loss routing, fused), and each type's class defaults are its vetted setting; any other key you set is your own assembly (e.g. `[orchestrator.algo.roles.user] alpha = 0.1` for echo — setting any echo role replaces the whole role table). `hierarchical_grpo` is only valid with a proposer-solver env: it compares solvers with attempts on the same proposed problem and proposers with other proposals in the group. There is no preset layer, and no config hook that points at user code — a new algorithm is a named class in the repo (subclass `Algorithm`, register it). Per-source override: `[orchestrator.train.source.algo] type = "opd"` (the source assembles its own algorithm). prime-rl only hosts the trainable policy; frozen models are inline external endpoints on the algorithm, named where the model is used — `[orchestrator.algo.teacher]` for opd (the frozen model scored against), `[orchestrator.algo.sampling.source]` for sft (the model it samples from), each with `name` + `base_url`. There is no shared `teacher` slot. opsd declares no model — it self-distills against the live policy. See `docs/algorithms.md`. +**Reward shaping** — `[[orchestrator.reward_shaping]] type = "length_penalty"` (a list; per-source override via `reward_shaping` on the source, `[]` disables). Shapers run before the algorithm on every finalized group and add terms to the training reward; `trace.reward` stays the env's verdict. The length penalty is success-gated (`success_threshold`, default `1.0`) and pass-rate scaled — see `docs/algorithms.md#reward-shaping`. + **`BaseModel | None` fields** — bare flag enables defaults; nested override enables and sets: ```bash diff --git a/src/prime_rl/orchestrator/algo/grpo.py b/src/prime_rl/orchestrator/algo/grpo.py index 232d3c48b4..be1ae31ef8 100644 --- a/src/prime_rl/orchestrator/algo/grpo.py +++ b/src/prime_rl/orchestrator/algo/grpo.py @@ -1,45 +1,20 @@ from __future__ import annotations -from typing import TYPE_CHECKING - import torch import verifiers.v1 as vf -from prime_rl.configs.algorithm import GRPOAlgoConfig from prime_rl.orchestrator.algo.base import Algorithm, iter_trainable_traces -from prime_rl.orchestrator.algo.routing import assign_advantages - -if TYPE_CHECKING: - from prime_rl.orchestrator.clients import InferenceClient +from prime_rl.orchestrator.algo.routing import assign_advantages, training_reward class GRPOAlgorithm(Algorithm): """Group Relative Policy Optimization: sample a group of rollouts from the - policy per example; credit = reward minus the group mean (optionally - length-shaped); action tokens feed the ``rl`` loss.""" - - def __init__(self, config: GRPOAlgoConfig, clients: InferenceClient): - super().__init__(config, clients) - self.length_penalty = config.length_penalty + policy per example; credit = training reward minus the group mean; action + tokens feed the ``rl`` loss.""" async def score_group(self, episodes: list[vf.Episode]) -> None: traces = [trace for _, trace in iter_trainable_traces(episodes)] - rewards = torch.tensor([trace.reward for trace in traces], dtype=torch.float32) - length_penalty = self.length_penalty - if length_penalty is None: - advantages = rewards - rewards.mean() - else: - output = torch.tensor([trace.num_output_tokens for trace in traces], dtype=rewards.dtype) - total = torch.tensor([trace.num_total_tokens for trace in traces], dtype=rewards.dtype) - turns = torch.tensor([trace.num_turns for trace in traces], dtype=rewards.dtype) - input = total - output - penalty_frac = ( - length_penalty.num_output_tokens_weight * (output / output.max().clamp(min=1)) - + length_penalty.num_input_tokens_weight * (input / input.max().clamp(min=1)) - + length_penalty.num_turns_weight * (turns / turns.max().clamp(min=1)) - ) - penalty = rewards.mean() * penalty_frac - shaped_rewards = rewards - penalty - advantages = shaped_rewards - shaped_rewards.mean() + rewards = torch.tensor([training_reward(trace) for trace in traces], dtype=torch.float32) + advantages = rewards - rewards.mean() for trace, advantage in zip(traces, advantages.tolist(), strict=True): assign_advantages(trace, advantage) diff --git a/src/prime_rl/orchestrator/algo/hierarchical_grpo.py b/src/prime_rl/orchestrator/algo/hierarchical_grpo.py index dc0c508c00..5081e92df0 100644 --- a/src/prime_rl/orchestrator/algo/hierarchical_grpo.py +++ b/src/prime_rl/orchestrator/algo/hierarchical_grpo.py @@ -7,7 +7,7 @@ from prime_rl.configs.algorithm import HierarchicalGRPOAlgoConfig from prime_rl.orchestrator.algo.base import Algorithm, iter_trainable_traces -from prime_rl.orchestrator.algo.routing import assign_advantages +from prime_rl.orchestrator.algo.routing import assign_advantages, training_reward if TYPE_CHECKING: from prime_rl.orchestrator.clients import InferenceClient @@ -37,6 +37,7 @@ async def score_group(self, episodes: list[vf.Episode]) -> None: key = (trace.agent.name, episode.id if episode_scoped else None) peers[key].append(trace) for members in peers.values(): - baseline = sum(trace.reward for trace in members) / len(members) + rewards = {trace.id: training_reward(trace) for trace in members} + baseline = sum(rewards.values()) / len(members) for trace in members: - assign_advantages(trace, trace.reward - baseline) + assign_advantages(trace, rewards[trace.id] - baseline) diff --git a/src/prime_rl/orchestrator/algo/max_rl.py b/src/prime_rl/orchestrator/algo/max_rl.py index 2eb93c4b8d..0c1ea25007 100644 --- a/src/prime_rl/orchestrator/algo/max_rl.py +++ b/src/prime_rl/orchestrator/algo/max_rl.py @@ -4,7 +4,7 @@ import verifiers.v1 as vf from prime_rl.orchestrator.algo.base import Algorithm, iter_trainable_traces -from prime_rl.orchestrator.algo.routing import assign_advantages +from prime_rl.orchestrator.algo.routing import assign_advantages, training_reward class MaxRLAlgorithm(Algorithm): @@ -21,7 +21,7 @@ class MaxRLAlgorithm(Algorithm): async def score_group(self, episodes: list[vf.Episode]) -> None: traces = [trace for _, trace in iter_trainable_traces(episodes)] - rewards = torch.tensor([trace.reward for trace in traces], dtype=torch.float32) + rewards = torch.tensor([training_reward(trace) for trace in traces], dtype=torch.float32) mean = rewards.mean() advantages = torch.zeros_like(rewards) if mean <= 0 else (rewards - mean) / mean for trace, advantage in zip(traces, advantages.tolist(), strict=True): diff --git a/src/prime_rl/orchestrator/algo/rae.py b/src/prime_rl/orchestrator/algo/rae.py index bbf0689458..0697893cd7 100644 --- a/src/prime_rl/orchestrator/algo/rae.py +++ b/src/prime_rl/orchestrator/algo/rae.py @@ -7,7 +7,7 @@ from prime_rl.configs.algorithm import RAEAlgoConfig from prime_rl.orchestrator.algo.base import Algorithm, iter_trainable_traces -from prime_rl.orchestrator.algo.routing import assign_advantages +from prime_rl.orchestrator.algo.routing import assign_advantages, training_reward if TYPE_CHECKING: from prime_rl.orchestrator.clients import InferenceClient @@ -38,6 +38,7 @@ def __init__(self, config: RAEAlgoConfig, clients: InferenceClient): async def score_group(self, episodes: list[vf.Episode]) -> None: for _, trace in iter_trainable_traces(episodes): + reward = training_reward(trace) baseline = self.baselines[trace.agent.name] - assign_advantages(trace, trace.reward - baseline) - self.baselines[trace.agent.name] = self.decay * baseline + (1.0 - self.decay) * trace.reward + assign_advantages(trace, reward - baseline) + self.baselines[trace.agent.name] = self.decay * baseline + (1.0 - self.decay) * reward diff --git a/src/prime_rl/orchestrator/algo/routing.py b/src/prime_rl/orchestrator/algo/routing.py index f1322f4168..ab448febea 100644 --- a/src/prime_rl/orchestrator/algo/routing.py +++ b/src/prime_rl/orchestrator/algo/routing.py @@ -50,6 +50,15 @@ def assign_reference_logprobs(branch: vf.Branch, values: list[float]) -> None: offset = end +def record_shaping(trace: vf.Trace, name: str, value: float) -> None: + trace.reward_shaping[name] = float(value) + + +def training_reward(trace: vf.Trace) -> float: + """The reward an algorithm assigns credit against: env reward plus every shaping term.""" + return trace.reward + sum(trace.reward_shaping.values()) + + def scalar_advantage(trace: vf.Trace) -> float | None: """Mean nonzero token advantage, or zero for an assigned-zero trace.""" advantages = [value for node in trace.nodes for value in node.advantages or []] diff --git a/src/prime_rl/orchestrator/envs.py b/src/prime_rl/orchestrator/envs.py index 55ddbdf2b2..5397264628 100644 --- a/src/prime_rl/orchestrator/envs.py +++ b/src/prime_rl/orchestrator/envs.py @@ -31,6 +31,7 @@ from prime_rl.configs.orchestrator import EnvConfig, EvalSourceConfig, TrainSourceConfig from prime_rl.orchestrator.algo import Algorithm, build_algorithm from prime_rl.orchestrator.generation_source import GenerationSource +from prime_rl.orchestrator.reward_shaping import RewardShaper, build_reward_shapers from prime_rl.utils.logger import format_time, get_logger # Max wait for the env server to answer health. Generous because the launcher spawns @@ -126,10 +127,12 @@ def __init__( address: str, generation_source: GenerationSource, algorithm: Algorithm, + reward_shapers: list[RewardShaper], ): super().__init__(config, address) self.generation_source = generation_source self.algorithm = algorithm + self.reward_shapers = reward_shapers self.sampling_args = generation_source.sampling_args(config.sampling.to_sampling_args()) @@ -198,12 +201,16 @@ def __init__( self._envs: dict[str, TrainEnv] = {} for config in configs: assert config.algo is not None, "TrainSourceConfig.algo must be resolved before env construction" + assert config.reward_shaping is not None, ( + "TrainSourceConfig.reward_shaping must be resolved before env construction" + ) get_logger().info(f"Initializing {config.algo.type} algorithm for {config.resolved_name}") env = TrainEnv( config, addresses[("train", config.resolved_name)], GenerationSource(config.algo.sampling, clients, renderer_config), build_algorithm(config.algo, clients), + build_reward_shapers(config.reward_shaping), ) self._envs[env.name] = env diff --git a/src/prime_rl/orchestrator/metrics.py b/src/prime_rl/orchestrator/metrics.py index 3eb25991f5..9eec4c5e83 100644 --- a/src/prime_rl/orchestrator/metrics.py +++ b/src/prime_rl/orchestrator/metrics.py @@ -2,7 +2,7 @@ from __future__ import annotations -from collections.abc import Callable, Iterator +from collections.abc import Callable, Iterator, Mapping from dataclasses import dataclass from typing import Any, Literal @@ -153,21 +153,21 @@ class CustomMetrics(StatGroup): def __init__( self, records: list[TraceRecord], - attr: str, + source: Callable[[vf.Trace], Mapping[str, Any]], value: Callable[[Any], float] = float, ) -> None: super().__init__(records) - self.attr = attr + self.source = source self.value = value def stats(self) -> dict[str, Stat]: - names = sorted({name for trace in self.traces for name in getattr(trace, self.attr)}) + names = sorted({name for trace in self.traces for name in self.source(trace)}) return { name: Stat( [ self.value(scores[name]) if scores[name] is not None else 0.0 for trace in self.traces - if name in (scores := getattr(trace, self.attr)) + if name in (scores := self.source(trace)) ] ) for name in names @@ -197,11 +197,15 @@ def timing(self) -> TimingMetrics: @property def metrics(self) -> CustomMetrics: - return CustomMetrics(self.records, "metrics") + return CustomMetrics(self.records, lambda trace: trace.metrics) @property def rewards(self) -> CustomMetrics: - return CustomMetrics(self.records, "rewards", value=lambda reward: reward.value) + return CustomMetrics(self.records, lambda trace: trace.rewards, value=lambda reward: reward.value) + + @property + def reward_shaping(self) -> CustomMetrics: + return CustomMetrics(self.records, lambda trace: trace.reward_shaping) @property def has_error(self) -> Stat: @@ -251,6 +255,7 @@ def to_dict(self, prefix: str, *, subset: Subset) -> dict[str, float]: out |= self.timing.to_dict(f"{prefix}/timing") out |= self.metrics.to_dict(f"{prefix}/metrics") out |= self.rewards.to_dict(f"{prefix}/rewards") + out |= self.reward_shaping.to_dict(f"{prefix}/reward_shaping") if subset == "all": out[f"{prefix}/has_error/mean"] = self.has_error.mean() out[f"{prefix}/cancelled/mean"] = self.cancelled.mean() diff --git a/src/prime_rl/orchestrator/train_sink.py b/src/prime_rl/orchestrator/train_sink.py index bfba54d403..f18b615f5c 100644 --- a/src/prime_rl/orchestrator/train_sink.py +++ b/src/prime_rl/orchestrator/train_sink.py @@ -237,6 +237,8 @@ async def process_group(self, group_id: str) -> None: survivors = [trace for _, trace in iter_trainable_traces(group)] if survivors: + for shaper in env.reward_shapers: + shaper.shape_group(group) await env.algorithm.finalize_group(group) admitted = self._admit(group) if group else False if not survivors or not admitted: diff --git a/tests/unit/orchestrator/test_advantage.py b/tests/unit/orchestrator/test_advantage.py index f8aa3a5c75..51c04f1e35 100644 --- a/tests/unit/orchestrator/test_advantage.py +++ b/tests/unit/orchestrator/test_advantage.py @@ -3,14 +3,12 @@ import pytest import verifiers.v1 as vf -from prime_rl.configs.algorithm import ( - GRPOAlgoConfig, - LinearLengthPenaltyConfig, - MaxRLAlgoConfig, -) +from prime_rl.configs.algorithm import GRPOAlgoConfig, MaxRLAlgoConfig +from prime_rl.configs.reward_shaping import LengthPenaltyConfig from prime_rl.orchestrator.algo.grpo import GRPOAlgorithm from prime_rl.orchestrator.algo.max_rl import MaxRLAlgorithm -from prime_rl.orchestrator.algo.routing import assign_advantages +from prime_rl.orchestrator.algo.routing import assign_advantages, training_reward +from prime_rl.orchestrator.reward_shaping import LengthPenalty from prime_rl.orchestrator.trajectories import trace_to_samples @@ -151,9 +149,12 @@ def _scalar(episode: vf.Episode) -> float: raise AssertionError("episode has no trainable token") -def _grpo(group: list[vf.Episode], length_penalty=None) -> list[float]: - """Drive ``GRPOAlgorithm.score_group`` and read back each per-rollout scalar.""" - algo = GRPOAlgorithm(GRPOAlgoConfig(length_penalty=length_penalty), clients=None) +def _grpo(group: list[vf.Episode], length_penalty: LengthPenaltyConfig | None = None) -> list[float]: + """Shape the group (if asked), drive ``GRPOAlgorithm.score_group`` and read back + each per-rollout scalar — the sink's order: shapers first, then credit.""" + if length_penalty is not None: + LengthPenalty(length_penalty).shape_group(group) + algo = GRPOAlgorithm(GRPOAlgoConfig(), clients=None) asyncio.run(algo.score_group(group)) return [_scalar(episode) for episode in group] @@ -191,47 +192,50 @@ def test_max_rl_mean_normalized(): # -------------------------------------------------------------------------- -# GRPO linear length penalty: pass_rate-scaled penalty before the baseline. +# Length penalty reward shaping: success-gated, pass_rate-scaled term recorded +# before the baseline; the env reward itself is untouched. # -------------------------------------------------------------------------- -def test_linear_equal_lengths_reduce_to_plain_grpo(): - """Equal completion length and turns → every rollout takes the same penalty - fraction, so subtracting it leaves the centered advantages unchanged.""" - penalized = _grpo( - _make_group(rewards=[1.0, 0.0, 1.0], completion_lengths=[10, 10, 10], num_turns=[2, 2, 2]), - length_penalty=LinearLengthPenaltyConfig(), - ) - plain = _grpo(_make_group(rewards=[1.0, 0.0, 1.0], completion_lengths=[10, 10, 10], num_turns=[2, 2, 2])) +def test_length_penalty_equal_lengths_reduce_to_plain_grpo(): + """Equal completion length and turns → every successful rollout takes the same + penalty fraction, so subtracting it leaves the centered advantages unchanged.""" + group = _make_group(rewards=[1.0, 1.0, 1.0], completion_lengths=[10, 10, 10], num_turns=[2, 2, 2]) + penalized = _grpo(group, length_penalty=LengthPenaltyConfig()) + plain = _grpo(_make_group(rewards=[1.0, 1.0, 1.0], completion_lengths=[10, 10, 10], num_turns=[2, 2, 2])) assert penalized == pytest.approx(plain, abs=1e-6) + # shaping lands on the trace as a term, not on the env reward + trace = group[0].traces[0] + assert trace.reward == 1.0 + assert trace.reward_shaping["length_penalty"] < 0 + assert training_reward(trace) == pytest.approx(1.0 + trace.reward_shaping["length_penalty"]) -def test_linear_completion_term_penalizes_longer(): +def test_length_penalty_completion_term_penalizes_longer(): """With only the completion term, longer completions get a larger penalty and a lower advantage; advantages stay zero-mean.""" - cfg = LinearLengthPenaltyConfig(num_output_tokens_weight=0.25, num_input_tokens_weight=0.0, num_turns_weight=0.0) + cfg = LengthPenaltyConfig(num_output_tokens_weight=0.25, num_input_tokens_weight=0.0, num_turns_weight=0.0) advs = _grpo(_make_group(rewards=[1.0, 1.0, 1.0], completion_lengths=[10, 20, 30]), length_penalty=cfg) assert advs[0] > advs[1] > advs[2] assert sum(advs) == pytest.approx(0.0, abs=1e-6) -def test_linear_context_term_penalizes_more_context(): +def test_length_penalty_context_term_penalizes_more_context(): """The context term penalizes non-completion (prompt / tool-response) tokens: at equal completion length, more context tokens yields a lower advantage.""" - cfg = LinearLengthPenaltyConfig(num_output_tokens_weight=0.0, num_input_tokens_weight=0.25, num_turns_weight=0.0) + cfg = LengthPenaltyConfig(num_output_tokens_weight=0.0, num_input_tokens_weight=0.25, num_turns_weight=0.0) group = [ _build_episode(1.0, sampled_lengths=[10], obs_lengths=[]), _build_episode(1.0, sampled_lengths=[10], obs_lengths=[100]), ] - asyncio.run(GRPOAlgorithm(GRPOAlgoConfig(length_penalty=cfg), clients=None).score_group(group)) - advs = [_scalar(episode) for episode in group] + advs = _grpo(group, length_penalty=cfg) assert advs[0] > advs[1] assert sum(advs) == pytest.approx(0.0, abs=1e-6) -def test_linear_turns_term_penalizes_more_turns(): +def test_length_penalty_turns_term_penalizes_more_turns(): """The turns term penalizes higher turn counts at equal token lengths.""" - cfg = LinearLengthPenaltyConfig(num_output_tokens_weight=0.0, num_input_tokens_weight=0.0, num_turns_weight=0.25) + cfg = LengthPenaltyConfig(num_output_tokens_weight=0.0, num_input_tokens_weight=0.0, num_turns_weight=0.25) advs = _grpo( _make_group(rewards=[1.0, 1.0], completion_lengths=[100, 100], num_turns=[1, 4]), length_penalty=cfg, @@ -240,6 +244,23 @@ def test_linear_turns_term_penalizes_more_turns(): assert sum(advs) == pytest.approx(0.0, abs=1e-6) +def test_length_penalty_gates_on_success(): + """Failed rollouts pay no penalty: at the default threshold a long wrong answer + and a short wrong answer get the same advantage, while a long right answer is + penalized relative to a short right one. Lowering the threshold admits + partial-credit rollouts to the penalty.""" + cfg = LengthPenaltyConfig(num_output_tokens_weight=0.25, num_input_tokens_weight=0.0, num_turns_weight=0.0) + advs = _grpo(_make_group(rewards=[0.0, 0.0, 1.0, 1.0], completion_lengths=[10, 100, 10, 100]), length_penalty=cfg) + assert advs[0] == pytest.approx(advs[1]) + assert advs[2] > advs[3] + + group = _make_group(rewards=[0.5, 0.5], completion_lengths=[10, 100]) + assert _grpo(group, length_penalty=cfg) == pytest.approx([0.0, 0.0]) + group = _make_group(rewards=[0.5, 0.5], completion_lengths=[10, 100]) + advs = _grpo(group, length_penalty=cfg.model_copy(update={"success_threshold": 0.5})) + assert advs[0] > advs[1] + + # -------------------------------------------------------------------------- # assign_advantages: scalar broadcast over the rollout's trainable tokens. # -------------------------------------------------------------------------- diff --git a/tests/unit/orchestrator/test_metrics.py b/tests/unit/orchestrator/test_metrics.py index c4cc4718ab..3ed5c26f36 100644 --- a/tests/unit/orchestrator/test_metrics.py +++ b/tests/unit/orchestrator/test_metrics.py @@ -56,6 +56,7 @@ def mk( last_error=SimpleNamespace(type=error_type) if has_error else None, stop_condition=stop_condition, metrics=metrics or {}, + reward_shaping={}, agent=SimpleNamespace(trainable=trainable, name=agent_name), nodes=[SimpleNamespace(advantages=[1.0] if is_trainable else [0.0])], timing=SimpleNamespace( diff --git a/tests/unit/test_configs.py b/tests/unit/test_configs.py index 82fd522c61..a08da2f400 100644 --- a/tests/unit/test_configs.py +++ b/tests/unit/test_configs.py @@ -248,9 +248,10 @@ def test_env_algo_overrides_top_level(): { "renderer": {"name": "qwen3"}, # echo needs the renderer's role attribution "algo": {"type": "echo"}, + "reward_shaping": [{"type": "length_penalty", "num_output_tokens_weight": 0.0}], "train": { "source": [ - {"env": {"taskset": {"id": "reverse-text"}}, "algo": {"type": "grpo"}}, + {"env": {"taskset": {"id": "reverse-text"}}, "algo": {"type": "grpo"}, "reward_shaping": []}, {"env": {"taskset": {"id": "reverse-text"}}, "name": "b"}, ] }, @@ -260,6 +261,10 @@ def test_env_algo_overrides_top_level(): # Env a sets its own algorithm; only env b inherits the top-level echo algorithm. assert env_a.algo is not None and env_a.algo.type == "grpo" assert env_b.algo is not None and env_b.algo.type == "echo" + # Same rule for reward shaping: an explicit [] disables it, unset inherits the top-level list. + assert env_a.reward_shaping == [] + assert env_b.reward_shaping is not None and [s.type for s in env_b.reward_shaping] == ["length_penalty"] + assert env_b.reward_shaping[0].num_output_tokens_weight == 0.0 # Resolved configs round-trip. dumped = config.model_dump(exclude_none=True)