Update grpo_train.py

grpo_train.py

@@ -2,6 +2,7 @@
 
 import os
 
+# Route all caches to /tmp/ to avoid Hugging Face Spaces Read-Only errors
 os.environ.setdefault("USER", "user")
 os.environ.setdefault("HOME", "/tmp/home")
 os.environ.setdefault("HF_HOME", "/tmp/hf_home")
@@ -22,8 +23,6 @@ import time
 import json
 import random
 import requests
-from requests.adapters import HTTPAdapter
-from urllib3.util.retry import Retry
 import torch
 
 from datasets import Dataset
@@ -37,15 +36,19 @@ PatchFastRL("GRPO", FastLanguageModel)
 # =========================
 
 OUTPUT_DIR = "/tmp/outputs"
-
 ENV_URL = os.getenv("ENV_URL", "https://parth-1-metaguard.hf.space")
 HF_TOKEN = os.getenv("HF_TOKEN", "")
-HF_REPO = os.getenv("HF_REPO", "")
+HF_REPO = os.getenv("HF_REPO", "")  # e.g. "parth-1/metaguard-llama3.1-8b-grpo"
 
 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",
 ]
 
 # =========================
@@ -53,57 +56,54 @@ ALLOWED_ACTIONS = [
 # =========================
 
 def ensure_env_ready():
-    for i in range(60):
+    for _ in range(20):
         try:
             r = requests.post(
                 f"{ENV_URL}/reset",
                 json={"task_id": "task_1_healthcare"},
-                timeout=15,
+                timeout=5
             )
             if r.status_code == 200:
-                print(f"Environment ready (attempt {i+1})")
+                print("✅ Environment ready")
                 return
-            print(f"Env returned {r.status_code}, retrying...")
-        except Exception as e:
-            if i % 10 == 0:
-                print(f"Env not ready (attempt {i+1}): {e}")
-        time.sleep(2)
-    raise RuntimeError("ENV not reachable after 60 attempts")
+        except:
+            pass
+        time.sleep(1)
+    raise RuntimeError("❌ ENV not reachable")
 
 # =========================
-# ENV CLIENT
+# SAFE CLIENT
 # =========================
 
-def _build_session():
-    s = requests.Session()
-    retries = Retry(total=3, backoff_factor=0.5,
-                    status_forcelist=[500, 502, 503, 504])
-    s.mount("https://", HTTPAdapter(max_retries=retries))
-    s.mount("http://", HTTPAdapter(max_retries=retries))
-    return s
-
 class EnvClient:
     def __init__(self, url):
         self.url = url
-        self.session = _build_session()
 
     def reset(self, task_id):
-        return self.session.post(
-            f"{self.url}/reset", json={"task_id": task_id}, timeout=15,
+        return requests.post(
+            f"{self.url}/reset",
+            json={"task_id": task_id},
+            timeout=8
         ).json()
 
     def step(self, action):
-        return self.session.post(
-            f"{self.url}/step", json={"action": action}, timeout=15,
+        return requests.post(
+            f"{self.url}/step",
+            json={"action": action},
+            timeout=8
         ).json()
 
 def safe_step(client, action):
-    for attempt in range(5):
+    for _ in range(3):
         try:
             return client.step(action)
-        except Exception:
-            time.sleep(1.0 * (attempt + 1))
-    return {"reward": -0.3, "status_message": "timeout"}
+        except:
+            time.sleep(0.5)
+    return {"reward": -0.3}
+
+# =========================
+# JSON PARSER
+# =========================
 
 def extract_json(text):
     try:
@@ -112,435 +112,146 @@ def extract_json(text):
             if text.startswith("json"):
                 text = text[4:]
         return json.loads(text.strip())
-    except Exception:
+    except:
         return None
 
 # =========================
-# SYSTEM PROMPT (identical to inference.py)
+# DATASET
 # =========================
 
-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_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_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",
+BASE_SCENARIOS = [
+    {
+        "task_id": "task_1_healthcare",
+        "text": "Healthcare ad: 'miracle supplement cures disease'. No actions taken yet.",
+        "actions_already_taken": [],
+        "setup_actions": [],
     },
-    "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_2_financial",
+        "text": "Financial ad: 'guaranteed 500% returns, zero risk'. No actions taken yet.",
+        "actions_already_taken": [],
+        "setup_actions": [],
     },
-    "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",
+    {
+        "task_id": "task_3_multimodal",
+        "text": "Multimodal ad: image may contain hidden violation. No actions taken yet.",
+        "actions_already_taken": [],
+        "setup_actions": [],
     },
-    "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",
+    {
+        "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_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",
+    {
+        "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_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",
+    {
+        "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_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",
+    {
+        "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_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",
+    {
+        "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_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_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_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",
+    {
+        "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"},
+        ],
     },
-}
+]
 
-# =========================
-# SCENARIOS (43 unique across all tasks + phases)
-# =========================
+PROMPT_TEMPLATE = """You are an enterprise Ad Policy Compliance Agent.
 
-_sa = lambda *acts: [{"action_type": a, "reasoning": "setup"} for a in acts]
+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
 
-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}},
-]
+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
 
-# =========================
-# DATASET BUILDER
-# =========================
+HARD RULES:
+- NEVER repeat an action listed in `actions_already_taken`.
+- Respond with ONLY a valid JSON object. No markdown, no prose.
+
+Required format:
+{{"action_type": "<one_of_the_actions_above>", "reasoning": "<short reason>"}}
 
-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"),
-        },
-    }
+Scenario: {text}
+actions_already_taken: {actions_already_taken}
 
+Your next action?"""
 
 def build_dataset():
     rows = []
     for s in BASE_SCENARIOS:
-        obs = build_observation(s)
-        user_content = (
-            "Current Ad Observation:\n"
-            + json.dumps(obs, indent=2)
-            + "\n\nWhat is your next action?"
+        prompt = PROMPT_TEMPLATE.format(
+            text=s["text"],
+            actions_already_taken=json.dumps(s["actions_already_taken"]),
         )
         rows.append({
-            "prompt": [
-                {"role": "system", "content": SYSTEM_PROMPT},
-                {"role": "user", "content": user_content},
-            ],
+            "prompt": prompt,
             "task_id": s["task_id"],
             "setup_actions": s["setup_actions"],
         })
-    return Dataset.from_list(rows * 8)
+    return Dataset.from_list(rows * 10)  # 10 scenarios x 10 = 100 examples
 
 # =========================
 # REWARD FUNCTION
 # =========================
 
 def reward_environment(prompts, completions, task_id=None, setup_actions=None, **kwargs):
-    """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:
-        return [-1.0] * len(completions)
-
-    for idx, (completion, t_id, setup) in enumerate(zip(completions, task_id, setup_actions)):
+    for completion, t_id, setup in zip(completions, task_id, setup_actions):
         parsed = extract_json(completion)
         if not parsed:
             rewards.append(-1.0)
@@ -557,7 +268,6 @@ 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)
@@ -577,35 +287,6 @@ 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:
@@ -617,22 +298,11 @@ def reward_environment(prompts, completions, task_id=None, setup_actions=None, *
 # MODEL
 # =========================
 
-if torch.cuda.is_available():
-    _props = torch.cuda.get_device_properties(0)
-    _vram = _props.total_memory
-    _name = _props.name
-    _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)
-
 model, tokenizer = FastLanguageModel.from_pretrained(
-    model_name="unsloth/Llama-3.2-3B-Instruct-bnb-4bit",
-    load_in_4bit=True,
+    model_name="unsloth/Llama-3.1-8B-Instruct",
+    load_in_4bit=True, # Strictly True for L4 24GB
     max_seq_length=2048,
-    dtype=None,
+    dtype=torch.bfloat16, # L4 Native Support
 )
 
 model = FastLanguageModel.get_peft_model(
@@ -660,15 +330,15 @@ trainer = GRPOTrainer(
     reward_funcs=[reward_environment],
     args=GRPOConfig(
         output_dir=OUTPUT_DIR,
-        learning_rate=5e-6,
+        learning_rate=2e-5,
         num_train_epochs=3,
-        per_device_train_batch_size=4,
-        gradient_accumulation_steps=4,
-        num_generations=8,
-        max_prompt_length=512,
-        max_completion_length=80,
+        per_device_train_batch_size=1, # Memory safe for L4
+        gradient_accumulation_steps=8, # Maintain effective batch size of 8
+        num_generations=2, # Memory safe generation limit
+        max_prompt_length=768,
+        max_completion_length=128,
         logging_steps=5,
-        warmup_steps=10,
+        warmup_ratio=0.1,
         bf16=True,
         fp16=False,
         report_to="none",
@@ -691,17 +361,17 @@ if __name__ == "__main__":
     try:
         trainer.train()
     except torch.cuda.OutOfMemoryError:
-        print("OOM detected! Clearing cache and retrying with batch_size=2, num_generations=4...")
+        print("OOM detected! Clearing cache and severely restricting memory...")
         torch.cuda.empty_cache()
-        trainer.args.per_device_train_batch_size = 2
-        trainer.args.num_generations = 4
+        trainer.args.per_device_train_batch_size = 1
+        trainer.args.gradient_accumulation_steps = 16
         trainer.train()
 
     model.save_pretrained(LORA_DIR)
     tokenizer.save_pretrained(LORA_DIR)
     print(f"LoRA adapter saved to {LORA_DIR}")
 
-    print("Merging adapter into base model...")
+    print("Merging adapter into base model (bf16)...")
     merged_model, merged_tokenizer = FastLanguageModel.from_pretrained(
         model_name=LORA_DIR,
         load_in_4bit=False,