Upload rl/reward.py

rl/reward.py

@@ -1,219 +1,131 @@
 """
-GRPO-compatible reward hook using Impact Oracle as reward source.
-Includes an offline policy comparator for when training is infeasible.
+Reward module - GRPO-compatible reward hook using Impact Oracle.
 """
-
-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
-
+from typing import Any, Dict, List, Optional
 
-@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)
+import sys
+from pathlib import Path
+sys.path.insert(0, str(Path(__file__).parent.parent))
+from oracle.oracle import ImpactOracle
 
 
 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)
+    Converts Impact Oracle scores into RL rewards.
+    Compatible with TRL GRPOTrainer via reward_funcs parameter.
     """
 
     def __init__(
         self,
-        oracle,
-        ledger,
-        broker,
-        mode: str = "code",
-        agent_id: str = "default_agent",
+        oracle: Optional[ImpactOracle] = None,
+        mode: str = "retrieval_qa",
+        compute_budget: float = 10000.0,
+        target_accuracy: float = 0.8,
     ):
-        self.oracle = oracle
-        self.ledger = ledger
-        self.broker = broker
+        self.oracle = oracle or ImpactOracle()
         self.mode = mode
-        self.agent_id = agent_id
-        self._trajectories: List[Trajectory] = []
+        self.compute_budget = compute_budget
+        self.target_accuracy = target_accuracy
+        self.trajectory_history: List[Dict[str, Any]] = []
 
     def compute_rewards(
         self,
         prompts: List[str],
         completions: List[str],
-        oracle_inputs: Optional[List[Dict]] = None,
+        answers: List[Optional[str]],
+        gold_answers: List[str],
+        confidences: List[float],
+        compute_costs: List[float],
+        agent_ids: Optional[List[str]] = 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)
+        Returns list of float rewards (one per completion).
         """
         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()))
+        agent_ids = agent_ids or ["agent_default"] * len(prompts)
 
+        for i in range(len(prompts)):
             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,
-                )
+                action={"abstained": answers[i] is None},
+                context={"gold_answer": gold_answers[i]},
+                result={
+                    "answer": answers[i],
+                    "confidence": confidences[i],
+                    "evidence": kwargs.get("evidences", [{}] * len(prompts))[i],
+                    "compute_cost": compute_costs[i],
+                },
+                agent_id=agent_ids[i],
             )
+            rewards.append(oracle_res.reward_value)
+            self.trajectory_history.append({
+                "prompt": prompts[i][:100],
+                "reward": oracle_res.reward_value,
+                "raw_score": oracle_res.raw_score,
+                "failure_tags": oracle_res.failure_tags,
+            })
 
         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:
+    def compute_reward_single(
+        self,
+        prompt: str,
+        completion: str,
+        answer: Optional[str],
+        gold_answer: str,
+        confidence: float,
+        compute_cost: float,
+        agent_id: str = "agent_default",
+        evidence: Optional[Dict[str, Any]] = None,
+    ) -> float:
+        """Compute reward for a single completion."""
+        oracle_res = self.oracle.score(
+            mode=self.mode,
+            action={"abstained": answer is None},
+            context={"gold_answer": gold_answer},
+            result={
+                "answer": answer,
+                "confidence": confidence,
+                "evidence": evidence or {},
+                "compute_cost": compute_cost,
+            },
+            agent_id=agent_id,
+        )
+        self.trajectory_history.append({
+            "prompt": prompt[:100],
+            "reward": oracle_res.reward_value,
+            "raw_score": oracle_res.raw_score,
+            "failure_tags": oracle_res.failure_tags,
+        })
+        return oracle_res.reward_value
+
+
+class OfflinePolicyComparator:
     """
-    Compare two policies using saved trajectories when online GRPO training
-    is infeasible due to compute constraints.
+    Compare two policies using offline trajectory data.
+    Useful when full GRPO training is not feasible.
     """
 
-    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")
+    def __init__(self, reward_hook: RewardHook):
+        self.reward_hook = reward_hook
 
-        # 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,
-        }
+    def compare(
+        self,
+        policy_a_trajectories: List[Dict[str, Any]],
+        policy_b_trajectories: List[Dict[str, Any]],
+    ) -> Dict[str, Any]:
+        """Compare two policies on same test set."""
+        rewards_a = [t["reward"] for t in policy_a_trajectories]
+        rewards_b = [t["reward"] for t in policy_b_trajectories]
 
-    @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,
+            "mean_reward_a": sum(rewards_a) / len(rewards_a),
+            "mean_reward_b": sum(rewards_b) / len(rewards_b),
+            "win_rate": sum(1 for a, b in zip(rewards_a, rewards_b) if a > b) / len(rewards_a),
+            "improvement": (sum(rewards_a) - sum(rewards_b)) / max(abs(sum(rewards_b)), 1e-6),
+            "policy_a_failures": sum(1 for t in policy_a_trajectories if t.get("failure_tags")),
+            "policy_b_failures": sum(1 for t in policy_b_trajectories if t.get("failure_tags")),
         }
-
-    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")