6.8 Reward Engineering, Verifiers & Sandboxes¶

r"""
math_verifier.py — robust math answer checker for RL training.

Handles:
  - Fraction strings: "1/2", "3 1/4" (mixed numbers)
  - LaTeX: r"\frac{3}{4}", r"\sqrt{2}", r"2^{10}"
  - Percentages: "50%" -> 0.5 comparison
  - Sets/tuples: "{1, 2}" == "{2, 1}"
  - Symbolic via sympy with numeric fallback
"""

import re
import math
from fractions import Fraction
from typing import Optional, Union
import sympy
from sympy.parsing.latex import parse_latex
from sympy import simplify, N


def match_braces(s: str, open_idx: int) -> Optional[int]:
    """
    Given the index of a '{' in s, return the index of its matching '}'.

    Brace matching (not regex) is mandatory here: competition answers nest
    braces, so a lazy regex over a boxed group stops at the FIRST '}' and
    turns \\boxed{\\frac{1}{2}} into the garbage string '\\frac{1'. That
    single bug is a common cause of a phantom accuracy ceiling in RLVR runs.
    """
    depth = 0
    for i in range(open_idx, len(s)):
        if s[i] == '{':
            depth += 1
        elif s[i] == '}':
            depth -= 1
            if depth == 0:
                return i
    return None  # unbalanced (truncated generation)


def strip_boxed(s: str) -> str:
    """Replace the first \\boxed{...} with its (brace-balanced) contents."""
    m = re.search(r'\\boxed\s*\{', s)
    if not m:
        return s
    open_idx = m.end() - 1
    close_idx = match_braces(s, open_idx)
    if close_idx is None:
        return s
    return s[:m.start()] + s[open_idx + 1:close_idx] + s[close_idx + 1:]


def normalize_latex_number(s: str) -> str:
    """Strip LaTeX boilerplate that doesn't change value."""
    s = s.strip()
    # Remove \boxed{...} (brace-balanced, so \boxed{\frac{1}{2}} survives)
    s = strip_boxed(s)
    # Remove dollar signs
    s = s.replace('$', '').strip()
    # Normalize whitespace
    s = re.sub(r'\s+', ' ', s)
    return s


def try_numeric(s: str) -> Optional[float]:
    """
    Try to evaluate s as a number.
    Returns float or None if not parseable as a pure number.
    """
    s = normalize_latex_number(s)
    # Handle percentages: "50%" → 0.5
    if s.endswith('%'):
        try:
            return float(s[:-1]) / 100.0
        except ValueError:
            pass
    # Handle fractions: "3/4"
    try:
        return float(Fraction(s))
    except (ValueError, ZeroDivisionError):
        pass
    # Plain float
    try:
        return float(s)
    except ValueError:
        pass
    return None


def try_sympy(s: str) -> Optional[sympy.Expr]:
    """Parse a LaTeX string with sympy."""
    try:
        expr = parse_latex(normalize_latex_number(s))
        return expr
    except Exception:
        return None


def math_equivalent(pred: str, gold: str, rtol: float = 1e-6) -> bool:
    """
    Return True iff pred and gold represent the same mathematical value.

    Strategy:
      1. Exact string match after normalization
      2. Both are numeric → compare with relative tolerance
      3. Both parse as sympy expressions → simplify(pred - gold) == 0
      4. Sympy + numeric fallback (N(expr))
    """
    pred = normalize_latex_number(pred)
    gold = normalize_latex_number(gold)

    # 1. Exact match
    if pred == gold:
        return True

    # 2. Numeric
    pn = try_numeric(pred)
    gn = try_numeric(gold)
    if pn is not None and gn is not None:
        if gn == 0:
            return abs(pn) < 1e-9
        return abs(pn - gn) / (abs(gn) + 1e-12) < rtol

    # 3. Symbolic
    pe = try_sympy(pred)
    ge = try_sympy(gold)
    if pe is not None and ge is not None:
        diff = simplify(pe - ge)
        if diff == 0:
            return True
        # Numeric fallback: evaluate to float
        try:
            pf = float(N(pe, 30))
            gf = float(N(ge, 30))
            if gf == 0:
                return abs(pf) < 1e-9
            return abs(pf - gf) / (abs(gf) + 1e-12) < rtol
        except Exception:
            pass

    return False


def extract_answer_from_completion(completion: str) -> str:
    """
    Extract the model's final answer.
    Expects the model to write \\boxed{answer} or <answer>...</answer>.
    Returns empty string if not found.
    """
    # Try \boxed{...} (LaTeX math). Search the LAST \boxed{}: models often
    # restate the format in their reasoning before committing to an answer.
    starts = [m.start() for m in re.finditer(r'\\boxed\s*\{', completion)]
    for start in reversed(starts):
        open_idx = completion.index('{', start)
        close_idx = match_braces(completion, open_idx)
        if close_idx is not None:
            return completion[open_idx + 1:close_idx].strip()
    # Try <answer>...</answer> tag
    m = re.search(r'<answer>(.*?)</answer>', completion, re.DOTALL)
    if m:
        return m.group(1).strip()
    # Fallback: last number on a "The answer is X" line
    m = re.search(r'[Tt]he answer is[:\s]+([^\n.]+)', completion)
    if m:
        return m.group(1).strip()
    return ""


def math_reward(prompt: str, completion: str, gold_answer: str) -> float:
    """
    Main entry point: return reward in {0.0, 0.1, 1.0} for a math response.
    A small partial reward of 0.1 is returned when the format is correct
    but the answer is wrong, to encourage the model to use \boxed{}.
    """
    pred = extract_answer_from_completion(completion)
    if not pred:
        return 0.0  # No answer extracted

    if math_equivalent(pred, gold_answer):
        return 1.0
    else:
        # Small reward for correct format, zero for content
        return 0.1  # format reward (optional — see section on reward shaping)

# pip install math-verify
from math_verify import parse, verify

gold = parse(r"$\frac{1}{2}$")                      # list of candidate readings
pred = parse(r"So the answer is $\boxed{0.5}$.")    # extraction + parsing in one call
print(verify(gold, pred))                           # True: 1/2 == 0.5

"""
code_verifier.py — unit-test-based reward for code generation.

This module orchestrates sandboxed execution (see sandboxes section)
and returns a float reward based on test outcomes.
"""

import re
import subprocess
import sys
import tempfile
import os
import json
import textwrap
from dataclasses import dataclass


@dataclass
class TestResult:
    passed: int         # number of tests that passed
    total: int          # total tests
    timed_out: bool     # did any test time out?
    error: str          # compilation/import error if any


def run_tests_in_subprocess(
    code: str,
    test_suite: str,
    timeout_seconds: float = 5.0,
) -> TestResult:
    """
    Execute generated code + test suite in a temporary directory.
    Returns TestResult.

    IMPORTANT: For production use, wrap this in a proper sandbox
    (see the Sandboxed Execution section). This bare subprocess
    approach is only safe on air-gapped machines or inside a container.
    """
    with tempfile.TemporaryDirectory() as tmpdir:
        # Write generated code
        code_path = os.path.join(tmpdir, "solution.py")
        with open(code_path, "w") as f:
            f.write(code)

        # Write test runner that imports solution.py.
        # NOTE: build the template flush-left and indent the embedded
        # test_suite block explicitly — mixing textwrap.dedent() with a
        # multi-line f-string substitution is fragile: dedent() computes a
        # single common leading-whitespace prefix over the *whole* string,
        # but only the first line of an embedded multi-line value inherits
        # the template's indentation (subsequent lines don't), so the
        # inserted test functions end up misaligned relative to the
        # surrounding `try:` block and raise IndentationError.
        indented_tests = textwrap.indent(test_suite, "    ")
        runner = f'''import sys, json, traceback
sys.path.insert(0, {repr(tmpdir)})

results = {{"passed": 0, "total": 0, "error": ""}}
try:
    from solution import *
{indented_tests}
    # Each test function is test_<name>; discover and run
    test_fns = [v for k, v in globals().items() if k.startswith("test_")]
    results["total"] = len(test_fns)
    for fn in test_fns:
        try:
            fn()
            results["passed"] += 1
        except Exception:
            pass
except Exception as e:
    results["error"] = traceback.format_exc()
    results["total"] = 1  # At least one test failed

print(json.dumps(results))
'''
        runner_path = os.path.join(tmpdir, "runner.py")
        with open(runner_path, "w") as f:
            f.write(runner)

        try:
            proc = subprocess.run(
                # Use sys.executable, not a bare "python": many Linux/CI
                # images only have "python3" on PATH, so a hardcoded
                # "python" raises FileNotFoundError there.
                [sys.executable, runner_path],
                capture_output=True, text=True,
                timeout=timeout_seconds,
            )
            if proc.returncode != 0 and not proc.stdout:
                return TestResult(0, 1, False, proc.stderr[:500])
            data = json.loads(proc.stdout.strip())
            return TestResult(
                passed=data["passed"],
                total=max(data["total"], 1),
                timed_out=False,
                error=data.get("error", ""),
            )
        except subprocess.TimeoutExpired:
            return TestResult(0, 1, timed_out=True, error="timeout")
        except json.JSONDecodeError as e:
            return TestResult(0, 1, False, f"JSON decode error: {e}")


def extract_code_block(completion: str) -> str:
    """
    Pull the code out of a chat completion. Models wrap solutions in markdown
    fences; take the LAST fenced block, since a model often shows a failed
    attempt first and the final block is the one it stands behind.
    Falls back to the raw completion if there is no fence.
    """
    # Build the fence pattern from a variable rather than writing three
    # literal backticks: this file is itself embedded in markdown, and a
    # literal fence inside a fence terminates the block early.
    fence = "`" * 3
    pattern = fence + r"(?:python|py)?\n(.*?)" + fence
    blocks = re.findall(pattern, completion, re.DOTALL)
    return blocks[-1] if blocks else completion


def code_reward(
    prompt: str,
    completion: str,
    test_suite: str,
    partial_credit: bool = True,
) -> float:
    """
    Reward for code generation:
      1.0  if all tests pass
      k/n  if partial_credit and k of n tests pass
      0.0  if syntax error, timeout, or all tests fail
    """
    code = extract_code_block(completion)

    result = run_tests_in_subprocess(code, test_suite)

    if result.timed_out:
        return 0.0
    if result.error and result.passed == 0:
        return 0.0
    if partial_credit:
        return result.passed / result.total
    else:
        return 1.0 if result.passed == result.total else 0.0

"""
sandbox_pool.py — reusable sandbox pool using Docker containers.

In production you would use Firecracker or gVisor microVMs;
Docker containers are shown here for clarity.

Each sandbox is a long-lived container. We send code + tests
over stdin and get results back over stdout, avoiding the per-run
container startup cost (~200ms) by reusing warm containers.
"""

import threading
import queue
import select
import subprocess
import json
import time
from dataclasses import dataclass, field
from typing import Optional


SANDBOX_IMAGE = "python:3.11-slim"
# Dockerfile for the sandbox image would add no extra packages
# and run as a non-root user, but we keep it simple here.

SANDBOX_RUNNER = """
import sys, json, traceback, signal, resource

# Hard memory limit: 256 MB
resource.setrlimit(resource.RLIMIT_AS, (256 * 1024 * 1024, 256 * 1024 * 1024))

# Read tasks from stdin until EOF
for line in sys.stdin:
    task = json.loads(line)
    code = task["code"]
    tests = task["tests"]
    result = {"passed": 0, "total": 0, "error": ""}

    try:
        exec_globals = {}
        exec(compile(code, "<generated>", "exec"), exec_globals)
        exec(compile(tests, "<tests>", "exec"), exec_globals)
        fns = [v for k, v in exec_globals.items() if k.startswith("test_")]
        result["total"] = len(fns)
        for fn in fns:
            try:
                fn()
                result["passed"] += 1
            except Exception:
                pass
    except Exception:
        result["error"] = traceback.format_exc()[-500:]
        result["total"] = 1

    print(json.dumps(result), flush=True)
"""


@dataclass
class SandboxWorker:
    proc: subprocess.Popen
    lock: threading.Lock = field(default_factory=threading.Lock)
    last_used: float = field(default_factory=time.time)


class SandboxPool:
    """
    A pool of warm sandbox processes.
    Thread-safe: acquire() blocks until a worker is free.
    """

    def __init__(self, pool_size: int = 8, timeout: float = 5.0):
        self.timeout = timeout
        self._available: queue.Queue = queue.Queue()
        self._all_workers = []
        for _ in range(pool_size):
            w = self._spawn()
            self._all_workers.append(w)
            self._available.put(w)

    def _spawn(self) -> SandboxWorker:
        """Start a sandbox process that reads JSON tasks from stdin."""
        proc = subprocess.Popen(
            [
                "docker", "run", "--rm", "--interactive",
                "--network=none",           # No network
                "--memory=256m",            # 256 MB RAM limit
                "--cpus=1",                 # 1 vCPU
                "--pids-limit=50",          # No fork bombs
                "--read-only",              # Read-only root filesystem
                "--tmpfs=/tmp:size=64m",    # Small writable /tmp
                SANDBOX_IMAGE,
                "python", "-c", SANDBOX_RUNNER,
            ],
            stdin=subprocess.PIPE,
            stdout=subprocess.PIPE,
            stderr=subprocess.DEVNULL,
            text=True,
        )
        return SandboxWorker(proc=proc)

    def execute(self, code: str, tests: str) -> dict:
        """
        Run code + tests in a sandbox. Blocks until a worker is available.
        Returns {"passed": int, "total": int, "error": str}.
        """
        worker: SandboxWorker = self._available.get(timeout=30.0)
        try:
            task = json.dumps({"code": code, "tests": tests}) + "\n"
            worker.proc.stdin.write(task)
            worker.proc.stdin.flush()
            # Read response with a timeout. A subprocess.Popen(text=True)
            # pipe is a TextIOWrapper over a plain OS pipe, not a socket,
            # so it has no `_sock` attribute to call settimeout() on —
            # select() is the correct way to bound the wait on a pipe fd.
            ready, _, _ = select.select([worker.proc.stdout], [], [], self.timeout)
            if not ready:
                raise TimeoutError(f"sandbox worker timed out after {self.timeout}s")
            line = worker.proc.stdout.readline()
            return json.loads(line)
        except Exception as e:
            # Worker is dead (or wedged); kill and replace it, keeping the
            # registry in sync so shutdown can reap every live container.
            worker.proc.kill()
            self._all_workers.remove(worker)
            worker = self._spawn()
            self._all_workers.append(worker)
            return {"passed": 0, "total": 1, "error": str(e)}
        finally:
            worker.last_used = time.time()
            self._available.put(worker)  # Return worker to pool

"""
llm_judge.py — LLM-as-a-judge reward for open-ended tasks.

Speaks the OpenAI chat-completions API, which is also what vLLM, SGLang and
TGI expose. Default here is a *local* judge: `vllm serve <model> --port 8000`
gives you an OpenAI-compatible /v1 endpoint, so an RL run can call a judge
thousands of times per step at zero marginal cost and with no rate limits.
Point JUDGE_BASE_URL at a hosted provider instead if you want a frontier judge.
"""

import os
import re
import asyncio
import openai
from typing import Optional


# One client per process. Constructing an AsyncOpenAI inside the request path
# creates (and leaks) a fresh connection pool on every judged completion.
JUDGE_BASE_URL = os.environ.get("JUDGE_BASE_URL", "http://localhost:8000/v1")
JUDGE_MODEL = os.environ.get("JUDGE_MODEL", "local-judge")  # --served-model-name
_client = openai.AsyncOpenAI(
    base_url=JUDGE_BASE_URL,
    api_key=os.environ.get("OPENAI_API_KEY", "EMPTY"),  # vLLM ignores the key
)


JUDGE_SYSTEM = """
You are an expert evaluator for mathematical reasoning and problem solving.
You will be given a problem, a reference solution, and a model response.
Your task is to evaluate the model response on a scale of 0 to 10.

Evaluation criteria:
- Correctness (5 pts): Is the final answer correct?
- Reasoning quality (3 pts): Is the chain-of-thought valid, with each step justified?
- Clarity (2 pts): Is the solution well-organized and readable?

Output format:
<thinking>
[Your analysis of the response]
</thinking>
<score>N</score>

Where N is an integer from 0 to 10.
""".strip()


JUDGE_USER_TEMPLATE = """
PROBLEM:
{problem}

REFERENCE SOLUTION:
{reference}

MODEL RESPONSE:
{response}

Please evaluate the model response using the criteria above.
""".strip()


async def judge_single(
    problem: str,
    response: str,
    reference: str,
    judge_model: str = JUDGE_MODEL,
    temperature: float = 0.0,
) -> float:
    """
    Returns a normalized reward in [0, 1] for a single (problem, response) pair.
    Calls the judge LLM asynchronously.
    """
    messages = [
        {"role": "system", "content": JUDGE_SYSTEM},
        {"role": "user", "content": JUDGE_USER_TEMPLATE.format(
            problem=problem,
            reference=reference,
            response=response,
        )},
    ]

    resp = await _client.chat.completions.create(
        model=judge_model,
        messages=messages,
        temperature=temperature,
        max_tokens=512,
    )

    text = resp.choices[0].message.content
    m = re.search(r"<score>(\d+(?:\.\d+)?)</score>", text)
    if m:
        raw_score = float(m.group(1))
        return min(max(raw_score / 10.0, 0.0), 1.0)

    # Fallback: try to parse any integer at the end of the response
    nums = re.findall(r"\b(\d+)\b", text)
    if nums:
        raw = float(nums[-1])
        return min(max(raw / 10.0, 0.0), 1.0)

    return 0.5  # Uncertain — return neutral score


async def batch_judge(
    problems: list[str],
    responses: list[str],
    references: list[str],
    judge_model: str = JUDGE_MODEL,
    concurrency: int = 32,
) -> list[float]:
    """
    Evaluate a batch of completions concurrently.
    concurrency limits simultaneous API calls to avoid rate limits.
    """
    sem = asyncio.Semaphore(concurrency)

    async def bounded_judge(p, r, ref):
        async with sem:
            return await judge_single(p, r, ref, judge_model)

    tasks = [
        bounded_judge(p, r, ref)
        for p, r, ref in zip(problems, responses, references)
    ]
    return await asyncio.gather(*tasks)

"""
prm_reward.py — step-level reward using a process reward model.

The PRM is a separate model (often a smaller classifier fine-tuned
on step-level correctness annotations) that scores each reasoning step.
"""

import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification


class ProcessRewardModel:
    """
    Wraps a PRM that scores individual reasoning steps.
    Input: a (problem, partial_solution_so_far) pair.
    Output: a scalar score in [0, 1] for the most recent step.
    """

    def __init__(self, model_name: str, device: str = "cuda"):
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        self.model = AutoModelForSequenceClassification.from_pretrained(
            model_name, num_labels=1,
        ).to(device).eval()
        self.device = device

    @torch.no_grad()
    def score_steps(
        self,
        problem: str,
        steps: list[str],
        batch_size: int = 32,
    ) -> list[float]:
        """
        Score each step in a chain-of-thought.
        steps: list of individual reasoning steps (not cumulative).
        Returns a list of scalar rewards, one per step.
        """
        scores = []
        # Build cumulative context for each step
        contexts = []
        for i in range(len(steps)):
            # Include all steps up to and including step i
            partial = problem + "\n\n" + "\n".join(steps[:i+1])
            contexts.append(partial)

        # Batch inference
        for start in range(0, len(contexts), batch_size):
            batch = contexts[start:start + batch_size]
            enc = self.tokenizer(
                batch,
                return_tensors="pt",
                padding=True,
                truncation=True,
                max_length=2048,
            ).to(self.device)
            logits = self.model(**enc).logits.squeeze(-1)
            # Sigmoid to get probability of "step is correct"
            probs = torch.sigmoid(logits).cpu().tolist()
            if isinstance(probs, float):
                probs = [probs]
            scores.extend(probs)

        return scores

    def aggregate(
        self,
        step_scores: list[float],
        method: str = "min",
    ) -> float:
        """
        Aggregate step scores into a single episode reward.

        'min': reward = min step score (weakest-link; used in Math-Shepherd)
        'last': reward = score of final step (common in practice)
        'mean': reward = mean of all step scores
        """
        if not step_scores:
            return 0.0
        if method == "min":
            return min(step_scores)
        elif method == "last":
            return step_scores[-1]
        elif method == "mean":
            return sum(step_scores) / len(step_scores)
        raise ValueError(f"Unknown aggregation method: {method}")

"""
reward_combiner.py — multi-objective reward combining and normalization.
"""

import numpy as np
from dataclasses import dataclass
from typing import Callable


@dataclass
class RewardComponent:
    name: str
    fn: Callable                  # (prompt, completion, metadata) -> float
    weight: float                 # Relative importance
    clip_min: float = -10.0       # Clip extreme values
    clip_max: float = 10.0


class MultiObjectiveReward:
    """
    Combines multiple reward signals, normalizes per group,
    and returns the final scalar reward for each completion.
    """

    def __init__(self, components: list[RewardComponent]):
        self.components = components

    def compute(
        self,
        prompts: list[str],
        completions: list[str],
        metadata: list[dict],
        normalize_per_group: bool = True,
    ) -> tuple[list[float], dict[str, list[float]]]:
        """
        Returns (total_rewards, per_component_rewards).

        Assumes each (prompt, completions) group has the same prompt
        and completions are ordered: all completions for prompt 0,
        then all for prompt 1, etc.

        metadata: list of dicts with per-sample info (e.g., gold_answer).
        normalize_per_group: apply Z-score normalization within groups.
        """
        n = len(completions)
        component_rewards = {c.name: np.zeros(n) for c in self.components}

        # Evaluate each component
        for comp in self.components:
            for i, (prompt, completion, meta) in enumerate(
                zip(prompts, completions, metadata)
            ):
                raw = comp.fn(prompt, completion, meta)
                clipped = np.clip(raw, comp.clip_min, comp.clip_max)
                component_rewards[comp.name][i] = clipped

        # Weighted sum
        total = np.zeros(n)
        for comp in self.components:
            total += comp.weight * component_rewards[comp.name]

        # Normalize per group if requested
        if normalize_per_group:
            # Group by unique prompt (assumes they come in contiguous blocks)
            unique_prompts = []
            seen = {}
            for p in prompts:
                if p not in seen:
                    seen[p] = len(unique_prompts)
                    unique_prompts.append(p)
            for prompt in unique_prompts:
                idx = [i for i, p in enumerate(prompts) if p == prompt]
                group = total[idx]
                mu, sigma = group.mean(), group.std()
                total[idx] = (group - mu) / (sigma + 1e-8)

        return total.tolist(), {k: v.tolist() for k, v in component_rewards.items()}

from datasets import Dataset
from trl import GRPOConfig, GRPOTrainer
from math_verifier import math_reward, extract_answer_from_completion

def correctness_reward(completions, gold_answer, **kwargs) -> list[float]:
    """One float per completion. `gold_answer` comes from the dataset column."""
    texts = [c[0]["content"] if isinstance(c, list) else c for c in completions]
    return [math_reward("", t, g) for t, g in zip(texts, gold_answer)]

def format_reward(completions, **kwargs) -> list[float]:
    texts = [c[0]["content"] if isinstance(c, list) else c for c in completions]
    return [1.0 if extract_answer_from_completion(t) else 0.0 for t in texts]

ds = Dataset.from_dict({
    "prompt": ["What is 1/2 + 1/2?"],
    "gold_answer": ["1"],
})

cfg = GRPOConfig(
    output_dir="grpo-math",
    num_generations=8,              # group size G for the group-relative baseline
    reward_weights=[1.0, 0.1],      # correctness weight 1.0, format weight lambda_f
    max_completion_length=1024,
)
trainer = GRPOTrainer(
    model="HuggingFaceTB/SmolLM2-135M-Instruct",
    reward_funcs=[correctness_reward, format_reward],
    args=cfg,
    train_dataset=ds,
)
# trainer.train()

# reward.py — veRL-style per-sample scorer
from math_verifier import math_reward
from code_verifier import code_reward

def compute_score(data_source, solution_str, ground_truth, extra_info=None):
    if data_source.startswith("math"):
        return math_reward("", solution_str, ground_truth)
    if data_source.startswith("code"):
        return code_reward("", solution_str, ground_truth)  # ground_truth = tests
    raise ValueError(f"no verifier registered for data_source={data_source}")

"""
reward_server.py — minimal HTTP reward server for RL training.

Deploy behind a load balancer; run multiple instances for throughput.
"""

import asyncio
import hashlib
import json
from contextlib import asynccontextmanager
from typing import Optional

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel

from math_verifier import math_reward
from code_verifier import extract_code_block
from sandbox_pool import SandboxPool
from llm_judge import batch_judge

# Global sandbox pool — initialized once at startup
_sandbox_pool: Optional[SandboxPool] = None
# Simple in-memory cache (use Redis in production)
_cache: dict[str, float] = {}


@asynccontextmanager
async def lifespan(app: FastAPI):
    """
    Startup/shutdown hook. Use lifespan, not @app.on_event("startup"):
    the on_event decorators have been deprecated since FastAPI 0.93.
    Warming the pool here means the first training step does not pay
    container boot latency.
    """
    global _sandbox_pool
    _sandbox_pool = SandboxPool(pool_size=16, timeout=10.0)
    yield
    for w in _sandbox_pool._all_workers:
        w.proc.kill()


app = FastAPI(title="Reward Server", lifespan=lifespan)


class RewardRequest(BaseModel):
    task_type: str          # "math", "code", "judge"
    prompt: str
    completion: str
    metadata: dict          # gold_answer, test_suite, reference, etc.


class RewardResponse(BaseModel):
    reward: float
    cached: bool
    breakdown: dict         # Per-component rewards for logging


def _cache_key(req: RewardRequest) -> str:
    raw = json.dumps({
        "type": req.task_type,
        "prompt": req.prompt,
        "completion": req.completion,
        "meta": req.metadata,
    }, sort_keys=True)
    return hashlib.sha256(raw.encode()).hexdigest()


@app.post("/reward", response_model=RewardResponse)
async def compute_reward(req: RewardRequest):
    key = _cache_key(req)
    if key in _cache:
        return RewardResponse(reward=_cache[key], cached=True, breakdown={})

    breakdown = {}
    total = 0.0

    if req.task_type == "math":
        gold = req.metadata.get("gold_answer", "")
        r = math_reward(req.prompt, req.completion, gold)
        breakdown["correctness"] = r
        total = r

    elif req.task_type == "code":
        tests = req.metadata.get("test_suite", "")
        code = extract_code_block(req.completion)
        # SandboxPool.execute() blocks on a pipe to a container, so hand it to
        # the default thread pool; otherwise one slow test suite freezes the
        # event loop and every other in-flight reward request with it.
        out = await asyncio.get_running_loop().run_in_executor(
            None, lambda: _sandbox_pool.execute(code, tests)
        )
        r = out["passed"] / max(out["total"], 1)
        breakdown["code_tests"] = r
        total = r

    elif req.task_type == "judge":
        reference = req.metadata.get("reference", "")
        rewards = await batch_judge(
            [req.prompt], [req.completion], [reference]
        )
        r = rewards[0]
        breakdown["llm_judge"] = r
        total = r

    else:
        raise HTTPException(400, f"Unknown task_type: {req.task_type}")

    _cache[key] = total
    return RewardResponse(reward=total, cached=False, breakdown=breakdown)


@app.post("/reward/batch")
async def compute_rewards_batch(requests: list[RewardRequest]):
    """Process a batch of reward requests concurrently."""
    tasks = [compute_reward(req) for req in requests]
    return await asyncio.gather(*tasks)

"""
rl_reward_pipeline.py — end-to-end reward pipeline for a math RLVR run.

This orchestrates: rollout → reward computation → advantage normalization.
Intended to run on the CPU reward server, called from the training loop.
"""

import re
import asyncio
import numpy as np
from typing import NamedTuple


class RolloutBatch(NamedTuple):
    prompts: list[str]            # One per prompt group
    completions: list[list[str]]  # completions[i] = list of G completions for prompt i
    gold_answers: list[str]       # Reference answers for verifier


class RewardBatch(NamedTuple):
    rewards: np.ndarray           # Shape (N,) flat — all completions
    advantages: np.ndarray        # Shape (N,) after group normalization
    component_log: dict           # For logging / dashboards


async def compute_rewards_for_batch(
    batch: RolloutBatch,
    format_lambda: float = 0.1,
) -> RewardBatch:
    """
    Main pipeline:
    1. Compute correctness reward (rule-based math verifier)
    2. Compute format reward
    3. Combine with weights
    4. Normalize per group to get advantages
    """
    from math_verifier import math_reward, extract_answer_from_completion

    flat_prompts = []
    flat_completions = []
    flat_golds = []
    group_ids = []  # Which group each completion belongs to

    for g, (prompt, comps, gold) in enumerate(
        zip(batch.prompts, batch.completions, batch.gold_answers)
    ):
        for comp in comps:
            flat_prompts.append(prompt)
            flat_completions.append(comp)
            flat_golds.append(gold)
            group_ids.append(g)

    n = len(flat_completions)
    correctness = np.zeros(n)
    format_r = np.zeros(n)

    for i, (prompt, comp, gold) in enumerate(
        zip(flat_prompts, flat_completions, flat_golds)
    ):
        # Format reward: 1 if \boxed{} is present, 0 otherwise
        has_boxed = bool(re.search(r'\\boxed\s*\{', comp))
        format_r[i] = 1.0 if has_boxed else 0.0
        # Correctness reward
        correctness[i] = math_reward(prompt, comp, gold)

    combined = correctness + format_lambda * format_r

    # Normalize per group → advantages
    advantages = np.zeros(n)
    num_groups = len(batch.prompts)
    for g in range(num_groups):
        idx = [i for i, gid in enumerate(group_ids) if gid == g]
        group = combined[idx]
        mu, sigma = group.mean(), group.std()
        advantages[idx] = (group - mu) / (sigma + 1e-8)

    return RewardBatch(
        rewards=combined,
        advantages=advantages,
        component_log={
            "correctness": correctness.tolist(),
            "format": format_r.tolist(),
        },
    )