""" Per-completion proxy reward for online GRPO. Why a proxy reward? ------------------- TRL's :class:`trl.GRPOTrainer` (verified against the installed TRL 1.2.0 source — see ``trl/trainer/grpo_trainer.py``) calls ``unwrapped_model.generate(...)`` itself and then forwards those *fresh* completions to ``reward_func(prompts, completions, ...)``. Any reward function that tries to look up the completion in a pre-collected ``{(prompt, completion) -> reward}`` table will *miss on every single step* — the lookup key (the freshly generated completion) is never the same string as the recorded one. That's the bug the original notebook had: ``make_reward_fn`` returned a constant ``-0.01`` for almost every step → zero advantage signal → no learning. This module replaces it with a **verifiable, per-completion** reward that scores any (prompt, completion) pair WITHOUT touching the FraudArena server, by: 1. **Format / schema validity** — does the completion parse as JSON and validate as :class:`AdReviewAction`? Currently the dominant source of fallback at training time, so worth ~50% of the budget. 2. **Coherence** — when the action references an ``ad_id`` / ``linked_ad_id``, that ad must actually appear in the prompt's pending list / current-focus block. 3. **Action-class matches recorded gold** — when the dataset row carries a recorded gold-action class (verdict / investigate / link_accounts), generations of the right class get a small bonus. 4. **Decision matches recorded gold** — when the row's ``terminal_grader_score`` says the recorded episode succeeded, verdicts / investigation_targets / linked_ad_ids that match the recorded ones get a larger bonus (it was a high-quality demo). When the recorded episode FAILED, copying the recorded action is mildly *penalised* (don't imitate failure modes). This is consistent with how TRL's verifiable-reward GRPO recipes look (e.g. ``open-r1`` / ``oat-math``): a fast, deterministic scorer that captures schema correctness + a small task-specific signal. """ from __future__ import annotations import json import re from typing import Any, Dict, List, Optional from pydantic import ValidationError from counterfeint.models import AdReviewAction _JSON_FENCE_RE = re.compile(r"```(?:json)?\s*\n(.*?)```", re.DOTALL) def _extract_json_text(raw: str) -> str: text = (raw or "").strip() m = _JSON_FENCE_RE.search(text) if m: return m.group(1).strip() if text.startswith("```"): lines = [l for l in text.split("\n") if not l.strip().startswith("```")] return "\n".join(lines).strip() return text def _parse_completion(completion: str) -> Optional[AdReviewAction]: text = _extract_json_text(completion) if not text: return None try: data = json.loads(text) except json.JSONDecodeError: return None if not isinstance(data, dict): return None try: return AdReviewAction.model_validate(data) except ValidationError: return None def _action_class(action_type: str) -> str: return "verdict" if action_type in {"verdict", "link_accounts"} else "investigate" # Lightweight {key: value} extraction from the recorded action_repr # string the rollout collector stores in metadata. We only need a # handful of fields and avoid eval()/AST on untrusted strings. _REPR_FIELD_RE = re.compile(r"(\w+)=(?:'([^']*)'|([^,)\s]+))") def _gold_fields_from_metadata(meta: Dict[str, Any]) -> Dict[str, Optional[str]]: """Best-effort extraction of (action_type, ad_id, verdict, target, linked) from the dataset row's recorded metadata.""" out: Dict[str, Optional[str]] = { "action_type": None, "ad_id": None, "verdict": None, "investigation_target": None, "linked_ad_id": None, } repr_str = meta.get("action_repr") if not isinstance(repr_str, str): return out for match in _REPR_FIELD_RE.finditer(repr_str): k = match.group(1) v = match.group(2) if match.group(2) is not None else match.group(3) if k in out: out[k] = v return out def _coherent_with_prompt(text: str, prompt: str) -> bool: """Soft check: the referenced ad_id appears verbatim in the prompt.""" return bool(text) and text in prompt def proxy_reward_one( prompt: str, completion: str, *, gold: Dict[str, Optional[str]], gold_episode_score: float, ) -> float: """Score a single (prompt, completion) pair on the [-0.5, 2.5] range. GRPO needs reward *variance* within each generation group to compute non-zero advantages. This function uses a mix of binary gates AND continuous components so that similar-but-not-identical completions get meaningfully different scores. """ action = _parse_completion(completion) if action is None: raw = _extract_json_text(completion) if not raw: return -0.5 # Partial credit: the model tried to produce JSON but it didn't # validate. Give a small continuous score based on how "close" # it was - this creates gradient between "total garbage" and # "almost valid JSON". partial = -0.3 if raw.startswith("{"): partial += 0.05 if "action_type" in raw: partial += 0.05 if "ad_id" in raw: partial += 0.05 if raw.rstrip().endswith("}"): partial += 0.05 return partial reward = 0.0 # 1. Schema validity — binary gate. reward += 0.6 # 2. Coherence — the action references real IDs the prompt mentions. if action.ad_id and _coherent_with_prompt(action.ad_id, prompt): reward += 0.15 if action.linked_ad_id and _coherent_with_prompt(action.linked_ad_id, prompt): reward += 0.15 # 3. Action-class matches the recorded gold class. gold_at = gold.get("action_type") if gold_at and _action_class(action.action_type) == _action_class(gold_at): reward += 0.2 # 4. Decision matches recorded gold, scaled by episode quality. quality = max(0.0, min(1.0, gold_episode_score)) if quality > 0.0: if action.action_type == "verdict" and gold.get("verdict") == action.verdict: reward += 0.6 * quality if ( action.action_type == "investigate" and gold.get("investigation_target") == action.investigation_target ): reward += 0.5 * quality if ( action.action_type == "link_accounts" and gold.get("linked_ad_id") == action.linked_ad_id ): reward += 0.6 * quality # ---- CONTINUOUS components (break ties among valid completions) ---- # 5. Confidence value — continuous [0, 0.15]. Rewards higher # confidence on verdicts (the grader rewards decisive agents). if action.action_type == "verdict" and action.confidence is not None: reward += 0.15 * float(action.confidence) # 6. Rationale evidence density — count how many tokens from the # prompt's findings block appear in the rationale. More evidence # citations = better rationale = higher reward. Continuous. if action.rationale and action.action_type in ("verdict", "link_accounts"): rat_lower = action.rationale.lower() evidence_hits = 0 for marker in ("pmt_", "reg_", "fsdp-", "similarity", "%", ".com", ".net", ".org"): if marker in rat_lower: evidence_hits += 1 reward += min(0.2, evidence_hits * 0.04) # 7. Conciseness bonus — shorter valid completions are better (less # wasted tokens, less chance of trailing garbage). Continuous. comp_len = len(completion.strip()) if comp_len < 150: reward += 0.1 elif comp_len < 300: reward += 0.05 else: reward -= 0.05 # 8. Deterministic hash tiebreaker — last-resort variance injection. # Maps completion text to [0, 0.02] so no two identical-scoring # completions produce exactly the same reward. import hashlib h = int(hashlib.md5(completion.encode()).hexdigest()[:8], 16) reward += 0.02 * (h / 0xFFFFFFFF) return reward def make_proxy_reward_fn( *, gold_lookup: Dict[str, Dict[str, Any]], ): """Build a TRL-compatible reward function. ``gold_lookup`` maps each ``prompt`` string in the dataset to its gold metadata + ``terminal_grader_score`` (constructed once at dataset-build time; see :func:`build_gold_lookup`). """ def _extract_user_text(prompt: Any) -> str: """Extract the raw user prompt text for gold_lookup key. TRL passes chat-formatted prompts as lists of dicts ``[{role: system, ...}, {role: user, content: ...}]``, but our gold_lookup is keyed by the raw user content string. """ if isinstance(prompt, list): for msg in prompt: if isinstance(msg, dict) and msg.get("role") == "user": return msg.get("content", "") return str(prompt) return prompt def _to_str(val: Any) -> str: if isinstance(val, str): return val if isinstance(val, list): return " ".join(str(x) for x in val) return str(val) def reward_fn(prompts, completions, **_: Any) -> List[float]: out: List[float] = [] for prompt, completion in zip(prompts, completions): completion = _to_str(completion) prompt_key = _extract_user_text(prompt) prompt_text = _to_str(prompt_key) gold = gold_lookup.get(prompt_key) if gold is None: # Prompt the trainer batched but we never recorded — # only score schema validity + coherence. out.append( proxy_reward_one( prompt_text, completion, gold={"action_type": None, "ad_id": None, "verdict": None, "investigation_target": None, "linked_ad_id": None}, gold_episode_score=0.0, ) ) continue out.append( proxy_reward_one( prompt_text, completion, gold=gold["fields"], gold_episode_score=float(gold["episode_score"]), ) ) return out return reward_fn def build_gold_lookup(samples: List[Any]) -> Dict[str, Dict[str, Any]]: """Build the prompt → gold map from a list of :class:`counterfeint.training.rollout.InvestigatorTrainingSample`. Most-recent recording wins on duplicate prompts (rare; happens only if the same observation is reached twice in different episodes). """ out: Dict[str, Dict[str, Any]] = {} for s in samples: out[s.prompt] = { "fields": _gold_fields_from_metadata(s.metadata), "episode_score": float(s.terminal_grader_score), } return out __all__ = [ "build_gold_lookup", "make_proxy_reward_fn", "proxy_reward_one", ]