Upload grpo_hook.py

grpo_hook.py

@@ -1,113 +1,124 @@
 """
-Minimal GRPO-compatible reward hook demonstration.
+GRPO-compatible reward hook for TRL.
 
-If full GRPO training is feasible, use this with TRL GRPOTrainer.
-If not, use OfflineComparator for policy evaluation.
+This module provides a reward function factory that wraps the OCC
+ImpactOracle into a TRL GRPOTrainer-compatible callable.
+
+Usage with TRL::
+
+    from grpo_hook import make_occ_reward_func
+    from trl import GRPOTrainer
+
+    reward_fn = make_occ_reward_func(mode="code", compute_budget=1e5)
+    trainer = GRPOTrainer(
+        model="Qwen/Qwen2.5-0.5B-Instruct",
+        reward_funcs=reward_fn,
+        train_dataset=ds,   # must have a "prompt" column
+    )
+
+The reward function signature follows TRL conventions:
+    def reward_fn(completions, **kwargs) -> list[float]
 """
 
 import json
 from pathlib import Path
 from typing import Dict, List, Optional
 
-import numpy as np
-
 from oracle.oracle import ImpactOracle
 from ledger.ledger import CreditLedger
 from broker.broker import ResourceBroker
-from rl.reward import RewardHook, OfflineComparator
+from rl.reward import RewardHook, OfflinePolicyComparator
 
 
-def demo_grpo_hook():
+def make_occ_reward_func(
+    mode: str = "retrieval_qa",
+    compute_budget: float = 1e5,
+    qa_weights: Optional[Dict] = None,
+    code_weights: Optional[Dict] = None,
+    debate_weights: Optional[Dict] = None,
+) -> callable:
     """
-    Demonstrate the reward hook with synthetic completions.
-    This is a toy loop showing how GRPO reward computation would work.
+    Factory for a TRL-compatible reward function.
+
+    Returns a function with signature (completions, **kwargs) -> list[float].
     """
-    oracle = ImpactOracle(compute_budget=1e5)
-    ledger = CreditLedger(decay_lambda=0.05)
-    broker = ResourceBroker()
-    hook = RewardHook(oracle, ledger, broker, mode="code", agent_id="demo_agent")
-
-    # Simulate a group of completions (as in GRPO)
-    prompts = [
-        "def add(a, b):\n    return",
-        "def add(a, b):\n    return",
-        "def add(a, b):\n    return",
-    ]
-    completions = [
-        "a + b",
-        "a * b",
-        "a + b + 0",
+    oracle = ImpactOracle(
+        compute_budget=compute_budget,
+        qa_weights=qa_weights,
+        code_weights=code_weights,
+        debate_weights=debate_weights,
+    )
+    hook = RewardHook(oracle=oracle, mode=mode)
+
+    def _reward_fn(completions, **kwargs):
+        """
+        TRL calls this with completions as list[str] (standard format)
+        or list[list[dict]] (conversational format).
+        We extract text and look for answer tags.
+        """
+        texts = []
+        for comp in completions:
+            if isinstance(comp, list) and len(comp) > 0 and isinstance(comp[0], dict):
+                # Conversational format: [{"role":"assistant","content":"..."}]
+                texts.append(comp[0].get("content", ""))
+            elif isinstance(comp, str):
+                texts.append(comp)
+            else:
+                texts.append(str(comp))
+
+        answers = []
+        confidences = []
+        compute_costs = []
+
+        for txt in texts:
+            if "<answer>" in txt and "</answer>" in txt:
+                start = txt.find("<answer>") + len("<answer>")
+                end = txt.find("</answer>")
+                ans = txt[start:end].strip()
+            else:
+                # Fallback: last token or empty
+                parts = txt.strip().split()
+                ans = parts[-1] if parts else ""
+            answers.append(ans)
+            confidences.append(0.7 if len(ans) > 0 else 0.3)
+            compute_costs.append(len(txt.split()))
+
+        gold_answers = kwargs.get("answers", [""] * len(texts))
+        if not gold_answers:
+            gold_answers = [""] * len(texts)
+
+        rewards = hook.compute_rewards(
+            prompts=kwargs.get("prompts", [""] * len(texts)),
+            completions=texts,
+            answers=answers,
+            gold_answers=gold_answers,
+            confidences=confidences,
+            compute_costs=compute_costs,
+            agent_ids=kwargs.get("agent_ids", None),
+        )
+        return rewards
+
+    return _reward_fn
+
+
+def demo_offline():
+    """Offline comparison of two policies using the reward hook."""
+    hook = RewardHook(oracle=ImpactOracle(compute_budget=1e5), mode="retrieval_qa")
+    comparator = OfflinePolicyComparator(reward_hook=hook)
+
+    policy_a = [
+        {"reward": 0.5 + i * 0.02, "failure_tags": []}
+        for i in range(10)
     ]
-    oracle_inputs = [
-        {
-            "action": {"text": c},
-            "context": {"previous_passed": False},
-            "result": {"passed": True, "hidden_passed": True, "compute_cost": 5.0},
-            "task_id": "task_1",
-            "action_id": f"comp_{i}",
-        }
-        for i, c in enumerate(completions)
+    policy_b = [
+        {"reward": 0.4 + i * 0.01, "failure_tags": []}
+        for i in range(10)
     ]
-    # Fix the wrong one
-    oracle_inputs[1]["result"]["passed"] = False
-    oracle_inputs[1]["result"]["hidden_passed"] = False
 
-    rewards = hook.compute_rewards(prompts, completions, oracle_inputs)
-    print("GRPO Hook Demo")
-    print("Prompts:", prompts)
-    print("Completions:", completions)
-    print("Rewards:", rewards)
-
-    # Save trajectories for offline comparison
-    hook.save_trajectories("/app/occ/reports/demo_trajectories.jsonl")
-    print("Saved trajectories to reports/demo_trajectories.jsonl")
-
-    return hook
-
-
-def demo_offline_comparison():
-    """
-    Compare two policies using offline trajectory comparison.
-    """
-    # Create baseline policy trajectories
-    baseline_trajs = []
-    for i in range(10):
-        t = type("T", (), {
-            "prompt": f"prompt_{i}",
-            "completion": f"baseline_completion_{i}",
-            "reward": 0.5 + np.random.rand() * 0.3,
-            "compute_cost": 100.0,
-            "mode": "code",
-            "metadata": {},
-        })()
-        baseline_trajs.append(t)
-
-    # Create candidate policy trajectories
-    candidate_trajs = []
-    for i in range(10):
-        t = type("T", (), {
-            "prompt": f"prompt_{i}",
-            "completion": f"candidate_completion_{i}",
-            "reward": 0.6 + np.random.rand() * 0.3,
-            "compute_cost": 70.0,
-            "mode": "code",
-            "metadata": {},
-        })()
-        candidate_trajs.append(t)
-
-    comparator = OfflineComparator()
-    comparator.save_baseline(baseline_trajs, "/app/occ/reports/baseline_trajectories.jsonl")
-
-    comparator2 = OfflineComparator("/app/occ/reports/baseline_trajectories.jsonl")
-    result = comparator2.compare(candidate_trajs)
-
-    print("\nOffline Comparison Demo")
+    result = comparator.compare(policy_a, policy_b)
     print(json.dumps(result, indent=2, default=str))
     return result
 
 
 if __name__ == "__main__":
-    Path("/app/occ/reports").mkdir(parents=True, exist_ok=True)
-    demo_grpo_hook()
-    print()
-    demo_offline_comparison()
+    demo_offline()