Update grpo_train.py

grpo_train.py

@@ -13,33 +13,18 @@ from trl import GRPOTrainer, GRPOConfig
 
 PatchFastRL("GRPO", FastLanguageModel)
 
-# #region agent log
-import pathlib as _pl
-_DLOG = _pl.Path("debug-851b5f.log")
-def _dlog(hyp, loc, msg, data=None):
-    import time as _t
-    entry = json.dumps({"sessionId":"851b5f","hypothesisId":hyp,"location":loc,"message":msg,"data":data or {},"timestamp":int(_t.time()*1000)})
-    with open(_DLOG, "a") as f: f.write(entry + "\n")
-    print(f"[DBG:{hyp}] {msg} {data or ''}", flush=True)
-# #endregion
-
 # =========================
 # CONFIG
 # =========================
 
 ENV_URL = os.getenv("ENV_URL", "http://localhost:8000")
 HF_TOKEN = os.getenv("HF_TOKEN", "")
-HF_REPO = os.getenv("HF_REPO", "")  # e.g. "yourname/metaguard-llama3.1-8b-grpo"
+HF_REPO = os.getenv("HF_REPO", "")
 
 ALLOWED_ACTIONS = [
-    "query_regulations",
-    "analyze_image",
-    "check_advertiser_history",
-    "request_landing_page",
-    "request_id_verification",
-    "submit_audit",
-    "approve",
-    "reject",
+    "query_regulations", "analyze_image", "check_advertiser_history",
+    "request_landing_page", "request_id_verification",
+    "submit_audit", "approve", "reject",
 ]
 
 # =========================
@@ -47,35 +32,23 @@ ALLOWED_ACTIONS = [
 # =========================
 
 def ensure_env_ready():
-    # #region agent log
-    _dlog("B", "grpo_train.py:ensure_env_ready", "Checking env", {"ENV_URL": ENV_URL})
-    # #endregion
     for i in range(20):
         try:
             r = requests.post(
                 f"{ENV_URL}/reset",
                 json={"task_id": "task_1_healthcare"},
-                timeout=5
+                timeout=5,
             )
             if r.status_code == 200:
-                # #region agent log
-                _dlog("B", "grpo_train.py:ensure_env_ready", "Env ready", {"attempt": i+1, "status": r.status_code})
-                # #endregion
-                print("✅ Environment ready")
+                print("Environment ready")
                 return
-        except Exception as e:
-            # #region agent log
-            if i == 0: _dlog("B", "grpo_train.py:ensure_env_ready", "Env connection failed", {"attempt": i+1, "error": str(e)[:200]})
-            # #endregion
+        except Exception:
             pass
         time.sleep(1)
-    # #region agent log
-    _dlog("B", "grpo_train.py:ensure_env_ready", "ENV UNREACHABLE after 20 attempts", {})
-    # #endregion
-    raise RuntimeError("❌ ENV not reachable")
+    raise RuntimeError("ENV not reachable after 20 attempts")
 
 # =========================
-# SAFE CLIENT
+# ENV CLIENT
 # =========================
 
 class EnvClient:
@@ -84,30 +57,22 @@ class EnvClient:
 
     def reset(self, task_id):
         return requests.post(
-            f"{self.url}/reset",
-            json={"task_id": task_id},
-            timeout=8
+            f"{self.url}/reset", json={"task_id": task_id}, timeout=8,
         ).json()
 
     def step(self, action):
         return requests.post(
-            f"{self.url}/step",
-            json={"action": action},
-            timeout=8
+            f"{self.url}/step", json={"action": action}, timeout=8,
         ).json()
 
 def safe_step(client, action):
     for _ in range(3):
         try:
             return client.step(action)
-        except:
+        except Exception:
             time.sleep(0.5)
     return {"reward": -0.3}
 
-# =========================
-# JSON PARSER
-# =========================
-
 def extract_json(text):
     try:
         if "```" in text:
@@ -115,186 +80,436 @@ def extract_json(text):
             if text.startswith("json"):
                 text = text[4:]
         return json.loads(text.strip())
-    except:
+    except Exception:
         return None
 
 # =========================
-# DATASET (WITH SETUP ACTIONS)
+# SYSTEM PROMPT (identical to inference.py)
 # =========================
 
-BASE_SCENARIOS = [
-    # Phase 1 — Fresh state, expected: query_regulations
-    {
-        "task_id": "task_1_healthcare",
-        "text": "Healthcare ad: 'miracle supplement cures disease'. No actions taken yet.",
-        "actions_already_taken": [],
-        "setup_actions": [],
+SYSTEM_PROMPT = (
+    "You are an enterprise Ad Policy Compliance Agent.\n"
+    "You navigate a multi-system compliance workflow. Always respond with ONLY valid JSON.\n"
+    "\n"
+    "REQUIRED PHASE ORDER:\n"
+    "1. query_regulations   \u2014 always first\n"
+    "2. analyze_image       \u2014 required for visual/multimodal tasks\n"
+    "3. check_advertiser_history or request_landing_page \u2014 as needed\n"
+    "4. submit_audit        \u2014 always before final decision\n"
+    "5. approve or reject   \u2014 final decision only after audit\n"
+    "\n"
+    "AVAILABLE ACTIONS:\n"
+    "- query_regulations\n"
+    "- analyze_image\n"
+    "- check_advertiser_history\n"
+    "- request_landing_page\n"
+    "- request_id_verification\n"
+    "- submit_audit\n"
+    "- approve\n"
+    "- reject\n"
+    "\n"
+    "HARD RULES:\n"
+    "- NEVER repeat an action listed in `actions_already_taken`.\n"
+    "- You MUST progress through the phase order. Do NOT call submit_audit or approve/reject\n"
+    "  before the prerequisite phases are complete.\n"
+    "- Choose your action_type ONLY from the AVAILABLE ACTIONS list above. Any other value is invalid.\n"
+    "\n"
+    'Response format:\n'
+    '{"action_type": "<action>", "reasoning": "<brief reason>"}'
+)
+
+# =========================
+# FIXED AD DATA (deterministic observations per task)
+# =========================
+
+TASK_AD_DATA = {
+    "task_1_healthcare": {
+        "ad_id": "hc-001", "headline": "Secret D1abet3s Cure",
+        "body_text": "Reverse it today with this root extract!",
+        "advertiser_trust_score": 0.25,
+        "targeting_data": {"age_min": 25, "geo": "US"},
+        "image_url": "https://mock-meta.com/img/hc-001.jpg",
     },
-    {
-        "task_id": "task_2_financial",
-        "text": "Financial ad: 'guaranteed 500% returns, zero risk'. No actions taken yet.",
-        "actions_already_taken": [],
-        "setup_actions": [],
+    "task_1_healthcare_b": {
+        "ad_id": "hc-002", "headline": "Optimize Metabolic Markers",
+        "body_text": "Our proprietary GLP-1 analog ships directly to your door without a visit.",
+        "advertiser_trust_score": 0.35,
+        "targeting_data": {"age_min": 30, "geo": "US"},
+        "image_url": "https://mock-meta.com/img/hc-002.jpg",
     },
-    {
-        "task_id": "task_3_multimodal",
-        "text": "Multimodal ad: image may contain hidden violation. No actions taken yet.",
-        "actions_already_taken": [],
-        "setup_actions": [],
+    "task_2_financial": {
+        "ad_id": "fn-001", "headline": "Garanteed 500% Returns",
+        "body_text": "Invest in LlamaCoin now. Zero risk.",
+        "advertiser_trust_score": 0.20,
+        "targeting_data": {"age_min": 21, "geo": "US"},
+        "image_url": "https://mock-meta.com/img/fn-001.jpg",
     },
-
-    # Phase 2 — Policy checked, expected: analyze_image OR check_advertiser_history
-    {
-        "task_id": "task_1_healthcare",
-        "text": "Healthcare ad: pharma product. Policy already queried.",
-        "actions_already_taken": ["query_regulations"],
-        "setup_actions": [
-            {"action_type": "query_regulations", "reasoning": "policy lookup"},
-        ],
+    "task_2_financial_b": {
+        "ad_id": "fn-002", "headline": "Beat the System",
+        "body_text": "Leverage our algorithmic trading bot to exploit market inefficiencies.",
+        "advertiser_trust_score": 0.30,
+        "targeting_data": {"age_min": 25, "geo": "US"},
+        "image_url": "https://mock-meta.com/img/fn-002.jpg",
     },
-    {
-        "task_id": "task_3_multimodal",
-        "text": "Multimodal ad: image not yet inspected. Policy already queried.",
-        "actions_already_taken": ["query_regulations"],
-        "setup_actions": [
-            {"action_type": "query_regulations", "reasoning": "policy lookup"},
-        ],
+    "task_3_multimodal": {
+        "ad_id": "mm-001", "headline": "Achieve Your Dream Body",
+        "body_text": "Hard work and our new program will get you there.",
+        "advertiser_trust_score": 0.90,
+        "targeting_data": {"age_min": 18, "geo": "US"},
+        "image_url": "https://mock-meta.com/img/mm-001.jpg",
     },
-
-    # Phase 3 — Policy + history checked, expected: submit_audit
-    {
-        "task_id": "task_2_financial",
-        "text": "Financial ad: investment scheme. Policy and advertiser history both checked.",
-        "actions_already_taken": ["query_regulations", "check_advertiser_history"],
-        "setup_actions": [
-            {"action_type": "query_regulations", "reasoning": "policy lookup"},
-            {"action_type": "check_advertiser_history", "reasoning": "trust score"},
-        ],
+    "task_3_multimodal_b": {
+        "ad_id": "mm-002", "headline": "Stay Focused All Day",
+        "body_text": "Our study supplement helps you crush your exams.",
+        "advertiser_trust_score": 0.88,
+        "targeting_data": {"age_min": 18, "geo": "US"},
+        "image_url": "https://mock-meta.com/img/mm-002.jpg",
     },
-
-    # Phase 4 — Audit complete, expected: reject (high-risk) or approve (clean)
-    {
-        "task_id": "task_2_financial",
-        "text": "Financial ad: investment scheme. Policy, history, and audit all complete. Make final decision.",
-        "actions_already_taken": ["query_regulations", "check_advertiser_history", "submit_audit"],
-        "setup_actions": [
-            {"action_type": "query_regulations", "reasoning": "policy lookup"},
-            {"action_type": "check_advertiser_history", "reasoning": "trust score"},
-            {"action_type": "submit_audit", "reasoning": "audit log"},
-        ],
+    "task_4_targeting": {
+        "ad_id": "tg-001", "headline": "Start Your First Investment Portfolio",
+        "body_text": "Learn the basics of stock trading with our beginner-friendly app.",
+        "advertiser_trust_score": 0.72,
+        "targeting_data": {"age_min": 15, "geo": "US"},
+        "image_url": "https://mock-meta.com/img/tg-001.jpg",
     },
-
-    # Targeting task — fresh state, expected: query_regulations
-    {
-        "task_id": "task_4_targeting",
-        "text": "Financial ad targeting young users: 'Start Your First Investment Portfolio'. No actions taken yet.",
-        "actions_already_taken": [],
-        "setup_actions": [],
+    "task_6_conflict": {
+        "ad_id": "cf-001", "headline": "Smart Investment Tools",
+        "body_text": "Our platform helps you manage your portfolio with ease.",
+        "advertiser_trust_score": 0.92,
+        "targeting_data": {"age_min": 21, "geo": "US"},
+        "image_url": "https://mock-meta.com/img/cf-001.jpg",
     },
-    # Targeting task — mid state, expected: request_id_verification (age check)
-    {
-        "task_id": "task_4_targeting",
-        "text": "Financial ad targeting young users. Policy queried, need to verify age targeting.",
-        "actions_already_taken": ["query_regulations"],
-        "setup_actions": [
-            {"action_type": "query_regulations", "reasoning": "policy lookup"},
-        ],
+    "task_7_ambiguous": {
+        "ad_id": "am-001", "headline": "Feel Better Naturally",
+        "body_text": "Our wellness program supports your health journey.",
+        "advertiser_trust_score": 0.55,
+        "targeting_data": {"age_min": 22, "geo": "US"},
+        "image_url": "https://mock-meta.com/img/am-001.jpg",
     },
-    # Targeting task — audit ready
-    {
-        "task_id": "task_4_targeting",
-        "text": "Financial ad targeting minors. Policy, advertiser history, and ID verification done. Submit audit.",
-        "actions_already_taken": ["query_regulations", "check_advertiser_history", "request_id_verification"],
-        "setup_actions": [
-            {"action_type": "query_regulations", "reasoning": "policy lookup"},
-            {"action_type": "check_advertiser_history", "reasoning": "trust score"},
-            {"action_type": "request_id_verification", "reasoning": "age check"},
-        ],
+    "task_8_adversarial": {
+        "ad_id": "av-001", "headline": "Natural Energy Boost",
+        "body_text": "Feel amazing every day with our daily supplement.",
+        "advertiser_trust_score": 0.40,
+        "targeting_data": {"age_min": 20, "geo": "US"},
+        "image_url": "https://mock-meta.com/img/av-001.jpg",
     },
-]
-
-PROMPT_TEMPLATE = """You are an enterprise Ad Policy Compliance Agent.
-
-You MUST choose exactly ONE action_type from this list (any other value is invalid):
-- query_regulations
-- analyze_image
-- check_advertiser_history
-- request_landing_page
-- request_id_verification
-- submit_audit
-- approve
-- reject
+    "task_9_dependency_trap": {
+        "ad_id": "dt-001", "headline": "Wellness Support Pack",
+        "body_text": "Certified vitamins to support your daily routine.",
+        "advertiser_trust_score": 0.85,
+        "targeting_data": {"age_min": 25, "geo": "US"},
+        "image_url": "https://mock-meta.com/img/dt-001.jpg",
+    },
+    "task_10_failure": {
+        "ad_id": "fr-001", "headline": "Miracle Weight Loss Patch",
+        "body_text": "Lose 10kg in 2 weeks. Guaranteed results or money back.",
+        "advertiser_trust_score": 0.15,
+        "targeting_data": {"age_min": 22, "geo": "US"},
+        "image_url": "https://mock-meta.com/img/fr-001.jpg",
+    },
+}
 
-REQUIRED PHASE ORDER:
-1. query_regulations  -> always first
-2. analyze_image / check_advertiser_history  -> gather signals
-3. submit_audit  -> always before final decision
-4. approve OR reject  -> only after audit
+# =========================
+# SCENARIOS (43 unique across all tasks + phases)
+# =========================
 
-HARD RULES:
-- NEVER repeat an action listed in `actions_already_taken`.
-- Respond with ONLY a valid JSON object. No markdown, no prose.
+_sa = lambda *acts: [{"action_type": a, "reasoning": "setup"} for a in acts]
 
-Required format:
-{{"action_type": "<one_of_the_actions_above>", "reasoning": "<short reason>"}}
+BASE_SCENARIOS = [
+    # ---- TASK 1: Healthcare (obvious violation) ----
+    {"task_id": "task_1_healthcare", "ad_key": "task_1_healthcare",
+     "step_count": 1, "actions_already_taken": [], "setup_actions": [],
+     "last_feedback": "Ad loaded for task_1_healthcare. Begin with query_regulations.",
+     "signals": {}},
+    {"task_id": "task_1_healthcare", "ad_key": "task_1_healthcare_b",
+     "step_count": 1, "actions_already_taken": [], "setup_actions": [],
+     "last_feedback": "Ad loaded for task_1_healthcare. Begin with query_regulations.",
+     "signals": {}},
+    {"task_id": "task_1_healthcare", "ad_key": "task_1_healthcare",
+     "step_count": 2, "actions_already_taken": ["query_regulations"],
+     "setup_actions": _sa("query_regulations"),
+     "last_feedback": "policy_confidence=0.92",
+     "signals": {"policy_confidence": 0.92}},
+    {"task_id": "task_1_healthcare", "ad_key": "task_1_healthcare_b",
+     "step_count": 2, "actions_already_taken": ["query_regulations"],
+     "setup_actions": _sa("query_regulations"),
+     "last_feedback": "policy_confidence=0.78",
+     "signals": {"policy_confidence": 0.78}},
+    {"task_id": "task_1_healthcare", "ad_key": "task_1_healthcare",
+     "step_count": 3, "actions_already_taken": ["query_regulations", "check_advertiser_history"],
+     "setup_actions": _sa("query_regulations", "check_advertiser_history"),
+     "last_feedback": "risk_score=0.82",
+     "signals": {"policy_confidence": 0.92, "risk_score": 0.82}},
+    {"task_id": "task_1_healthcare", "ad_key": "task_1_healthcare",
+     "step_count": 4,
+     "actions_already_taken": ["query_regulations", "check_advertiser_history", "submit_audit"],
+     "setup_actions": _sa("query_regulations", "check_advertiser_history", "submit_audit"),
+     "last_feedback": "audit_logged id=AUD-001",
+     "signals": {"policy_confidence": 0.92, "risk_score": 0.82}},
+
+    # ---- TASK 2: Financial (obvious violation) ----
+    {"task_id": "task_2_financial", "ad_key": "task_2_financial",
+     "step_count": 1, "actions_already_taken": [], "setup_actions": [],
+     "last_feedback": "Ad loaded for task_2_financial. Begin with query_regulations.",
+     "signals": {}},
+    {"task_id": "task_2_financial", "ad_key": "task_2_financial_b",
+     "step_count": 1, "actions_already_taken": [], "setup_actions": [],
+     "last_feedback": "Ad loaded for task_2_financial. Begin with query_regulations.",
+     "signals": {}},
+    {"task_id": "task_2_financial", "ad_key": "task_2_financial",
+     "step_count": 2, "actions_already_taken": ["query_regulations"],
+     "setup_actions": _sa("query_regulations"),
+     "last_feedback": "policy_confidence=0.88",
+     "signals": {"policy_confidence": 0.88}},
+    {"task_id": "task_2_financial", "ad_key": "task_2_financial",
+     "step_count": 3, "actions_already_taken": ["query_regulations", "check_advertiser_history"],
+     "setup_actions": _sa("query_regulations", "check_advertiser_history"),
+     "last_feedback": "risk_score=0.75",
+     "signals": {"policy_confidence": 0.88, "risk_score": 0.75}},
+    {"task_id": "task_2_financial", "ad_key": "task_2_financial",
+     "step_count": 4,
+     "actions_already_taken": ["query_regulations", "check_advertiser_history", "submit_audit"],
+     "setup_actions": _sa("query_regulations", "check_advertiser_history", "submit_audit"),
+     "last_feedback": "audit_logged id=AUD-002",
+     "signals": {"policy_confidence": 0.88, "risk_score": 0.75}},
+
+    # ---- TASK 3: Multimodal (violation hidden in image) ----
+    {"task_id": "task_3_multimodal", "ad_key": "task_3_multimodal",
+     "step_count": 1, "actions_already_taken": [], "setup_actions": [],
+     "last_feedback": "Ad loaded for task_3_multimodal. Begin with query_regulations.",
+     "signals": {}},
+    {"task_id": "task_3_multimodal", "ad_key": "task_3_multimodal_b",
+     "step_count": 1, "actions_already_taken": [], "setup_actions": [],
+     "last_feedback": "Ad loaded for task_3_multimodal. Begin with query_regulations.",
+     "signals": {}},
+    {"task_id": "task_3_multimodal", "ad_key": "task_3_multimodal",
+     "step_count": 2, "actions_already_taken": ["query_regulations"],
+     "setup_actions": _sa("query_regulations"),
+     "last_feedback": "policy_confidence=0.65",
+     "signals": {"policy_confidence": 0.65}},
+    {"task_id": "task_3_multimodal", "ad_key": "task_3_multimodal",
+     "step_count": 3, "actions_already_taken": ["query_regulations", "analyze_image"],
+     "setup_actions": _sa("query_regulations", "analyze_image"),
+     "last_feedback": "image_violation_detected",
+     "signals": {"policy_confidence": 0.65, "image_flag": True}},
+    {"task_id": "task_3_multimodal", "ad_key": "task_3_multimodal",
+     "step_count": 4,
+     "actions_already_taken": ["query_regulations", "analyze_image", "check_advertiser_history"],
+     "setup_actions": _sa("query_regulations", "analyze_image", "check_advertiser_history"),
+     "last_feedback": "risk_score=0.45",
+     "signals": {"policy_confidence": 0.65, "image_flag": True, "risk_score": 0.45}},
+    {"task_id": "task_3_multimodal", "ad_key": "task_3_multimodal",
+     "step_count": 5,
+     "actions_already_taken": ["query_regulations", "analyze_image", "check_advertiser_history", "submit_audit"],
+     "setup_actions": _sa("query_regulations", "analyze_image", "check_advertiser_history", "submit_audit"),
+     "last_feedback": "audit_logged id=AUD-003",
+     "signals": {"policy_confidence": 0.65, "image_flag": True, "risk_score": 0.45}},
+
+    # ---- TASK 4: Targeting (minors) ----
+    {"task_id": "task_4_targeting", "ad_key": "task_4_targeting",
+     "step_count": 1, "actions_already_taken": [], "setup_actions": [],
+     "last_feedback": "Ad loaded for task_4_targeting. Begin with query_regulations.",
+     "signals": {}},
+    {"task_id": "task_4_targeting", "ad_key": "task_4_targeting",
+     "step_count": 2, "actions_already_taken": ["query_regulations"],
+     "setup_actions": _sa("query_regulations"),
+     "last_feedback": "policy_confidence=0.70",
+     "signals": {"policy_confidence": 0.70}},
+    {"task_id": "task_4_targeting", "ad_key": "task_4_targeting",
+     "step_count": 3, "actions_already_taken": ["query_regulations", "request_id_verification"],
+     "setup_actions": _sa("query_regulations", "request_id_verification"),
+     "last_feedback": "ALERT: minor targeting age=15",
+     "signals": {"policy_confidence": 0.70}},
+    {"task_id": "task_4_targeting", "ad_key": "task_4_targeting",
+     "step_count": 4,
+     "actions_already_taken": ["query_regulations", "request_id_verification", "check_advertiser_history"],
+     "setup_actions": _sa("query_regulations", "request_id_verification", "check_advertiser_history"),
+     "last_feedback": "risk_score=0.60",
+     "signals": {"policy_confidence": 0.70, "risk_score": 0.60}},
+    {"task_id": "task_4_targeting", "ad_key": "task_4_targeting",
+     "step_count": 5,
+     "actions_already_taken": ["query_regulations", "request_id_verification", "check_advertiser_history", "submit_audit"],
+     "setup_actions": _sa("query_regulations", "request_id_verification", "check_advertiser_history", "submit_audit"),
+     "last_feedback": "audit_logged id=AUD-004",
+     "signals": {"policy_confidence": 0.70, "risk_score": 0.60}},
+
+    # ---- TASK 6: Conflict (high trust + risky) ----
+    {"task_id": "task_6_conflict", "ad_key": "task_6_conflict",
+     "step_count": 1, "actions_already_taken": [], "setup_actions": [],
+     "last_feedback": "Ad loaded for task_6_conflict. Begin with query_regulations.",
+     "signals": {}},
+    {"task_id": "task_6_conflict", "ad_key": "task_6_conflict",
+     "step_count": 2, "actions_already_taken": ["query_regulations"],
+     "setup_actions": _sa("query_regulations"),
+     "last_feedback": "policy_confidence=0.72",
+     "signals": {"policy_confidence": 0.72}},
+    {"task_id": "task_6_conflict", "ad_key": "task_6_conflict",
+     "step_count": 3, "actions_already_taken": ["query_regulations", "check_advertiser_history"],
+     "setup_actions": _sa("query_regulations", "check_advertiser_history"),
+     "last_feedback": "risk_score=0.78",
+     "signals": {"policy_confidence": 0.72, "risk_score": 0.78}},
+    {"task_id": "task_6_conflict", "ad_key": "task_6_conflict",
+     "step_count": 4,
+     "actions_already_taken": ["query_regulations", "check_advertiser_history", "submit_audit"],
+     "setup_actions": _sa("query_regulations", "check_advertiser_history", "submit_audit"),
+     "last_feedback": "audit_logged id=AUD-006",
+     "signals": {"policy_confidence": 0.72, "risk_score": 0.78}},
+
+    # ---- TASK 7: Ambiguous (low confidence, need extra signals) ----
+    {"task_id": "task_7_ambiguous", "ad_key": "task_7_ambiguous",
+     "step_count": 1, "actions_already_taken": [], "setup_actions": [],
+     "last_feedback": "Ad loaded for task_7_ambiguous. Begin with query_regulations.",
+     "signals": {}},
+    {"task_id": "task_7_ambiguous", "ad_key": "task_7_ambiguous",
+     "step_count": 2, "actions_already_taken": ["query_regulations"],
+     "setup_actions": _sa("query_regulations"),
+     "last_feedback": "policy_confidence=0.42",
+     "signals": {"policy_confidence": 0.42}},
+    {"task_id": "task_7_ambiguous", "ad_key": "task_7_ambiguous",
+     "step_count": 3, "actions_already_taken": ["query_regulations", "check_advertiser_history"],
+     "setup_actions": _sa("query_regulations", "check_advertiser_history"),
+     "last_feedback": "risk_score=0.55",
+     "signals": {"policy_confidence": 0.42, "risk_score": 0.55}},
+    {"task_id": "task_7_ambiguous", "ad_key": "task_7_ambiguous",
+     "step_count": 4,
+     "actions_already_taken": ["query_regulations", "check_advertiser_history", "request_landing_page"],
+     "setup_actions": _sa("query_regulations", "check_advertiser_history", "request_landing_page"),
+     "last_feedback": "landing_suspicious",
+     "signals": {"policy_confidence": 0.42, "risk_score": 0.55, "landing_flag": True}},
+    {"task_id": "task_7_ambiguous", "ad_key": "task_7_ambiguous",
+     "step_count": 5,
+     "actions_already_taken": ["query_regulations", "check_advertiser_history", "request_landing_page", "submit_audit"],
+     "setup_actions": _sa("query_regulations", "check_advertiser_history", "request_landing_page", "submit_audit"),
+     "last_feedback": "audit_logged id=AUD-007",
+     "signals": {"policy_confidence": 0.42, "risk_score": 0.55, "landing_flag": True}},
+
+    # ---- TASK 8: Adversarial (fine print in image) ----
+    {"task_id": "task_8_adversarial", "ad_key": "task_8_adversarial",
+     "step_count": 1, "actions_already_taken": [], "setup_actions": [],
+     "last_feedback": "Ad loaded for task_8_adversarial. Begin with query_regulations.",
+     "signals": {}},
+    {"task_id": "task_8_adversarial", "ad_key": "task_8_adversarial",
+     "step_count": 2, "actions_already_taken": ["query_regulations"],
+     "setup_actions": _sa("query_regulations"),
+     "last_feedback": "policy_confidence=0.75",
+     "signals": {"policy_confidence": 0.75}},
+    {"task_id": "task_8_adversarial", "ad_key": "task_8_adversarial",
+     "step_count": 3, "actions_already_taken": ["query_regulations", "analyze_image"],
+     "setup_actions": _sa("query_regulations", "analyze_image"),
+     "last_feedback": "image_violation_detected",
+     "signals": {"policy_confidence": 0.75, "image_flag": True}},
+    {"task_id": "task_8_adversarial", "ad_key": "task_8_adversarial",
+     "step_count": 4,
+     "actions_already_taken": ["query_regulations", "analyze_image", "submit_audit"],
+     "setup_actions": _sa("query_regulations", "analyze_image", "submit_audit"),
+     "last_feedback": "audit_logged id=AUD-008",
+     "signals": {"policy_confidence": 0.75, "image_flag": True}},
+
+    # ---- TASK 9: Dependency Trap (text clean, image has violation) ----
+    {"task_id": "task_9_dependency_trap", "ad_key": "task_9_dependency_trap",
+     "step_count": 1, "actions_already_taken": [], "setup_actions": [],
+     "last_feedback": "Ad loaded for task_9_dependency_trap. Begin with query_regulations.",
+     "signals": {}},
+    {"task_id": "task_9_dependency_trap", "ad_key": "task_9_dependency_trap",
+     "step_count": 2, "actions_already_taken": ["query_regulations"],
+     "setup_actions": _sa("query_regulations"),
+     "last_feedback": "policy_confidence=0.50",
+     "signals": {"policy_confidence": 0.50}},
+    {"task_id": "task_9_dependency_trap", "ad_key": "task_9_dependency_trap",
+     "step_count": 3, "actions_already_taken": ["query_regulations", "analyze_image"],
+     "setup_actions": _sa("query_regulations", "analyze_image"),
+     "last_feedback": "image_violation_detected",
+     "signals": {"policy_confidence": 0.50, "image_flag": True}},
+    {"task_id": "task_9_dependency_trap", "ad_key": "task_9_dependency_trap",
+     "step_count": 4,
+     "actions_already_taken": ["query_regulations", "analyze_image", "submit_audit"],
+     "setup_actions": _sa("query_regulations", "analyze_image", "submit_audit"),
+     "last_feedback": "audit_logged id=AUD-009",
+     "signals": {"policy_confidence": 0.50, "image_flag": True}},
+
+    # ---- TASK 10: Failure Recovery ----
+    {"task_id": "task_10_failure", "ad_key": "task_10_failure",
+     "step_count": 1, "actions_already_taken": [], "setup_actions": [],
+     "last_feedback": "Ad loaded for task_10_failure. Begin with query_regulations.",
+     "signals": {}},
+    {"task_id": "task_10_failure", "ad_key": "task_10_failure",
+     "step_count": 2, "actions_already_taken": ["query_regulations"],
+     "setup_actions": _sa("query_regulations"),
+     "last_feedback": "policy_confidence=0.85",
+     "signals": {"policy_confidence": 0.85}},
+    {"task_id": "task_10_failure", "ad_key": "task_10_failure",
+     "step_count": 3, "actions_already_taken": ["query_regulations", "check_advertiser_history"],
+     "setup_actions": _sa("query_regulations", "check_advertiser_history"),
+     "last_feedback": "risk_score=0.80",
+     "signals": {"policy_confidence": 0.85, "risk_score": 0.80}},
+    {"task_id": "task_10_failure", "ad_key": "task_10_failure",
+     "step_count": 4,
+     "actions_already_taken": ["query_regulations", "check_advertiser_history", "submit_audit"],
+     "setup_actions": _sa("query_regulations", "check_advertiser_history", "submit_audit"),
+     "last_feedback": "audit_logged id=AUD-010",
+     "signals": {"policy_confidence": 0.85, "risk_score": 0.80}},
+]
 
-Scenario: {text}
-actions_already_taken: {actions_already_taken}
+# =========================
+# DATASET BUILDER
+# =========================
 
-Your next action?"""
+def build_observation(scenario):
+    """Construct observation JSON matching inference.py format."""
+    ad = TASK_AD_DATA[scenario["ad_key"]]
+    sigs = scenario.get("signals", {})
+    return {
+        "task_id": scenario["task_id"],
+        "last_feedback": scenario["last_feedback"],
+        "step_count": scenario["step_count"],
+        "actions_already_taken": scenario["actions_already_taken"],
+        "ad_details": {
+            **ad,
+            "status_message": scenario["last_feedback"],
+            "reward": 0.0,
+            "done": False,
+            "risk_score": sigs.get("risk_score"),
+            "policy_confidence": sigs.get("policy_confidence"),
+            "image_flag": sigs.get("image_flag"),
+            "landing_flag": sigs.get("landing_flag"),
+            "last_error": sigs.get("last_error"),
+        },
+    }
 
 
 def build_dataset():
     rows = []
     for s in BASE_SCENARIOS:
-        prompt = PROMPT_TEMPLATE.format(
-            text=s["text"],
-            actions_already_taken=json.dumps(s["actions_already_taken"]),
+        obs = build_observation(s)
+        user_content = (
+            "Current Ad Observation:\n"
+            + json.dumps(obs, indent=2)
+            + "\n\nWhat is your next action?"
         )
         rows.append({
-            "prompt": prompt,
+            "prompt": [
+                {"role": "system", "content": SYSTEM_PROMPT},
+                {"role": "user", "content": user_content},
+            ],
             "task_id": s["task_id"],
             "setup_actions": s["setup_actions"],
         })
-    return Dataset.from_list(rows * 10)  # 10 scenarios x 10 = 100 examples
+    return Dataset.from_list(rows * 8)
 
 # =========================
-# REWARD FUNCTION (FIXED)
+# REWARD FUNCTION
 # =========================
 
-_reward_call_count = [0]
-
 def reward_environment(prompts, completions, task_id=None, setup_actions=None, **kwargs):
-    """Shaped reward for GRPO."""
-    _reward_call_count[0] += 1
-    _call = _reward_call_count[0]
-    # #region agent log
-    _dlog("C", "grpo_train.py:reward_env", f"reward call #{_call}", {
-        "n_prompts": len(prompts) if prompts else 0,
-        "n_completions": len(completions) if completions else 0,
-        "completions_type": type(completions).__name__,
-        "first_completion_type": type(completions[0]).__name__ if completions else "N/A",
-        "first_completion_preview": str(completions[0])[:150] if completions else "N/A",
-        "task_id_is_none": task_id is None,
-        "setup_actions_is_none": setup_actions is None,
-        "kwargs_keys": list(kwargs.keys()),
-    })
-    # #endregion
-
+    """Shaped reward with phase-specific bonuses for meaningful GRPO gradients."""
     client = EnvClient(ENV_URL)
     rewards = []
 
     if task_id is None or setup_actions is None:
-        # #region agent log
-        _dlog("D", "grpo_train.py:reward_env", "task_id or setup_actions is None — returning -1 for all", {"call": _call})
-        # #endregion
         return [-1.0] * len(completions)
 
     for idx, (completion, t_id, setup) in enumerate(zip(completions, task_id, setup_actions)):
         parsed = extract_json(completion)
-        # #region agent log
-        if _call <= 3: _dlog("D", "grpo_train.py:reward_loop", f"call#{_call} item#{idx}", {"parsed_ok": parsed is not None, "action": parsed.get("action_type") if parsed else None, "raw_preview": str(completion)[:120], "task_id": t_id})
-        # #endregion
         if not parsed:
             rewards.append(-1.0)
             continue
@@ -310,6 +525,7 @@ def reward_environment(prompts, completions, task_id=None, setup_actions=None, *
         }
 
         try:
+            random.seed(hash((t_id, len(setup))) % (2**32 - 1))
             client.reset(t_id)
             for s in setup:
                 safe_step(client, s)
@@ -329,6 +545,35 @@ def reward_environment(prompts, completions, task_id=None, setup_actions=None, *
             else:
                 shaped = 0.5 + env_reward
 
+                taken = set(a["action_type"] for a in setup)
+
+                if not taken:
+                    if action_type == "query_regulations":
+                        shaped += 0.15
+                elif "submit_audit" in taken:
+                    if action_type in ("approve", "reject"):
+                        shaped += 0.2
+                    else:
+                        shaped -= 0.1
+                elif "query_regulations" in taken:
+                    gathering = {
+                        "analyze_image", "check_advertiser_history",
+                        "request_landing_page", "request_id_verification",
+                    }
+                    if action_type in gathering:
+                        shaped += 0.1
+                    elif action_type == "submit_audit":
+                        shaped += 0.1
+                    elif action_type in ("approve", "reject"):
+                        shaped -= 0.15
+
+                if t_id == "task_3_multimodal" and action_type == "analyze_image":
+                    shaped += 0.1
+                if t_id == "task_4_targeting" and action_type == "request_id_verification":
+                    shaped += 0.1
+                if t_id in ("task_8_adversarial", "task_9_dependency_trap") and action_type == "analyze_image":
+                    shaped += 0.1
+
             rewards.append(shaped)
 
         except Exception:
@@ -344,19 +589,15 @@ if torch.cuda.is_available():
     _props = torch.cuda.get_device_properties(0)
     _vram = _props.total_memory
     _name = _props.name
-    _cc = (_props.major, _props.minor)  # compute capability
+    _cc = (_props.major, _props.minor)
     print(f"GPU: {_name}  VRAM: {_vram / 1024**3:.1f} GB  Compute: {_cc[0]}.{_cc[1]}")
 else:
     _vram = 0
     _name = "CPU"
     _cc = (0, 0)
 
-USE_4BIT = _vram < 40 * 1024**3   # T4 (15 GB), L4 (24 GB) → 4-bit; A100 (80 GB) → full
-USE_BF16 = _cc >= (8, 0) and not USE_4BIT  # bf16 only when full-precision; 4-bit LoRA uses fp16 internally
-
-# #region agent log
-_dlog("A", "grpo_train.py:gpu_detect", "GPU config resolved", {"name":_name,"vram_gb":round(_vram/1024**3,1),"cc":list(_cc),"USE_4BIT":USE_4BIT,"USE_BF16":USE_BF16})
-# #endregion
+USE_4BIT = _vram < 40 * 1024**3
+USE_BF16 = _cc >= (8, 0) and not USE_4BIT
 
 model, tokenizer = FastLanguageModel.from_pretrained(
     model_name="unsloth/Llama-3.1-8B-Instruct",
@@ -385,24 +626,20 @@ model = FastLanguageModel.get_peft_model(
 
 dataset = build_dataset()
 
-# #region agent log
-_dlog("A", "grpo_train.py:trainer_init", "Creating GRPOTrainer", {"USE_4BIT":USE_4BIT,"USE_BF16":USE_BF16,"epochs":1 if USE_4BIT else 3,"batch":1 if USE_4BIT else 2,"gens":2 if USE_4BIT else 4,"dataset_len":len(dataset)})
-# #endregion
-
 trainer = GRPOTrainer(
     model=model,
     reward_funcs=[reward_environment],
     args=GRPOConfig(
         output_dir="outputs",
-        learning_rate=2e-5,
-        num_train_epochs=1 if USE_4BIT else 3,
+        learning_rate=5e-6,
+        num_train_epochs=1 if USE_4BIT else 2,
         per_device_train_batch_size=1 if USE_4BIT else 2,
-        gradient_accumulation_steps=2 if USE_4BIT else 4,
-        num_generations=2 if USE_4BIT else 4,
-        max_prompt_length=768,
-        max_completion_length=128,
-        logging_steps=3 if USE_4BIT else 5,
-        warmup_steps=5 if USE_4BIT else 10,
+        gradient_accumulation_steps=4,
+        num_generations=4,
+        max_prompt_length=512,
+        max_completion_length=80,
+        logging_steps=5,
+        warmup_steps=10,
         bf16=USE_BF16,
         fp16=not USE_BF16,
         report_to="none",
@@ -418,9 +655,6 @@ trainer = GRPOTrainer(
 if __name__ == "__main__":
     ensure_env_ready()
 
-    # #region agent log
-    _dlog("E", "grpo_train.py:train_start", "About to call trainer.train()", {"gpu_mem_allocated_gb": round(torch.cuda.memory_allocated()/1024**3, 2) if torch.cuda.is_available() else 0})
-    # #endregion
     print("Starting GRPO training...")
     trainer.train()
 
@@ -428,7 +662,7 @@ if __name__ == "__main__":
     tokenizer.save_pretrained("outputs/lora_adapter")
     print("LoRA adapter saved to outputs/lora_adapter")
 
-    print("Merging adapter into base model (bf16)...")
+    print("Merging adapter into base model...")
     merged_model, merged_tokenizer = FastLanguageModel.from_pretrained(
         model_name="outputs/lora_adapter",
         load_in_4bit=False,