Upload rl/reward.py

rl/reward.py

@@ -0,0 +1,219 @@
+"""
+GRPO-compatible reward hook using Impact Oracle as reward source.
+Includes an offline policy comparator for when training is infeasible.
+"""
+
+import json
+import math
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Tuple
+
+import numpy as np
+
+
+@dataclass
+class Trajectory:
+    prompt: str = ""
+    completion: str = ""
+    oracle_result: Dict = field(default_factory=dict)
+    reward: float = 0.0
+    compute_cost: float = 0.0
+    mode: str = "code"
+    metadata: Dict = field(default_factory=dict)
+
+
+class RewardHook:
+    """
+    Wraps Impact Oracle + Ledger + Broker into a reward function
+    compatible with TRL GRPOTrainer.
+
+    Usage with GRPOTrainer:
+        reward_fn = RewardHook(oracle, ledger, broker).compute_rewards
+        trainer = GRPOTrainer(..., reward_func=reward_fn)
+    """
+
+    def __init__(
+        self,
+        oracle,
+        ledger,
+        broker,
+        mode: str = "code",
+        agent_id: str = "default_agent",
+    ):
+        self.oracle = oracle
+        self.ledger = ledger
+        self.broker = broker
+        self.mode = mode
+        self.agent_id = agent_id
+        self._trajectories: List[Trajectory] = []
+
+    def compute_rewards(
+        self,
+        prompts: List[str],
+        completions: List[str],
+        oracle_inputs: Optional[List[Dict]] = None,
+        **kwargs,
+    ) -> List[float]:
+        """
+        Compute rewards for a batch of completions.
+
+        Args:
+            prompts: list of prompt strings
+            completions: list of completion strings
+            oracle_inputs: optional list of dicts with keys:
+                {"action": ..., "context": ..., "result": ...}
+
+        Returns:
+            list of float rewards (same length as prompts/completions)
+        """
+        rewards = []
+        oracle_inputs = oracle_inputs or [{} for _ in prompts]
+
+        for prompt, completion, oin in zip(prompts, completions, oracle_inputs):
+            action = oin.get("action", {"text": completion})
+            context = oin.get("context", {})
+            result = oin.get("result", {})
+            result.setdefault("compute_cost", len(completion.split()))
+
+            oracle_res = self.oracle.score(
+                mode=self.mode,
+                action=action,
+                context=context,
+                result=result,
+                agent_id=self.agent_id,
+            )
+
+            reward = oracle_res.reward_value
+            rewards.append(reward)
+
+            # Ledger update
+            self.ledger.earn(
+                agent_id=self.agent_id,
+                task_id=oin.get("task_id", "default_task"),
+                action_id=oin.get("action_id", "default_action"),
+                amount=max(0.0, reward),
+                oracle_score=oracle_res.raw_score,
+                compute_cost=result["compute_cost"],
+                reason=oracle_res.reason,
+                capability_scope=oin.get("capability_scope", "general"),
+                task_scope=oin.get("task_scope", "global"),
+            )
+
+            self._trajectories.append(
+                Trajectory(
+                    prompt=prompt,
+                    completion=completion,
+                    oracle_result={
+                        "raw_score": oracle_res.raw_score,
+                        "cost_adjusted_score": oracle_res.cost_adjusted_score,
+                        "confidence": oracle_res.confidence,
+                        "reason": oracle_res.reason,
+                        "failure_tags": oracle_res.failure_tags,
+                    },
+                    reward=reward,
+                    compute_cost=result["compute_cost"],
+                    mode=self.mode,
+                    metadata=oin,
+                )
+            )
+
+        return rewards
+
+    def get_trajectories(self) -> List[Trajectory]:
+        return self._trajectories
+
+    def save_trajectories(self, path: str):
+        Path(path).parent.mkdir(parents=True, exist_ok=True)
+        with open(path, "w") as f:
+            for t in self._trajectories:
+                d = {
+                    "prompt": t.prompt,
+                    "completion": t.completion,
+                    "reward": t.reward,
+                    "compute_cost": t.compute_cost,
+                    "mode": t.mode,
+                    "metadata": t.metadata,
+                }
+                f.write(json.dumps(d, default=str) + "\n")
+
+
+class OfflineComparator:
+    """
+    Compare two policies using saved trajectories when online GRPO training
+    is infeasible due to compute constraints.
+    """
+
+    def __init__(self, baseline_path: Optional[str] = None):
+        self.baseline_path = baseline_path
+        self.baseline: List[Trajectory] = []
+        if baseline_path and Path(baseline_path).exists():
+            self._load(baseline_path)
+
+    def _load(self, path: str):
+        with open(path, "r") as f:
+            for line in f:
+                d = json.loads(line)
+                self.baseline.append(Trajectory(**d))
+
+    def compare(self, candidate_trajectories: List[Trajectory]) -> Dict:
+        """
+        Return comparative metrics between candidate and baseline.
+        """
+        if not self.baseline:
+            return self._summarize(candidate_trajectories, label="candidate")
+
+        base = self._summarize(self.baseline, label="baseline")
+        cand = self._summarize(candidate_trajectories, label="candidate")
+
+        # Paired comparison on common prompts if available
+        base_by_prompt = {t.prompt: t for t in self.baseline}
+        cand_by_prompt = {t.prompt: t for t in candidate_trajectories}
+        common = set(base_by_prompt.keys()) & set(cand_by_prompt.keys())
+
+        reward_diffs = []
+        cost_diffs = []
+        for p in common:
+            reward_diffs.append(cand_by_prompt[p].reward - base_by_prompt[p].reward)
+            cost_diffs.append(
+                cand_by_prompt[p].compute_cost - base_by_prompt[p].compute_cost
+            )
+
+        return {
+            "baseline": base,
+            "candidate": cand,
+            "common_prompts": len(common),
+            "mean_reward_diff": float(np.mean(reward_diffs)) if reward_diffs else None,
+            "mean_cost_diff": float(np.mean(cost_diffs)) if cost_diffs else None,
+            "reward_p_value": None,  # placeholder for t-test
+            "cost_p_value": None,
+        }
+
+    @staticmethod
+    def _summarize(trajectories: List[Trajectory], label: str) -> Dict:
+        rewards = [t.reward for t in trajectories]
+        costs = [t.compute_cost for t in trajectories]
+        return {
+            "label": label,
+            "n": len(trajectories),
+            "mean_reward": float(np.mean(rewards)) if rewards else 0.0,
+            "std_reward": float(np.std(rewards)) if rewards else 0.0,
+            "mean_cost": float(np.mean(costs)) if costs else 0.0,
+            "std_cost": float(np.std(costs)) if costs else 0.0,
+            "total_cost": float(np.sum(costs)) if costs else 0.0,
+            "success_rate": float(np.mean([r > 0.5 for r in rewards])) if rewards else 0.0,
+        }
+
+    def save_baseline(self, trajectories: List[Trajectory], path: str):
+        Path(path).parent.mkdir(parents=True, exist_ok=True)
+        with open(path, "w") as f:
+            for t in trajectories:
+                d = {
+                    "prompt": t.prompt,
+                    "completion": t.completion,
+                    "reward": t.reward,
+                    "compute_cost": t.compute_cost,
+                    "mode": t.mode,
+                    "metadata": t.metadata,
+                }
+                f.write(json.dumps(d, default=str) + "\n")