Upload grpo_hook.py

grpo_hook.py

@@ -0,0 +1,113 @@
+"""
+Minimal GRPO-compatible reward hook demonstration.
+
+If full GRPO training is feasible, use this with TRL GRPOTrainer.
+If not, use OfflineComparator for policy evaluation.
+"""
+
+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
+
+
+def demo_grpo_hook():
+    """
+    Demonstrate the reward hook with synthetic completions.
+    This is a toy loop showing how GRPO reward computation would work.
+    """
+    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_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)
+    ]
+    # 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")
+    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()