| """GRPO reward shaping and credit assignment, as pure functions over per-episode numbers. |
| |
| Settings follow the reference profile MiMo released with MiMo-V2.6 (recipes/arvo/REFERENCE_PENALTIES.json |
| in XiaomiMiMo/verl): |
| * reward: 1 correct and grounded, 0.5 correct but the answer never appeared in a tool result, 0 else; |
| * in-group length penalty on passing rollouts only, and only when more than half the group passes: |
| excess = max over (turns, tool-input tokens, generated tokens) of value / p30-of-passes - 1, |
| penalty = 0.2 * min(excess, 1) ** 1.5; |
| * advantage = shaped reward - group mean (no std normalization); groups whose shaped rewards are |
| all equal carry no signal and are dropped (dynamic sampling); |
| * segment-level penalty for bad tool-call turns (malformed, unknown tool, bad arguments, or a |
| repeated call): in positive episodes those tokens get no credit, in negative ones double blame, |
| and a batch-wide rescale of the clean tokens keeps each sign's total advantage mass unchanged. |
| """ |
| from __future__ import annotations |
|
|
| from dataclasses import dataclass |
|
|
| import numpy as np |
|
|
|
|
| @dataclass(frozen=True) |
| class LengthPenalty: |
| max_penalty: float = 0.2 |
| threshold: float = 0.0 |
| saturate: float = 1.0 |
| exponent: float = 1.5 |
| pass_threshold: float = 0.5 |
| anchor_quantile: float = 0.3 |
| min_pass_rate: float = 0.5 |
|
|
|
|
| def base_reward(correct: bool, grounded: bool, ungrounded: float = 0.5, invented: bool = False, |
| invented_penalty: float = 0.0) -> float: |
| """1 = correct and read from a tool result, `ungrounded` = correct but never seen, 0 = wrong, |
| -invented_penalty = wrong and seen nowhere (made up). So within a group, admitting NOT_FOUND |
| beats inventing an answer, while a correct answer still beats both.""" |
| if correct: |
| return 1.0 if grounded else ungrounded |
| return -invented_penalty if invented else 0.0 |
|
|
|
|
| def length_deltas(rewards: list[float], signals: list[dict], cfg: LengthPenalty = LengthPenalty()) -> list[float]: |
| """Non-positive reward deltas for one group. signals: per episode {metric: value}.""" |
| n = len(rewards) |
| out = [0.0] * n |
| passed = [i for i in range(n) if rewards[i] >= cfg.pass_threshold] |
| if not passed or len(passed) / n <= cfg.min_pass_rate: |
| return out |
| metrics = signals[0].keys() |
| anchor = {m: float(np.percentile([signals[i][m] for i in passed], cfg.anchor_quantile * 100)) for m in metrics} |
| for i in passed: |
| ex = [max(0.0, signals[i][m] / anchor[m] - 1.0) for m in metrics if anchor[m] > 0] |
| e = max(ex, default=0.0) |
| if e <= cfg.threshold: |
| continue |
| t = min((e - cfg.threshold) / (cfg.saturate - cfg.threshold), 1.0) |
| out[i] = -cfg.max_penalty * t ** cfg.exponent |
| return out |
|
|
|
|
| def group_advantages(rewards: list[float]) -> list[float]: |
| mu = sum(rewards) / len(rewards) |
| return [r - mu for r in rewards] |
|
|
|
|
| def has_signal(rewards: list[float], eps: float = 1e-9) -> bool: |
| return max(rewards) - min(rewards) > eps |
|
|
|
|
| def signed_rebalance(adv: list[float], n_flag: list[int], n_clean: list[int], kappa: float = 2.0, |
| min_scale: float = 0.5, max_scale: float = 2.0) -> list[tuple[float, float]]: |
| """Per episode (weight on clean tokens, weight on flagged tokens); token advantage = A * weight. |
| Batch-wide: positive episodes drop credit on flagged tokens and scale clean ones by alpha; |
| negative episodes put kappa x blame on flagged tokens and scale clean ones by beta, so that |
| (unless clipped) the total positive and negative advantage mass is conserved.""" |
| hp = sum(a * f for a, f in zip(adv, n_flag) if a > 0) |
| cp = sum(a * c for a, c in zip(adv, n_clean) if a > 0) |
| hn = sum(-a * f for a, f in zip(adv, n_flag) if a < 0) |
| cn = sum(-a * c for a, c in zip(adv, n_clean) if a < 0) |
| alpha = min(max_scale, 1 + hp / cp) if cp > 0 else 1.0 |
| beta = max(min_scale, 1 - (kappa - 1) * hn / cn) if cn > 0 else 1.0 |
| return [(alpha, 0.0) if a > 0 else (beta, kappa) if a < 0 else (0.0, 0.0) for a in adv] |
|
|
