Модуль 9. RLVR и GRPO: проверяемые награды, групповые оценки и устойчивость

from fractions import Fraction
import math
import re

NUMBER_PATTERN = re.compile(
    r"-?\d+\s*/\s*-?\d+|-?(?:\d+(?:\.\d*)?|\.\d+)(?:[eE][+-]?\d+)?"
)
BOXED_PATTERN = re.compile(r"\\boxed\{([^{}]*)\}")
THOUSANDS_SEPARATOR = re.compile(r"(?<=\d),(?=\d{3}(?:\D|$))")


def extract_numeric_answer(text: str) -> str | None:
    if not isinstance(text, str):
        raise TypeError("text must be a string")
    boxed = BOXED_PATTERN.findall(text)
    if boxed:
        candidate = boxed[-1].strip()
        return candidate or None
    normalized = THOUSANDS_SEPARATOR.sub("", text)
    matches = list(NUMBER_PATTERN.finditer(normalized))
    return matches[-1].group(0).strip() if matches else None


def parse_numeric(value: str) -> float:
    if not isinstance(value, str):
        raise TypeError("value must be a string")
    value = THOUSANDS_SEPARATOR.sub("", value).replace(" ", "").strip()
    if not value:
        raise ValueError("empty numeric value")
    if "/" in value:
        parts = value.split("/")
        if len(parts) != 2:
            raise ValueError("fraction must contain one slash")
        numerator, denominator = (int(part) for part in parts)
        if denominator == 0:
            raise ZeroDivisionError("fraction denominator must be non-zero")
        result = float(Fraction(numerator, denominator))
    else:
        result = float(value)
    if not math.isfinite(result):
        raise ValueError("numeric value must be finite")
    return result


def verify_numeric_answer(
    response: str,
    gold: str,
    rtol: float = 1e-6,
    atol: float = 1e-9,
) -> int:
    if rtol < 0 or atol < 0 or not math.isfinite(rtol) or not math.isfinite(atol):
        raise ValueError("rtol and atol must be finite and non-negative")
    answer = extract_numeric_answer(response)
    if answer is None:
        return 0
    try:
        observed = parse_numeric(answer)
        expected = parse_numeric(gold)
    except (TypeError, ValueError, ZeroDivisionError, OverflowError):
        return 0
    tolerance = atol + rtol * max(1.0, abs(expected))
    return int(abs(observed - expected) <= tolerance)

cases = [
    (r"\\boxed{42}", "42", 1),
    ("answer: 1/2", "0.5", 1),
    (r"wrong 42, final \\boxed{41}", "42", 0),
    ("1,000", "1000", 1),
    ("1e3", "1000", 1),
    (r"\\boxed{1/0}", "0", 0),
    ("no answer", "1", 0),
    ("42 items minus 1 = 41", "42", 0),
]
for response, gold, expected in cases:
    assert verify_numeric_answer(response, gold) == expected

import math
import numpy as np


def group_relative_advantages(
    rewards,
    scale: bool = True,
    eps: float = 1e-8,
) -> np.ndarray:
    r = np.asarray(rewards, dtype=float)
    if r.ndim != 1 or r.size == 0 or not np.all(np.isfinite(r)):
        raise ValueError("rewards must be a non-empty finite 1D array")
    if eps < 0 or not math.isfinite(eps):
        raise ValueError("eps must be finite and non-negative")
    centered = r - r.mean()
    if not scale:
        return centered
    std = float(r.std(ddof=0))
    if std == 0.0:
        return np.zeros_like(r)
    return centered / (std + eps)

assert np.allclose(
    group_relative_advantages([1, 0, 0, 0]),
    [np.sqrt(3), -1/np.sqrt(3), -1/np.sqrt(3), -1/np.sqrt(3)],
    atol=1e-7,
)
assert np.allclose(
    group_relative_advantages([0.9, 1.0, 0.8, 0.9], scale=False),
    [0.0, 0.1, -0.1, 0.0],
)
assert np.array_equal(group_relative_advantages([1, 1, 1, 1]), np.zeros(4))

import math
import numpy as np


def _paired_vectors(ratios, advantages):
    rho = np.asarray(ratios, dtype=float)
    adv = np.asarray(advantages, dtype=float)
    if rho.ndim != 1 or adv.shape != rho.shape or rho.size == 0:
        raise ValueError("ratios and advantages must be equal non-empty vectors")
    if not np.all(np.isfinite(rho)) or not np.all(np.isfinite(adv)):
        raise ValueError("ratios and advantages must be finite")
    if np.any(rho <= 0):
        raise ValueError("ratios must be strictly positive")
    return rho, adv


def clipped_terms(ratios, advantages, eps_low=0.2, eps_high=None):
    rho, adv = _paired_vectors(ratios, advantages)
    if eps_high is None:
        eps_high = eps_low
    if eps_low < 0 or eps_high < 0 or not math.isfinite(eps_low) or not math.isfinite(eps_high):
        raise ValueError("clip widths must be finite and non-negative")
    clipped = np.clip(rho, 1.0 - eps_low, 1.0 + eps_high)
    return np.minimum(rho * adv, clipped * adv)


def clip_activity(ratios, advantages, eps_low=0.2, eps_high=None):
    rho, adv = _paired_vectors(ratios, advantages)
    if eps_high is None:
        eps_high = eps_low
    if eps_low < 0 or eps_high < 0:
        raise ValueError("clip widths must be non-negative")
    return ((adv > 0) & (rho > 1.0 + eps_high)) | ((adv < 0) & (rho < 1.0 - eps_low))

ratios = np.array([1.50, 1.25, 0.70, 1.10])
advs = np.array([1.00, 1.00, -1.00, 1.00])
sym = clipped_terms(ratios, advs, eps_low=0.2)
asym = clipped_terms(ratios, advs, eps_low=0.2, eps_high=0.28)
assert np.allclose(sym, [1.2, 1.2, -0.8, 1.1])
assert np.allclose(asym, [1.28, 1.25, -0.8, 1.1])
assert np.array_equal(
    clip_activity(ratios, advs, eps_low=0.2, eps_high=0.28),
    [True, False, True, False],
)

import math
import numpy as np


def _rewards_and_lengths(rewards, lengths):
    r = np.asarray(rewards, dtype=float)
    lens = np.asarray(lengths, dtype=int)
    if r.ndim != 1 or r.size == 0 or not np.all(np.isfinite(r)):
        raise ValueError("rewards must be a non-empty finite vector")
    if lens.ndim != 1 or lens.shape != r.shape or np.any(lens <= 0):
        raise ValueError("lengths must be positive and match rewards")
    return r, lens


def _advantages(rewards, *, scale, eps=1e-8):
    centered = rewards - rewards.mean()
    if not scale:
        return centered
    if eps < 0 or not math.isfinite(eps):
        raise ValueError("eps must be finite and non-negative")
    std = float(rewards.std(ddof=0))
    return np.zeros_like(rewards) if std == 0.0 else centered / (std + eps)


def original_grpo_token_weights(rewards, lengths, *, eps=1e-8):
    r, lens = _rewards_and_lengths(rewards, lengths)
    adv = _advantages(r, scale=True, eps=eps)
    group_size = r.size
    return [
        np.full(int(length), a / (group_size * length), dtype=float)
        for a, length in zip(adv, lens)
    ]


def dr_grpo_token_weights(rewards, lengths, *, fixed_length):
    r, lens = _rewards_and_lengths(rewards, lengths)
    if fixed_length <= 0:
        raise ValueError("fixed_length must be positive")
    adv = _advantages(r, scale=False)
    group_size = r.size
    return [
        np.full(int(length), a / (group_size * fixed_length), dtype=float)
        for a, length in zip(adv, lens)
    ]


def dapo_token_weights(rewards, lengths, *, scale_rewards=True):
    r, lens = _rewards_and_lengths(rewards, lengths)
    adv = _advantages(r, scale=scale_rewards)
    total_active_tokens = int(lens.sum())
    return [
        np.full(int(length), a / total_active_tokens, dtype=float)
        for a, length in zip(adv, lens)
    ]


def response_contributions(token_weights):
    if not token_weights:
        raise ValueError("token_weights must be non-empty")
    return np.asarray([np.asarray(w, dtype=float).sum() for w in token_weights])

rewards, lengths = [1.0, 0.0], [10, 1000]
assert np.allclose(
    response_contributions(original_grpo_token_weights(rewards, lengths)),
    [0.5, -0.5],
)
assert np.allclose(
    response_contributions(dr_grpo_token_weights(rewards, lengths, fixed_length=1000)),
    [0.0025, -0.25],
)
assert np.allclose(
    response_contributions(dapo_token_weights(rewards, lengths)),
    [10/1010, -1000/1010],
    atol=1e-8,
)

import math
import numpy as np


def _active_log_ratio(new_logp, old_logp, mask=None):
    new = np.asarray(new_logp, dtype=float)
    old = np.asarray(old_logp, dtype=float)
    if new.ndim != 1 or old.shape != new.shape or new.size == 0:
        raise ValueError("new_logp and old_logp must be equal non-empty vectors")
    if not np.all(np.isfinite(new)) or not np.all(np.isfinite(old)):
        raise ValueError("log-probabilities must be finite")
    active = np.ones_like(new, dtype=bool)
    if mask is not None:
        m = np.asarray(mask)
        if m.shape != new.shape or np.any((m != 0) & (m != 1)):
            raise ValueError("mask must match and contain only 0/1")
        active = m.astype(bool)
    if not active.any():
        raise ValueError("at least one token must be active")
    return (new - old)[active]


def sequence_importance_ratio(new_logp, old_logp, mask=None):
    mean_log_ratio = float(_active_log_ratio(new_logp, old_logp, mask).mean())
    return float(np.exp(mean_log_ratio))


def gspo_surrogate_term(
    new_logp,
    old_logp,
    advantage,
    eps_low=3e-4,
    eps_high=4e-4,
    mask=None,
):
    if not math.isfinite(advantage):
        raise ValueError("advantage must be finite")
    if eps_low < 0 or eps_high < 0:
        raise ValueError("clip widths must be non-negative")
    ratio = sequence_importance_ratio(new_logp, old_logp, mask)
    clipped = float(np.clip(ratio, 1.0 - eps_low, 1.0 + eps_high))
    return float(min(ratio * advantage, clipped * advantage))

old = np.log([0.20, 0.40, 0.50, 0.30])
new = old + np.log([1.01, 0.99, 1.02, 10.0])
ratio = sequence_importance_ratio(new, old, mask=[1, 1, 1, 0])
expected = (1.01 * 0.99 * 1.02) ** (1 / 3)
assert np.isclose(ratio, expected)

# Длинная последовательность остаётся численно устойчивой: суммируются логи.
long_old = np.full(100_000, -20.0)
long_new = long_old + 1e-4
assert np.isclose(sequence_importance_ratio(long_new, long_old), np.exp(1e-4))

from dataclasses import dataclass
import math


@dataclass(frozen=True)
class RolloutRecord:
    prompt_id: str
    group_id: str
    token_ids: tuple[int, ...]
    action_mask: tuple[int, ...]
    old_logprobs: tuple[float, ...]
    reward_components: tuple[float, ...]
    stop_reason: str
    generator_version: str


def validate_rollout_record(record: RolloutRecord) -> None:
    if not record.prompt_id or not record.group_id or not record.generator_version:
        raise ValueError("prompt_id, group_id, and generator_version must be non-empty")
    n = len(record.token_ids)
    if n == 0 or len(record.action_mask) != n or len(record.old_logprobs) != n:
        raise ValueError("token_ids, action_mask, and old_logprobs must align")
    if any(m not in (0, 1) for m in record.action_mask) or not any(record.action_mask):
        raise ValueError("action_mask must contain 0/1 and at least one active action")
    if not all(math.isfinite(x) for x in record.old_logprobs):
        raise ValueError("old_logprobs must be finite")
    if not record.reward_components or not all(math.isfinite(x) for x in record.reward_components):
        raise ValueError("reward_components must be a non-empty finite sequence")
    if record.stop_reason not in {"eos", "stop_string", "truncated", "environment", "error"}:
        raise ValueError("unsupported stop_reason")

valid = RolloutRecord(
    prompt_id="p-17",
    group_id="g-4",
    token_ids=(101, 102, 103),
    action_mask=(0, 1, 1),
    old_logprobs=(-2.0, -0.4, -0.8),
    reward_components=(1.0, 0.1),
    stop_reason="eos",
    generator_version="checkpoint-1200",
)
validate_rollout_record(valid)

for invalid in (
    RolloutRecord("p", "g", (1, 2), (0, 0), (-1.0, -1.0), (1.0,), "eos", "v1"),
    RolloutRecord("p", "g", (1, 2), (1,), (-1.0, -1.0), (1.0,), "eos", "v1"),
    RolloutRecord("p", "g", (1,), (1,), (float("nan"),), (1.0,), "eos", "v1"),
):
    try:
        validate_rollout_record(invalid)
        raise AssertionError("invalid record was accepted")
    except ValueError:
        pass