code/local_train/reward.py

from __future__ import annotations

import re
from dataclasses import asdict, dataclass

from albedo_eval_service.judge_core import amputated_thinking, reserved_token_leak
from albedo_eval_service.shared.loop_check import loop_verdict
from albedo_eval_service.shared.observation_format import first_bash_block, is_truncated
from albedo_eval_service.simulator.prompt_simulator import COMPLETE_MARKER
from local_eval.live_protocol import is_live_submit
from sanity_service.checks import check_one

from .chain_heuristics import empty_submit_count_from_texts
from .pack import is_edit_command

_DUMMY_RE = re.compile(
    r"your_command_here|"
    r"cat\s+<<'EOF'\s*>\s*newfile\.py|"
    r"sed\s+-i\s+.*\bfilename\.py\b"
)


@dataclass
class RewardBreakdown:
    reward: float
    fatal: bool
    submitted: bool
    edited: bool
    submit_before_edit: bool
    dummy: bool
    reasons: list[str]

    def as_dict(self) -> dict:
        return asdict(self)


def score_texts(
    texts: list[str],
    *,
    submit_command: str = "",
    gold_paths: list[str] | None = None,
    require_edit_before_submit: bool = True,
) -> RewardBreakdown:
    """Measurement-gate shaped reward. Not local proxy_score."""
    reasons: list[str] = []
    if not texts:
        return RewardBreakdown(0.0, True, False, False, False, False, ["empty"])

    commands = [first_bash_block(text) for text in texts if first_bash_block(text)]
    looped = loop_verdict(texts)
    leak = reserved_token_leak("\n".join(texts))
    heuristic = check_one(texts[0])
    truncated = any(is_truncated(text) for text in texts)
    amputated = amputated_thinking("\n".join(texts))

    fatal = False
    if truncated:
        fatal, reasons = True, reasons + ["truncated"]
    if not heuristic.passed:
        fatal, reasons = True, reasons + [heuristic.reason]
    if looped.looped:
        fatal, reasons = True, reasons + list(looped.reasons or ["looped"])
    if leak:
        fatal, reasons = True, reasons + [f"reserved_token_leak:{leak}"]
    if not commands:
        fatal, reasons = True, reasons + ["no bash command"]
    if fatal:
        return RewardBreakdown(0.0, True, False, False, False, False, reasons)

    empty_submits = empty_submit_count_from_texts(texts, _marker(submit_command))
    if empty_submits >= 2:
        return RewardBreakdown(
            0.0, True, True, False, True, False, reasons + ["empty_double_submit"]
        )

    dummy = any(_DUMMY_RE.search(cmd) for cmd in commands)
    edited = any(is_edit_command(cmd) for cmd in commands)
    relevant = _relevant_edit(commands, gold_paths or [])
    submitted = _submitted(texts, submit_command)
    first_submit = next(
        (i for i, text in enumerate(texts) if _submitted([text], submit_command)), None
    )
    first_edit = next((i for i, cmd in enumerate(commands) if is_edit_command(cmd)), None)
    submit_before_edit = bool(
        require_edit_before_submit
        and submitted
        and (first_edit is None or (first_submit is not None and first_submit <= first_edit))
    )

    reward = 0.15
    if relevant:
        reward += 0.40
        reasons.append("edit_on_gold_path")
    elif edited:
        reward += 0.20
        reasons.append("edit")
    if submitted and not submit_before_edit:
        reward += 0.35
        reasons.append("submit")
    if submit_before_edit:
        reward -= 0.25
        reasons.append("submit_before_edit")
    if dummy:
        reward -= 0.35
        reasons.append("dummy_edit")
    if len(set(commands)) >= 2:
        reward += 0.05
        reasons.append("diverse")
    if looped.dup_cmd_ratio < 0.25:
        reward += 0.05
    if amputated:
        reward *= 0.5
        reasons.append("amputated_thinking")
    return RewardBreakdown(
        reward=round(max(0.0, min(1.0, reward)), 4),
        fatal=False,
        submitted=submitted,
        edited=edited,
        submit_before_edit=submit_before_edit,
        dummy=dummy,
        reasons=reasons,
    )


def reward_completions(
    completions: list[str],
    submit_command: list[str] | None = None,
    gold_paths: list[list[str] | str] | None = None,
    **_kwargs,
) -> list[float]:
    """TRL GRPO callback. Extra dataset columns arrive as aligned lists."""
    rewards: list[float] = []
    for index, completion in enumerate(completions):
        command = ""
        if submit_command:
            command = submit_command[index] if index < len(submit_command) else submit_command[0]
        paths = _paths_at(gold_paths, index)
        rewards.append(
            score_texts([completion], submit_command=command, gold_paths=paths).reward
        )
    return rewards


def _submitted(texts: list[str], submit_command: str) -> bool:
    marker = _marker(submit_command)
    return any(is_live_submit(text, command=submit_command, marker=marker) for text in texts)


def _relevant_edit(commands: list[str], gold_paths: list[str]) -> bool:
    if not gold_paths:
        return False
    needles = [path.rstrip("/").split("/")[-1] for path in gold_paths if path]
    for command in commands:
        if not is_edit_command(command):
            continue
        if any(needle and needle in command for needle in needles):
            return True
        if any(path in command for path in gold_paths):
            return True
    return False


def _marker(submit_command: str) -> str:
    if not submit_command:
        return COMPLETE_MARKER
    token = submit_command.split("&&")[0].strip()
    if token.startswith("echo "):
        return token[5:].strip().strip("'\"")
    return token


def _paths_at(gold_paths: list[list[str] | str] | None, index: int) -> list[str]:
    if not gold_paths:
        return []
    value = gold_paths[index] if index < len(gold_paths) else gold_paths[0]
    if isinstance(value, str):
        return [part for part in value.split("\n") if part.strip()]
    return list(value)