feat: reward shaping layer with success-gated length penalty

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

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[<type>]`; 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/<type>`.
 
-[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.

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). |
 

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

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

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

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

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]

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."""

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

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)

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)