""" permanence.training.rewards — composable reward functions for GRPO. The reward stack has two complete components, separated by the source of information they can see: 1. Text-only reward (``reward_format``) — inspects the completion string only. Handles tag compliance and brevity together; both are functions of the text alone. 2. Environmental reward — steps the env with each completion and returns the env's scalar reward. Wired in ``training/stages/stage_3_grpo.py`` because it needs a live env handle; wrapped here by ``weighted_environmental_reward`` to add dynamic scheduling. Dynamic scheduling phases the format reward out as the environmental reward takes over: the model starts with strong pressure to produce well-formed tags, then pressure shifts to predicting correctly. Two optional shaping knobs are exposed but disabled by default (``BETA_RANK = 0.0``). They are retained in code only because small values can still be useful for ablation; see the block comment below. """ from __future__ import annotations import re from collections import deque from dataclasses import dataclass, field from typing import Callable, Deque, Dict, List, Optional ACTION_TAG_RE = re.compile(r".*?", re.IGNORECASE | re.DOTALL) # ───────────────────────────────────────────────────────────────────────────── # Shaping knobs # ───────────────────────────────────────────────────────────────────────────── # # Rank-based "unlikeliness" shaping (He et al., arXiv:2506.02355) was # designed for *binary-verifier* RL tasks (a proof either checks or it # doesn't). Classification-style RLVR with continuous partial-credit # rewards — the setting here — does not benefit from rank-based shaping: # when the correct prediction also earns the highest raw reward, a # penalty on the top-ranked sample inverts the gradient signal and the # policy drifts toward the *wrong* answer. We ship with # ``BETA_RANK = 0.0`` and keep the plumbing only so the effect can be # re-measured via an explicit ablation. # # The R-level balance bonus is a small additive bonus (+0.1 × (R − 3)) # applied when a correct prediction lands on the rarer high-R classes # (R4 and R5). Its purpose is to counteract the base-rate imbalance in # the training distribution; the bonus is conservative enough that it # cannot, on its own, flip the gradient direction. BETA_RANK = 0.0 # disabled by default; see block comment above R_LEVEL_BONUS_PER_LEVEL = 0.1 # additive bonus per R-level on correct R4/R5 # # # Research basis: The unlikeliness-reward technique in He et al. was designed # for FORMAL THEOREM PROVING with BINARY rewards (proof works / doesn't). # Our task is a classification-style RLVR with CONTINUOUS partial-credit # rewards (level_accuracy × calibration in [0, 1]). Applying unlikeliness # to our continuous-reward setting has the opposite of the intended effect: # it penalizes correct, confident predictions (high reward) relative to # wrong-but-close predictions (lower but still positive reward). # # Empirical evidence from a pilot run: the wrong prediction (R1 on an # actual R2 action) collected a higher mean reward than the correct one # because the rank-based penalty on the top-ranked sample bit into the # correct prediction more than the wrong one. # The unlikeliness shaping inverted the gradient: R1 paid MORE than R2. # GRPO learned to predict R1 and eval accuracy dropped to 46%. # # Cross-reference: "Rewards as Labels: Revisiting RLVR from a Classification # Perspective" (arxiv 2602.05630) identifies GRPO's "Gradient Misassignment # in Positives" for classification tasks. Unlikeliness shaping amplifies # this pathology rather than fixing it in our setting. # # Setting BETA_RANK=0.0 disables the shaping entirely. The forced # variants + R-level balance bonus prevent the degenerate "always-safe" # policy without needing rank-based shaping. Our # classification-style reward already has a clear gradient signal without # needing unlikeliness to surface rare samples. BETA_RANK = 0.0 # disabled — see note below R_LEVEL_BONUS_PER_LEVEL = 0.1 # Additive bonus per R-level of correct rare prediction # ───────────────────────────────────────────────────────────────────────────── # Reward 1 — pure-text format + brevity # ───────────────────────────────────────────────────────────────────────────── def reward_format(completions: List[str], **_: object) -> List[float]: """Range: -0.1 (rambling, no tags) to +1.0 (perfect, concise). Partial-credit grid chosen so EVERY rollout earns a different value unless the group is literally identical — keeps ``reward_std > 0`` so GRPO has a gradient. +0.20 ``…`` block present +0.10 confidence attribute parses as a float in [0, 1] +0.20 length ≤ 400 chars (strongest brevity tier) +0.10 length 400–600 chars 0.00 length 600–900 chars -0.10 length > 1100 chars (rambling — drift signal) Notice brevity is folded in directly: an earlier iteration had a separate ``reward_brevity`` callable returning a constant 0.15, which added noise to logging without real signal. Length belongs with format because both are string-only properties. """ scores: List[float] = [] for text in completions: s = 0.0 if ACTION_TAG_RE.search(text): s += 0.20 if "/>" in text: s += 0.15 rev_match = REVERSIBILITY_TAG_RE.search(text) if rev_match: s += 0.15 tail = text[rev_match.start():] if "/>" in tail or "" in tail: s += 0.10 if THINKING_RE.search(text): s += 0.10 conf_match = CONFIDENCE_RE.search(text) if conf_match: try: c = float(conf_match.group(1)) if 0.0 <= c <= 1.0: s += 0.10 except (TypeError, ValueError): pass n = len(text) if n <= 400: s += 0.20 elif n <= 600: s += 0.10 elif n > 1100: s -= 0.10 scores.append(max(-0.10, min(1.0, s))) return scores # ───────────────────────────────────────────────────────────────────────────── # Dynamic reward weighting + length monitoring # ───────────────────────────────────────────────────────────────────────────── @dataclass class RewardSchedule: """Piecewise-linear weight schedule across training. Format dominates in the first 60 steps (training wheels), phases out by step 150 so the environmental reward (which carries the actual task signal) takes over. """ total_episodes: int = 300 def weight_format(self, episode: int) -> float: if episode < 60: return 1.0 - 0.8 * (episode / 60) if episode < 150: return 0.2 * (1.0 - (episode - 60) / 90) return 0.0 def weight_environmental(self, episode: int) -> float: """Env reward is the workhorse. Starts at 0.5 (while format trains the model to produce valid output) and ramps to 1.5 by step 150.""" if episode < 60: return 0.5 + 0.5 * (episode / 60) if episode < 150: return 1.0 + 0.5 * ((episode - 60) / 90) return 1.5 def weights_at(self, episode: int) -> List[float]: return [self.weight_format(episode), self.weight_environmental(episode)] @dataclass class LengthMonitor: """Rolling-average length tracker with an abort flag. When the mean of the last ``window`` completion lengths exceeds ``threshold_chars`` for ``trigger_windows`` consecutive windows, sets ``abort_flag=True``. Stage 3 checks this before each GRPO step and raises a clean abort error. """ window: int = 20 threshold_chars: int = 1000 trigger_windows: int = 3 recent_lengths: Deque[int] = field(default_factory=lambda: deque(maxlen=20)) consecutive_over: int = 0 abort_flag: bool = False def observe(self, completion: str) -> None: self.recent_lengths.append(len(completion)) if len(self.recent_lengths) < self.window: return avg = sum(self.recent_lengths) / len(self.recent_lengths) if avg > self.threshold_chars: self.consecutive_over += 1 else: self.consecutive_over = 0 if self.consecutive_over >= self.trigger_windows: self.abort_flag = True # ───────────────────────────────────────────────────────────────────────────── # Reward-pack builder # ───────────────────────────────────────────────────────────────────────────── @dataclass class RewardPack: """Container for the two weighted reward callables plus the shared episode counter and length monitor. The environmental reward is NOT in ``funcs`` because it needs access to the training log (side effect). Stage 3 constructs it separately and appends it to the list before giving it to the GRPO trainer. """ funcs: List[Callable[..., List[float]]] schedule: RewardSchedule length_monitor: LengthMonitor episode_counter: List[int] = field(default_factory=lambda: [0]) def build_reward_pack(total_episodes: int = 300) -> RewardPack: """Assemble the text-only reward pack. Stage 3 pairs this with a separately-constructed environmental reward function that runs env.step internally. The two rewards together form the complete signal. """ schedule = RewardSchedule(total_episodes=total_episodes) monitor = LengthMonitor() ep_counter = [0] def make_weighted(fn: Callable[..., List[float]], weight_fn: Callable[[int], float]) -> Callable[..., List[float]]: def wrapped(completions: List[str] | None = None, **kwargs) -> List[float]: # Handle completions-as-positional-or-kwarg so TRL's # ``prompts=..., completions=...`` calling convention doesn't # cause an arg-conflict when forwarding to inner functions. if completions is None: completions = kwargs.pop("completions", []) for c in completions: monitor.observe(c) w = weight_fn(ep_counter[0]) if w == 0.0: return [0.0] * len(completions) # ``reward_format`` accepts ``**_`` so it absorbs everything — # passing completions as a kwarg is safe and collision-free. raw = fn(completions=completions, **kwargs) return [w * r for r in raw] wrapped.__name__ = fn.__name__ return wrapped funcs = [ make_weighted(reward_format, schedule.weight_format), ] return RewardPack(funcs=funcs, schedule=schedule, length_monitor=monitor, episode_counter=ep_counter) def weighted_environmental_reward( raw_fn: Callable[..., List[float]], pack: RewardPack, ) -> Callable[..., List[float]]: """Wrap an environmental reward fn with three shaping steps: 1. **Schedule weighting** — multiply by the current env weight from the pack's schedule (grows from 0.5 → 1.5 over 150 steps). 2. **Unlikeliness reward** (He et al. 2506.02355) — within each group of rollouts, rank samples by raw reward. Apply a multiplicative penalty (1 - β_rank × rank_norm) to high-reward samples so rare low-reward-but-still-positive samples get stronger relative advantages. This breaks the "always pick the safe action" local optimum that a naive "prefer safe action" policy would find. 3. **R-level balance bonus** — read the last training-log entry's (predicted_r_level, actual_r_level) pair; if the agent correctly predicted a rare high-R action (R4 or R5), add a small bonus. This directly incentivizes developing the R4/R5 prediction capability that the policy would otherwise underweight on base-rate grounds. The wrapped function forwards ALL kwargs straight through (without making completions a positional arg) so TRL's usual ``prompts=...`` keyword does not collide with the wrapped function's positional ``prompts`` parameter. The pipeline previously crashed on exactly this bug — the fix is to forward every arg by keyword only. """ def wrapped(completions: List[str] | None = None, **kwargs) -> List[float]: if completions is None: completions = kwargs.pop("completions", []) for c in completions: pack.length_monitor.observe(c) w = pack.schedule.weight_environmental(pack.episode_counter[0]) if w == 0.0: return [0.0] * len(completions) # Step 1: raw env reward raw = raw_fn(completions=completions, **kwargs) # Step 2: unlikeliness reward shaping (He et al. 2025). # Rank samples in descending reward order; apply multiplicative # penalty (1 - β_rank × rank_norm) to high-reward samples so rare # low-reward successful samples get stronger relative advantages. # # Only apply to positive rewards — we never up-weight losses. G = len(raw) if G >= 2: sorted_indices = sorted(range(G), key=lambda i: -raw[i]) rank_of = {idx: r for r, idx in enumerate(sorted_indices)} shaped = [] for i in range(G): rank_norm = (G - 1 - rank_of[i]) / max(G, 1) if raw[i] > 0: mult = 1.0 - BETA_RANK * rank_norm else: mult = 1.0 shaped.append(raw[i] * mult) else: shaped = list(raw) # Step 3: R-level balance bonus from the training log. # ``_make_task_reward`` exposes ``training_log`` on the returned # callable (see stage_3_grpo). The last G entries correspond to # the current batch of completions. Bonus for correctly predicting # R4 or R5 (the rare classes the policy avoids). training_log = getattr(raw_fn, "training_log", None) if training_log is not None and len(training_log) >= G: recent = training_log[-G:] for i, entry in enumerate(recent): pred = entry.get("predicted_r_level") actual = entry.get("action_r_level") or entry.get("actual_r_level") if pred is None or actual is None: continue if pred == actual and actual >= 4: shaped[i] += R_LEVEL_BONUS_PER_LEVEL * (actual - 3) return [w * r for r in shaped] wrapped.__name__ = raw_fn.__name__ return wrapped