66 lines
2.4 KiB
Python
66 lines
2.4 KiB
Python
"""Verifiable rewards derived from deterministic code execution."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
from typing import Sequence
|
|
|
|
from .verifier import CodeTask, VerificationResult, VerificationStatus, verify_code
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class CodeReward:
|
|
total: float
|
|
correctness: float
|
|
execution: float
|
|
format: float
|
|
reason: str
|
|
|
|
|
|
def score_verification(result: VerificationResult) -> CodeReward:
|
|
"""Map verifier outcomes to a transparent reward in [-1, 1]."""
|
|
status = result.status
|
|
if status == VerificationStatus.PASSED:
|
|
return CodeReward(1.0, 0.8, 0.1, 0.1, "all tests passed")
|
|
if status == VerificationStatus.WRONG_ANSWER:
|
|
correctness = 0.8 * result.pass_rate
|
|
total = -0.2 + correctness + 0.1 + 0.1
|
|
return CodeReward(total, correctness, 0.1, 0.1, "partial test credit")
|
|
if status == VerificationStatus.RUNTIME_ERROR:
|
|
return CodeReward(-0.4, 0.0, -0.5, 0.1, "runtime error")
|
|
if status == VerificationStatus.TIMEOUT:
|
|
return CodeReward(-0.6, 0.0, -0.7, 0.1, "execution timeout")
|
|
if status == VerificationStatus.SYNTAX_ERROR:
|
|
return CodeReward(-0.8, 0.0, -0.8, 0.0, "syntax error")
|
|
if status in {VerificationStatus.NO_CODE, VerificationStatus.POLICY_VIOLATION}:
|
|
return CodeReward(-1.0, 0.0, 0.0, -1.0, status.value)
|
|
return CodeReward(-1.0, 0.0, -1.0, 0.0, "verifier failure")
|
|
|
|
|
|
def score_code_batch(
|
|
tasks: Sequence[CodeTask],
|
|
responses: Sequence[str],
|
|
*,
|
|
num_generations: int,
|
|
timeout_seconds: float = 2.0,
|
|
memory_mb: int = 256,
|
|
) -> list[CodeReward]:
|
|
"""Score grouped generations in task-major order for GRPO/CISPO."""
|
|
if num_generations < 1:
|
|
raise ValueError("num_generations must be at least 1")
|
|
expected = len(tasks) * num_generations
|
|
if len(responses) != expected:
|
|
raise ValueError(f"expected {expected} responses for {len(tasks)} tasks, got {len(responses)}")
|
|
rewards: list[CodeReward] = []
|
|
for task_index, task in enumerate(tasks):
|
|
start = task_index * num_generations
|
|
for response in responses[start:start + num_generations]:
|
|
result = verify_code(
|
|
response,
|
|
task,
|
|
timeout_seconds=timeout_seconds,
|
|
memory_mb=memory_mb,
|
|
)
|
|
rewards.append(score_verification(result))
|
|
return rewards
|