Upload rl/grpo_train_demo.py

rl/grpo_train_demo.py

@@ -1 +1,152 @@
-See /app/occ/rl/grpo_train_demo.py
+"""
+GRPO Training Demonstrator
+Uses Qwen2.5-0.5B-Instruct + DeepMath-103K dataset with cost-aware rewards.
+
+This is a minimal demonstrator showing how OCC rewards can be used
+with TRL's GRPOTrainer. If compute is available, it trains for a few
+steps; otherwise it falls back to offline comparison.
+"""
+import json
+import sys
+from pathlib import Path
+
+from datasets import load_dataset
+
+sys.path.insert(0, str(Path(__file__).parent.parent))
+from rl.reward import RewardHook
+from oracle.oracle import ImpactOracle
+
+
+def occ_reward_func(prompts, completions, **kwargs):
+    """OCC cost-aware reward function for GRPO."""
+    oracle = ImpactOracle(
+        qa_weights={
+            "correctness": 1.0,
+            "evidence_support": 0.5,
+            "calibration": 0.2,
+            "abstention_utility": 1.0,
+            "hallucination_penalty": 2.0,
+            "confident_wrong_penalty": 3.0,
+        }
+    )
+    reward_hook = RewardHook(oracle=oracle, mode="retrieval_qa")
+
+    answers = []
+    confidences = []
+    compute_costs = []
+
+    for comp in completions:
+        if "<answer>" in comp and "</answer>" in comp:
+            start = comp.find("<answer>") + len("<answer>")
+            end = comp.find("</answer>")
+            ans = comp[start:end].strip()
+        else:
+            ans = comp.strip().split()[-1] if comp.strip() else ""
+
+        answers.append(ans)
+        confidences.append(0.7 if len(ans) > 0 else 0.3)
+        compute_costs.append(len(comp.split()))
+
+    gold_answers = kwargs.get("answers", [""] * len(prompts))
+    if not gold_answers:
+        gold_answers = [""] * len(prompts)
+
+    rewards = reward_hook.compute_rewards(
+        prompts=prompts,
+        completions=completions,
+        answers=answers,
+        gold_answers=gold_answers,
+        confidences=confidences,
+        compute_costs=compute_costs,
+    )
+    return rewards
+
+
+def run_offline_demonstrator():
+    """Run offline policy comparison without actual model training."""
+    print("=" * 60)
+    print("GRPO OFFLINE DEMONSTRATOR")
+    print("=" * 60)
+    print("\nAttempting to load DeepMath-103K dataset...")
+    try:
+        ds = load_dataset("trl-lib/DeepMath-103K", split="train")
+        sample = ds.select(range(5))
+        print(f"Dataset loaded: {len(ds)} examples")
+        print(f"Columns: {sample.features}")
+        for i, ex in enumerate(sample):
+            print(f"\nExample {i}:")
+            prompt = ex.get("prompt", "")[:100]
+            solution = ex.get("solution", "")[:100]
+            print(f"  prompt: {prompt}...")
+            print(f"  solution: {solution}...")
+    except Exception as e:
+        print(f"Could not load dataset: {e}")
+        return {"status": "dataset_load_failed", "error": str(e)}
+
+    print("\n--- Simulating policy trajectories ---")
+    policy_a_completions = ["The answer is 42. <answer>42</answer>" for _ in range(10)]
+    policy_b_completions = ["I think it might be 42 or maybe 41. <answer>42</answer>" for _ in range(10)]
+
+    prompts = ["Solve: 20 + 22 = ?"] * 10
+    rewards_a = occ_reward_func(prompts, policy_a_completions, answers=["42"] * 10)
+    rewards_b = occ_reward_func(prompts, policy_b_completions, answers=["42"] * 10)
+
+    print(f"Policy A (concise, confident): mean reward = {sum(rewards_a)/len(rewards_a):.3f}")
+    print(f"Policy B (verbose, uncertain): mean reward = {sum(rewards_b)/len(rewards_b):.3f}")
+
+    from rl.reward import OfflinePolicyComparator
+    comparator = OfflinePolicyComparator(RewardHook(oracle=ImpactOracle(), mode="retrieval_qa"))
+
+    traj_a = [{"reward": r, "failure_tags": []} for r in rewards_a]
+    traj_b = [{"reward": r, "failure_tags": []} for r in rewards_b]
+    comparison = comparator.compare(traj_a, traj_b)
+
+    print(f"\nWin rate (A vs B): {comparison['win_rate']:.1%}")
+    print(f"Mean reward improvement: {comparison['improvement']:+.3f}")
+
+    return {
+        "status": "offline_demo_complete",
+        "policy_a_mean_reward": sum(rewards_a) / len(rewards_a),
+        "policy_b_mean_reward": sum(rewards_b) / len(rewards_b),
+        "comparison": comparison,
+    }
+
+
+def run_grpo_training(steps: int = 50):
+    """Run actual GRPO training if TRL is available."""
+    try:
+        from trl import GRPOTrainer
+        print("\n" + "=" * 60)
+        print("GRPO TRAINING DEMONSTRATION")
+        print("=" * 60)
+        print(f"Loading dataset and model for {steps} steps...")
+
+        ds = load_dataset("trl-lib/DeepMath-103K", split="train")
+
+        trainer = GRPOTrainer(
+            model="Qwen/Qwen2.5-0.5B-Instruct",
+            reward_funcs=occ_reward_func,
+            train_dataset=ds.select(range(100)),
+        )
+        trainer.train()
+        print("Training complete!")
+        return {"status": "training_complete", "steps": steps}
+    except ImportError:
+        print("TRL not installed. Falling back to offline demonstrator.")
+        return run_offline_demonstrator()
+    except Exception as e:
+        print(f"Training failed: {e}")
+        return {"status": "training_failed", "error": str(e)}
+
+
+def main():
+    results = run_grpo_training(steps=10)
+
+    Path("/app/occ/reports").mkdir(parents=True, exist_ok=True)
+    with open("/app/occ/reports/grpo_results.json", "w") as f:
+        json.dump(results, f, indent=2, default=str)
+    print("\nSaved to reports/grpo_results.json")
+
+
+if __name__ == "__main__":
+    main()