Grpo Rlvr Training

import re

def format_reward(completions, **kwargs) -> list[float]:
    """1.0 if the completion has a <reasoning>...</reasoning>
    block followed by an <answer>...</answer> block, else 0.0.
    This is a gate, not the correctness signal — a
    well-formed wrong answer still scores 0 on
    correctness_reward below.
    """
    pattern = re.compile(
        r"^<reasoning>.*?</reasoning>\s*<answer>.*?</answer>$",
        re.DOTALL,
    )
    return [1.0 if pattern.match(c.strip()) else 0.0 for c in completions]

def correctness_reward(completions, answer, **kwargs) -> list[float]:
    """`answer` is the ground-truth column from the
    training dataset, aligned index-for-index with
    `completions`. Extracts the <answer> block from
    format_reward's expected shape and compares.
    """
    rewards = []
    for completion, gold in zip(completions, answer):
        match = re.search(r"<answer>(.*?)</answer>", completion, re.DOTALL)
        predicted = match.group(1).strip() if match else None
        rewards.append(2.0 if predicted == gold.strip() else 0.0)
    return rewards

import json
from jsonschema import validate, ValidationError

def schema_reward(completions, output_schema, **kwargs) -> list[float]:
    """`output_schema` is a JSON Schema dict, either a
    single constant schema for the whole batch or a
    per-example list the same length as `completions`.
    Rewards valid, schema-conformant JSON; 0.0 for
    anything that doesn't parse or doesn't validate.
    """
    if isinstance(output_schema, dict):
        # Constant case: one schema dict for every completion —
        # zip()-ing a bare dict would iterate its keys instead,
        # not the schema itself, so normalize first.
        schemas = [output_schema] * len(completions)
    else:
        schemas = list(output_schema)
        if len(schemas) != len(completions):
            raise ValueError(
                f"schema_reward: {len(schemas)} schemas for "
                f"{len(completions)} completions"
            )
    rewards = []
    for completion, schema in zip(completions, schemas):
        try:
            parsed = json.loads(completion)
            validate(instance=parsed, schema=schema)
            rewards.append(1.0)
        except (json.JSONDecodeError, ValidationError):
            rewards.append(0.0)
    return rewards

import logging
import subprocess
import tempfile
from pathlib import Path

logger = logging.getLogger(__name__)

def test_execution_reward(
    completions, test_code, sandbox_cmd, timeout_s=10, **kwargs
) -> list[float]:
    """`test_code` is a per-example pytest snippet that
    imports the candidate under a fixed module name
    and asserts expected behavior. Runs each candidate
    in its own subprocess with a wall-clock timeout;
    an infinite loop or crash scores 0.0 instead of
    hanging the training loop.

    SECURITY: executes model-generated code. This
    function REQUIRES an isolation boundary — it does
    not run anything on the host by itself.

    `sandbox_cmd` (list[str], required) is a command
    prefix that wraps pytest in that boundary, e.g. a
    network-disabled, resource-capped Docker container:

        # sandbox_cmd = [
        #     "docker", "run", "--rm", "--network=none",
        #     "--memory=1g", "--cpus=1",
        #     "-v", f"{workdir}:/work:ro", "-w", "/work",
        #     "python:3.12-slim",
        # ]

    If `sandbox_cmd` is falsy, this function refuses to
    execute anything and returns 0.0 for every
    completion — it never falls back to running
    pytest on the host. The subprocess environment is
    scrubbed to a minimal PATH (no HF tokens,
    experiment-tracker keys, cloud credentials, or SSH
    keys). The timeout is a liveness guard for the
    training loop, NOT a security boundary — isolation
    comes entirely from `sandbox_cmd`; the temporary
    directory only confines harness writes, not what
    executed code can read or reach.
    """
    if not sandbox_cmd:
        logger.warning(
            "test_execution_reward: no sandbox boundary provided "
            "— refusing to execute model-generated code"
        )
        return [0.0 for _ in completions]

    scrubbed_env = {"PATH": "/usr/bin:/bin"}
    rewards = []
    for completion, tests in zip(completions, test_code):
        with tempfile.TemporaryDirectory() as tmp:
            candidate_path = Path(tmp) / "candidate.py"
            test_path = Path(tmp) / "test_candidate.py"
            candidate_path.write_text(completion)
            test_path.write_text(tests)
            try:
                result = subprocess.run(
                    [*sandbox_cmd, "python", "-m", "pytest",
                     str(test_path), "-q"],
                    cwd=tmp,
                    capture_output=True,
                    timeout=timeout_s,
                    env=scrubbed_env,
                )
                rewards.append(1.0 if result.returncode == 0 else 0.0)
            except subprocess.TimeoutExpired:
                rewards.append(0.0)
    return rewards

def with_length_penalty(reward_fn, target_len=512, penalty_per_token=0.001):
    """Returns a new reward function that subtracts a
    small per-token penalty for every token past
    `target_len`, applied on top of `reward_fn`'s
    output. Penalty is capped so it can't drive an
    otherwise-correct reward negative — it discourages
    padding without overriding correctness.
    """
    def wrapped(completions, **kwargs) -> list[float]:
        base_rewards = reward_fn(completions, **kwargs)
        adjusted = []
        for reward, completion in zip(base_rewards, completions):
            overflow = max(0, len(completion.split()) - target_len)
            penalty = min(reward, overflow * penalty_per_token)
            adjusted.append(reward - penalty)
        return adjusted
    return wrapped

def make_rubric_judge_reward(judge_client, rubric):
    """`judge_client` calls JUDGE_MODEL — a model from a
    *different* model family than the model under
    training, never the model being trained or a
    same-family relative of it. `rubric` is a fixed
    pass/fail criterion string, not a free-form
    quality prompt, so both are bound here rather than
    read from TRL-supplied per-example kwargs. Returns a
    reward function matching TRL's actual signature.
    """
    def rubric_judge_reward(completions, prompts, **kwargs) -> list[float]:
        """Returns 1.0/0.0 per completion, never an
        intermediate score."""
        rewards = []
        for prompt, completion in zip(prompts, completions):
            verdict = judge_client.judge(
                rubric=rubric,
                prompt=prompt,
                response=completion,
                output_format="pass_fail",   # binary only — no Likert scale
            )
            rewards.append(1.0 if verdict == "pass" else 0.0)
        return rewards
    return rubric_judge_reward