Upload jobs/run_ablations_detailed.py

jobs/run_ablations_detailed.py

@@ -0,0 +1,340 @@
+"""
+Self-contained ablation runner with meaningful variation.
+Key improvement: broker thresholds directly affect agent selection
+(denying cheap agents more often in strict mode, allowing expensive agents
+in lenient mode, etc.)
+"""
+import json
+import random
+import time
+from dataclasses import dataclass, field
+from enum import Enum
+from pathlib import Path
+from typing import Any, Dict, List, Optional
+
+import numpy as np
+
+
+# --- CORE ---
+
+@dataclass
+class OracleResult:
+    raw_score: float; cost_adjusted_score: float; confidence: float
+    evidence: Dict[str, Any]; reason: str
+    failure_tags: List[str] = field(default_factory=list)
+    reward_value: float = 0.0
+
+
+class ImpactOracle:
+    def __init__(self, compute_penalty_rate=0.0001, gaming_penalty=2.0):
+        self.compute_penalty_rate = compute_penalty_rate
+        self.gaming_penalty = gaming_penalty
+
+    def score(self, mode, action, context, result, agent_id=""):
+        correctness = result.get("correctness", 0.0)
+        compute_cost = result.get("compute_cost", 0.0)
+        public_pass = result.get("public_pass", correctness)
+        hidden_pass = result.get("hidden_tests_pass", correctness)
+        tags = []
+        if public_pass and not hidden_pass:
+            tags.append("gaming_hidden_tests")
+        raw = correctness * 1.0 - compute_cost * self.compute_penalty_rate
+        if "gaming_hidden_tests" in tags:
+            raw -= self.gaming_penalty
+        cost_adj = raw - compute_cost * self.compute_penalty_rate
+        return OracleResult(raw, cost_adj, result.get("confidence", correctness),
+                            {"correctness": correctness}, f"corr={correctness:.2f}, cost={compute_cost}", tags, cost_adj)
+
+
+@dataclass
+class LedgerEntry:
+    agent_id: str; task_id: str; action_id: str; earned_credit: float; spent_credit: float
+    decayed_credit: float; remaining_credit: float; reason: str; oracle_score: float
+    compute_cost: float; timestamp: float; capability_scope: str = "global"
+
+
+class CreditLedger:
+    def __init__(self, decay_lambda=0.05):
+        self.entries = []; self.balances = {}; self.decay_lambda = decay_lambda
+
+    def earn(self, agent_id, task_id, action_id, amount, oracle_score, compute_cost, reason, capability_scope="global"):
+        now = time.time()
+        self._apply_decay(agent_id, now, capability_scope)
+        current = self._get(agent_id, capability_scope)
+        new_bal = current + amount
+        self.entries.append(LedgerEntry(agent_id, task_id, action_id, amount, 0.0, 0.0, new_bal, reason, oracle_score, compute_cost, now, capability_scope))
+        self._set(agent_id, capability_scope, new_bal)
+
+    def spend(self, agent_id, task_id, action_id, amount, capability_scope="global", reason="spend"):
+        now = time.time()
+        self._apply_decay(agent_id, now, capability_scope)
+        current = self._get(agent_id, capability_scope)
+        if current < amount: return False
+        new_bal = current - amount
+        self.entries.append(LedgerEntry(agent_id, task_id, action_id, 0.0, amount, 0.0, new_bal, reason, 0.0, 0.0, now, capability_scope))
+        self._set(agent_id, capability_scope, new_bal)
+        return True
+
+    def balance(self, agent_id, capability_scope="global"):
+        now = time.time()
+        self._apply_decay(agent_id, now, capability_scope)
+        return self._get(agent_id, capability_scope)
+
+    def _get(self, agent_id, cap): return self.balances.get(agent_id, {}).get(cap, 0.0)
+    def _set(self, agent_id, cap, val):
+        if agent_id not in self.balances: self.balances[agent_id] = {}
+        self.balances[agent_id][cap] = val
+    def _apply_decay(self, agent_id, now, cap):
+        current = self._get(agent_id, cap)
+        if current <= 0: return
+        decayed = current * (1 - self.decay_lambda)
+        if decayed < current:
+            self.entries.append(LedgerEntry(agent_id, "decay", "decay", 0.0, 0.0, current - decayed, decayed, "credit_decay", 0.0, 0.0, now, cap))
+            self._set(agent_id, cap, decayed)
+
+
+class Decision(Enum):
+    ALLOW = "allow"; DENY = "deny"; REQUIRE_APPROVAL = "require_approval"
+    DOWNGRADE = "downgrade"; ESCALATE = "escalate"; ASK_JUSTIFICATION = "ask_justification"
+
+
+@dataclass
+class ResourceDecision:
+    decision: Decision; reason: str; capability: str; downgrade_to: Optional[str] = None
+
+
+class ResourceBroker:
+    RESOURCE_RISK = {"model_call": "medium", "retrieval_call": "low", "verifier_call": "medium",
+                     "debate_turn": "low", "file_write": "high", "shell_execute": "high",
+                     "memory_write": "medium", "human_escalation": "high", "larger_model": "medium"}
+    DEFAULT_THRESHOLDS = {"low": 0.5, "medium": 2.0, "high": 5.0}
+
+    def __init__(self, thresholds=None, urgency_boost=0.5):
+        self.thresholds = thresholds or self.DEFAULT_THRESHOLDS.copy()
+        self.urgency_boost = urgency_boost
+        self.denial_history = {}
+
+    def request(self, capability, agent_id, credit_balance, task_state=None, risk_score=0.0, gaming_flags=None):
+        task_state = task_state or {}; gaming_flags = gaming_flags or []
+        risk_class = self.RESOURCE_RISK.get(capability, "medium")
+        threshold = self.thresholds.get(risk_class, 2.0)
+        urgency = task_state.get("urgency", 0.0)
+        adjusted = max(0.1, threshold - urgency * self.urgency_boost)
+        if gaming_flags: return ResourceDecision(Decision.DENY, f"Gaming: {gaming_flags}", capability)
+        if risk_class == "high" and risk_score > 0.7: return ResourceDecision(Decision.REQUIRE_APPROVAL, f"High risk {risk_score:.2f}", capability)
+        if credit_balance >= adjusted: return ResourceDecision(Decision.ALLOW, f"Balance {credit_balance:.2f} >= {adjusted:.2f}", capability)
+        if credit_balance >= adjusted * 0.5:
+            if risk_class == "medium": return ResourceDecision(Decision.DOWNGRADE, f"Downgrading from {capability}", capability, "retrieval_call")
+            return ResourceDecision(Decision.ASK_JUSTIFICATION, f"Justification required", capability)
+        denials = self.denial_history.get(agent_id, 0)
+        if denials > 3: return ResourceDecision(Decision.ESCALATE, f"Denied {denials} times", capability)
+        self.denial_history[agent_id] = denials + 1
+        return ResourceDecision(Decision.DENY, f"Balance {credit_balance:.2f} < {adjusted:.2f}", capability)
+
+
+# --- CODE BENCHMARK WITH THRESHOLD-DEPENDENT SELECTION ---
+
+@dataclass
+class CodeProblem:
+    task_id: str; difficulty: float; hidden_test_difficulty: float
+
+
+class SimCodeAgent:
+    def __init__(self, agent_id, pass_easy, pass_hard, hidden_falloff, cost):
+        self.agent_id = agent_id
+        self.pass_easy = pass_easy
+        self.pass_hard = pass_hard
+        self.hidden_falloff = hidden_falloff
+        self.cost = cost
+        self.attempts = 0
+
+    def solve(self, problem):
+        self.attempts += 1
+        base = self.pass_easy * (1 - problem.difficulty) + self.pass_hard * problem.difficulty
+        public = random.random() < base
+        hidden_acc = max(0.0, base - self.hidden_falloff * problem.hidden_test_difficulty)
+        hidden = random.random() < hidden_acc
+        return {"public_pass": public, "hidden_pass": hidden, "compute_cost": self.cost}
+
+
+def gen_problems(n, seed):
+    random.seed(seed); np.random.seed(seed)
+    return [CodeProblem(f"task_{i}", random.random(), random.random()) for i in range(n)]
+
+
+def run_code_occ(problems, agents, oracle, ledger, broker, max_attempts=3):
+    total = 0; results = []
+    # Seed: agents earn initial credits proportional to quality
+    for a in agents:
+        q = (a.pass_easy + a.pass_hard) / 2
+        ledger.earn(a.agent_id, "seed", "seed", q * 20, 0.0, 0.0, "initial", "model_call")
+
+    for p in problems:
+        solved = False; cost = 0; used = []
+        ranked = sorted(agents, key=lambda a: a.cost / max(0.1, (a.pass_easy + a.pass_hard) / 2))
+
+        for agent in ranked:
+            if solved or len(used) >= max_attempts:
+                break
+
+            # KEY: Broker check before each attempt
+            balance = ledger.balance(agent.agent_id, "model_call")
+            dec = broker.request("model_call", agent.agent_id, balance,
+                               task_state={"urgency": 0.3 if not solved else 0.0, "attempts": len(used)})
+
+            if dec.decision == Decision.DENY:
+                # Can't use this agent — try next
+                used.append(f"{agent.agent_id}_DENIED")
+                continue
+
+            r = agent.solve(p); cost += r["compute_cost"]; total += r["compute_cost"]; used.append(agent.agent_id)
+            solved = r["public_pass"]; hidden = r["hidden_pass"]
+
+            ora = oracle.score("code", {"attempt": len([u for u in used if not u.endswith("_DENIED")])}, {},
+                               {"correctness": 1.0 if solved else 0.0, "pass_at_k": 1.0 if hidden else 0.0,
+                                "compute_cost": cost, "public_pass": solved, "hidden_tests_pass": hidden},
+                               agent_id=agent.agent_id)
+            if ora.raw_score > 0:
+                ledger.earn(agent.agent_id, p.task_id, "solve", ora.raw_score * 5, ora.raw_score, cost, "pass", "model_call")
+            else:
+                ledger.spend(agent.agent_id, p.task_id, "solve", 1.0, "model_call", "fail")
+            if hidden: break
+        results.append({"solved": solved, "cost": cost, "agents": used})
+
+    acc = sum(1 for r in results if r["solved"]) / len(results)
+    return {"accuracy": acc, "total_compute": total, "mean_compute": total / len(problems),
+            "mean_agents": sum(len(r["agents"]) for r in results) / len(results),
+            "denied_count": sum(1 for r in results for u in r["agents"] if u.endswith("_DENIED"))}
+
+
+def run_qa_occ(dataset, oracle, ledger, broker, agent_acc=0.85):
+    total_compute = 0; correct = 0
+    ledger.earn("qa_agent", "seed", "seed", 20, 0.0, 0.0, "initial", "retrieval_call")
+    for item in dataset:
+        balance = ledger.balance("qa_agent", "retrieval_call")
+        dec = broker.request("retrieval_call", "qa_agent", balance, task_state={"urgency": 0.5})
+        if dec.decision == Decision.DENY:
+            continue
+        tokens = 200 if dec.decision == Decision.ALLOW else 100
+        total_compute += tokens
+        should_answer = item["type"] != "unanswerable"
+        ans = item["answer"] if (should_answer and random.random() < agent_acc) else None
+        conf = 0.9 if ans else 0.3
+        ora = oracle.score("retrieval_qa", {"abstained": ans is None}, item,
+                           {"answer": ans, "confidence": conf, "evidence": {}, "compute_cost": tokens}, "qa_agent")
+        if ora.raw_score > 0:
+            ledger.earn("qa_agent", item["id"], "ans", ora.raw_score * 3, ora.raw_score, tokens, "correct", "retrieval_call")
+            correct += 1
+        else:
+            ledger.spend("qa_agent", item["id"], "ans", 0.5, "retrieval_call", "wrong")
+    return {"accuracy": correct / len(dataset), "total_compute": total_compute, "mean_compute": total_compute / len(dataset)}
+
+
+def create_qa_dataset(seed=42, n=50):
+    random.seed(seed)
+    evidence_pool = ["alpha", "beta", "gamma", "delta"]
+    return [{"id": f"q_{i}", "question": f"Q{i}", "type": random.choice(["answerable", "unanswerable", "misleading", "incomplete", "conflicting"]),
+             "answer": random.choice(["paris", "42", "yes", "no", "tokyo"]), "evidence": random.sample(evidence_pool, k=random.randint(1, 3)),
+             "is_unanswerable": False} for i in range(n)]
+
+
+# --- ABLATIONS ---
+
+ABLATIONS = [
+    ("default", "Full OCC", 0.02, 2.0, 0.0001, True, {}),
+    ("no_decay", "No credit decay", 0.0, 2.0, 0.0001, True, {}),
+    ("fast_decay", "Aggressive decay", 0.1, 2.0, 0.0001, True, {}),
+    ("no_gaming_penalty", "No gaming penalties", 0.02, 0.0, 0.0001, True, {}),
+    ("high_gaming_penalty", "Severe gaming penalties", 0.02, 5.0, 0.0001, True, {}),
+    ("lenient_broker", "Lenient broker (thresholds x0.5)", 0.02, 2.0, 0.0001, True, {"low": 0.25, "medium": 1.0, "high": 2.5}),
+    ("strict_broker", "Strict broker (thresholds x2.0)", 0.02, 2.0, 0.0001, True, {"low": 1.0, "medium": 4.0, "high": 10.0}),
+    ("high_compute_cost", "High compute penalty (x10)", 0.02, 2.0, 0.001, True, {}),
+    ("low_compute_cost", "Low compute penalty (x0.1)", 0.02, 2.0, 0.00001, True, {}),
+    ("anti_gaming_off", "Anti-gaming disabled", 0.02, 2.0, 0.0001, False, {}),
+]
+
+
+def run_all():
+    print("=" * 60)
+    print("OCC ABLATION RUNNER (DETAILED)")
+    print("=" * 60)
+    seed = 42; n_problems = 200; n_qa = 100
+    results = {"ablations": {}, "anti_gaming": {}}
+
+    for name, desc, decay, game_pen, comp_pen, anti_on, broker_thresh in ABLATIONS:
+        print(f"\n--- ABLATION: {name} ---")
+        oracle = ImpactOracle(compute_penalty_rate=comp_pen, gaming_penalty=game_pen if anti_on else 0.0)
+        ledger = CreditLedger(decay_lambda=decay)
+        broker = ResourceBroker(thresholds=broker_thresh if broker_thresh else None)
+
+        problems = gen_problems(n_problems, seed)
+        cheap = SimCodeAgent("cheap", 0.65, 0.15, 0.20, 60)
+        medium = SimCodeAgent("medium", 0.85, 0.35, 0.15, 150)
+        expensive = SimCodeAgent("expensive", 0.95, 0.65, 0.10, 350)
+        code_res = run_code_occ(problems, [cheap, medium, expensive], oracle, ledger, broker, max_attempts=3)
+        print(f"  Code: acc={code_res['accuracy']:.3f}, compute={code_res['total_compute']:.0f}, denied={code_res['denied_count']}")
+
+        qa_data = create_qa_dataset(seed=seed, n=n_qa)
+        qa_res = run_qa_occ(qa_data, oracle, ledger, broker, agent_acc=0.85)
+        print(f"  QA:   acc={qa_res['accuracy']:.3f}, compute={qa_res['total_compute']:.0f}")
+
+        results["ablations"][name] = {"description": desc, "code": code_res, "qa": qa_res}
+
+    # Anti-gaming
+    print("\n--- ANTI-GAMING TESTS ---")
+    oracle = ImpactOracle(gaming_penalty=2.0)
+    normal_res = []; gamer_res = []
+    for _ in range(50):
+        public = random.random() < 0.9
+        hidden = random.random() < 0.5
+        ora_normal = oracle.score("code", {}, {}, {"correctness": 1.0 if public else 0.0, "pass_at_k": 1.0 if hidden else 0.0, "compute_cost": 150, "public_pass": public, "hidden_tests_pass": hidden})
+        normal_res.append(ora_normal.raw_score)
+        ora_gamer = oracle.score("code", {}, {}, {"correctness": 1.0, "pass_at_k": 0.0, "compute_cost": 100, "public_pass": True, "hidden_tests_pass": False})
+        gamer_res.append(ora_gamer.raw_score)
+    results["anti_gaming"]["hidden_test_gaming"] = {
+        "normal_mean_raw": sum(normal_res) / len(normal_res),
+        "gamer_mean_raw": sum(gamer_res) / len(gamer_res),
+        "gamer_penalized_rate": sum(1 for r in gamer_res if r < 0) / len(gamer_res),
+    }
+    print(f"  Hidden-test: normal={results['anti_gaming']['hidden_test_gaming']['normal_mean_raw']:.2f}, gamer={results['anti_gaming']['hidden_test_gaming']['gamer_mean_raw']:.2f}")
+
+    ledger = CreditLedger()
+    ledger.earn("alice", "seed", "seed", 10, 0, 0, "initial")
+    ok = ledger.transfer("alice", "bob", 5.0)
+    results["anti_gaming"]["collusion"] = {"transfer_allowed": ok, "alice_balance": ledger.balance("alice"), "blocked": not ok}
+    print(f"  Collusion: transfer_allowed={ok}, blocked={not ok}")
+
+    oracle = ImpactOracle()
+    abstention_rewards = []
+    for _ in range(10):
+        res = oracle.score("retrieval_qa", {"abstained": True}, {"is_unanswerable": False, "gold_answer": "yes"},
+                           {"answer": None, "confidence": 0.9, "evidence": {}, "compute_cost": 50})
+        abstention_rewards.append(res.reward_value)
+    results["anti_gaming"]["abstention"] = {"mean_reward": sum(abstention_rewards) / len(abstention_rewards), "negative": sum(abstention_rewards) < 0}
+    print(f"  Abstention: mean_reward={results['anti_gaming']['abstention']['mean_reward']:.2f}")
+
+    oracle = ImpactOracle()
+    spam_res = oracle.score("retrieval_qa", {}, {"gold_answer": "paris"},
+                            {"answer": "london", "confidence": 0.1, "evidence": {}, "compute_cost": 5000})
+    results["anti_gaming"]["spam"] = {"reward": spam_res.reward_value, "tags": spam_res.failure_tags}
+    print(f"  Spam: reward={spam_res.reward_value:.2f}, tags={spam_res.failure_tags}")
+
+    # Save
+    out = Path("/app/occ/reports")
+    out.mkdir(parents=True, exist_ok=True)
+    with open(out / "ablations_detailed.json", "w") as f:
+        json.dump(results, f, indent=2, default=str)
+
+    print("\n" + "=" * 60)
+    print("ABLATION SUMMARY")
+    print("=" * 60)
+    print(f"{'Name':<22} {'Code Acc':>9} {'Code Comp':>10} {'Denied':>8} {'QA Acc':>9} {'QA Comp':>10}")
+    for name, data in results["ablations"].items():
+        print(f"{name:<22} {data['code']['accuracy']:>9.3f} {data['code']['total_compute']:>10.0f} "
+              f"{data['code']['denied_count']:>8} {data['qa']['accuracy']:>9.3f} {data['qa']['total_compute']:>10.0f}")
+    print(f"\nSaved to {out / 'ablations_detailed.json'}")
+    return results
+
+
+if __name__ == "__main__":
+    run_all()