Add W&B experiment tracking and structured logging

README.md

@@ -24,6 +24,7 @@ short_description: Meta pytorch hugging face hackathon
 | **HF Space (Live Environment)** | [godreign-policy2logic.hf.space](https://godreign-policy2logic.hf.space) |
 | **Training Notebook (Colab)** | [Open in Colab](https://colab.research.google.com/github/GodreignElgin/policy2logic/blob/main/training/colab_training.ipynb) |
 | **Writeup / Slides** | *TBD — add your link here* |
+| **Experiment Tracking (W&B)** | [Wandb Project](https://wandb.ai/YOUR_USERNAME/policy-to-logic-rl) |
 
 ---
 
@@ -44,6 +45,15 @@ improving agent behavior across episodes without weight updates.
 
 ---
 
+## 📈 Experiment Tracking
+
+All training runs are logged to Weights & Biases.
+Metrics tracked per episode: total reward, final accuracy, steps used, success rate, few-shot examples used.
+
+Live dashboard: [wandb.ai/YOUR_USERNAME/policy-to-logic-rl](https://wandb.ai/YOUR_USERNAME/policy-to-logic-rl)
+
+---
+
 ## 🧠 What This Is
 
 This project builds a **verifiable RL environment** where:

policy_to_logic_env/server/requirements.txt

@@ -3,3 +3,4 @@ pydantic>=2.0
 fastapi>=0.104.0
 uvicorn>=0.24.0
 requests>=2.25.0
+wandb>=0.16.0

pyproject.toml

@@ -13,6 +13,7 @@ dependencies = [
     "huggingface-hub>=1.12.0",
     "matplotlib>=3.7.0",
     "numpy>=1.24.0",
+    "wandb>=0.16.0",
 ]
 
 [project.optional-dependencies]

training/colab_training.ipynb

@@ -19,7 +19,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "# Policy-to-Logic RL Environment — Training Notebook\n",
+    "# Policy-to-Logic RL Environment \u2014 Training Notebook\n",
     "\n",
     "This notebook runs the **reward-guided trajectory optimization loop** against the deployed environment.\n",
     "\n",
@@ -27,7 +27,8 @@
     "1. Connects to the live HF Spaces environment\n",
     "2. Runs 8 episodes per task (3 tasks = 24 total episodes)\n",
     "3. Accumulates high-reward trajectories as few-shot examples\n",
-    "4. Generates training evidence plots (reward curve, accuracy curve, improvement chart)"
+    "4. Generates training evidence plots (reward curve, accuracy curve, improvement chart)\n",
+    "5. Logs everything to Weights & Biases"
    ]
   },
   {
@@ -37,7 +38,7 @@
    "outputs": [],
    "source": [
     "# Cell 1: Install dependencies\n",
-    "!pip install openai requests matplotlib numpy"
+    "!pip install openai requests matplotlib numpy wandb"
    ]
   },
   {
@@ -52,12 +53,16 @@
     "# SET THESE BEFORE RUNNING\n",
     "HF_TOKEN = \"\"  # Your Hugging Face token with inference access\n",
     "ENV_URL = \"https://godreign-policy2logic.hf.space\"  # Your deployed environment URL\n",
+    "WANDB_API_KEY = \"\" # Your Wandb API key\n",
     "\n",
     "os.environ[\"HF_TOKEN\"] = HF_TOKEN\n",
     "os.environ[\"ENV_BASE_URL\"] = ENV_URL\n",
+    "if WANDB_API_KEY:\n",
+    "    os.environ[\"WANDB_API_KEY\"] = WANDB_API_KEY\n",
     "\n",
     "print(f\"Environment URL: {ENV_URL}\")\n",
-    "print(f\"HF Token set: {'Yes' if HF_TOKEN else 'NO - MUST SET THIS'}\")"
+    "print(f\"HF Token set: {'Yes' if HF_TOKEN else 'NO - MUST SET THIS'}\")\n",
+    "print(f\"Wandb Token set: {'Yes' if WANDB_API_KEY else 'NO - WILL PROMPT'}\")"
    ]
   },
   {
@@ -84,18 +89,40 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Cell 4: Training loop implementation\n",
-    "# Full contents of training/trajectory_optimizer.py\n",
+    "# Cell 4: Wandb login \u2014 run this before training\n",
+    "import wandb\n",
+    "wandb.login()  # Will prompt for API key if WANDB_API_KEY is not set"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Cell 5: Training loop implementation (full trajectory_optimizer.py)\n",
+    "\"\"\"\n",
+    "Reward-Guided Trajectory Optimization Loop\n",
+    "==========================================\n",
+    "Optimizes agent behavior across episodes by accumulating high-reward\n",
+    "trajectories as few-shot examples. Uses environment reward signal to\n",
+    "drive improvement \u00e2\u20ac\u201d no weight updates required.\n",
+    "\n",
+    "This implements a policy improvement loop where:\n",
+    "  - reward_signal \u00e2\u2020\u2019 trajectory_selection \u00e2\u2020\u2019 context_construction \u00e2\u2020\u2019 improved_policy\n",
+    "\"\"\"\n",
     "\n",
     "import json\n",
     "import os\n",
     "import time\n",
     "import requests\n",
+    "import logging\n",
+    "import wandb\n",
     "from dataclasses import dataclass, field\n",
     "from typing import Optional\n",
     "from openai import OpenAI\n",
     "\n",
-    "# ── Configuration ────────────────────────────────────────────────────────────\n",
+    "# \u00e2\u201d\u20ac\u00e2\u201d\u20ac Configuration \u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\n",
     "\n",
     "ENV_BASE_URL = os.getenv(\"ENV_BASE_URL\", \"http://localhost:7860\")\n",
     "HF_TOKEN = os.getenv(\"HF_TOKEN\", \"\")\n",
@@ -103,11 +130,14 @@
     "TEMPERATURE = 0.3\n",
     "MAX_TOKENS = 1024\n",
     "\n",
-    "NUM_EPISODES_PER_TASK = 8\n",
-    "TOP_K_TRAJECTORIES = 3\n",
-    "MIN_REWARD_THRESHOLD = 0.3\n",
+    "# Training hyperparameters\n",
+    "NUM_EPISODES_PER_TASK = 8        # Episodes to run per task\n",
+    "TOP_K_TRAJECTORIES = 3           # Max few-shot examples to keep\n",
+    "MIN_REWARD_THRESHOLD = 0.3       # Minimum reward to store trajectory\n",
     "TASKS = [\"data_access\", \"resource_access\", \"transaction_approval\"]\n",
     "\n",
+    "# \u00e2\u201d\u20ac\u00e2\u201d\u20ac Data Structures \u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\n",
+    "\n",
     "@dataclass\n",
     "class Step:\n",
     "    step_number: int\n",
@@ -128,7 +158,10 @@
     "    success: bool = False\n",
     "\n",
     "    def to_few_shot_string(self) -> str:\n",
-    "        lines = [f\"=== Example Episode (reward={self.total_reward:.2f}, accuracy={self.final_accuracy:.2f}) ===\"]\n",
+    "        \"\"\"Convert trajectory to a few-shot example string for prompting.\"\"\"\n",
+    "        lines = [\n",
+    "            f\"=== Example Episode (reward={self.total_reward:.2f}, accuracy={self.final_accuracy:.2f}) ===\",\n",
+    "        ]\n",
     "        for s in self.steps:\n",
     "            lines.append(f\"Step {s.step_number}: action={s.action_type}\")\n",
     "            lines.append(f\"  Content: {s.action_content[:200]}\")\n",
@@ -137,6 +170,8 @@
     "                lines.append(f\"  Feedback: {s.feedback[:150]}\")\n",
     "        return \"\\n\".join(lines)\n",
     "\n",
+    "# \u00e2\u201d\u20ac\u00e2\u201d\u20ac Environment Client \u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\n",
+    "\n",
     "class EnvClient:\n",
     "    def __init__(self, base_url: str):\n",
     "        self.base_url = base_url.rstrip(\"/\")\n",
@@ -148,7 +183,10 @@
     "        return r.json()\n",
     "\n",
     "    def step(self, action_type: str, content: str) -> dict:\n",
-    "        r = self.session.post(f\"{self.base_url}/step\", json={\"action_type\": action_type, \"content\": content})\n",
+    "        r = self.session.post(f\"{self.base_url}/step\", json={\n",
+    "            \"action_type\": action_type,\n",
+    "            \"content\": content\n",
+    "        })\n",
     "        r.raise_for_status()\n",
     "        return r.json()\n",
     "\n",
@@ -159,18 +197,40 @@
     "        except Exception:\n",
     "            return False\n",
     "\n",
+    "# \u00e2\u201d\u20ac\u00e2\u201d\u20ac LLM Agent \u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\n",
+    "\n",
     "class Agent:\n",
     "    def __init__(self, hf_token: str):\n",
-    "        self.client = OpenAI(base_url=\"https://router.huggingface.co/v1\", api_key=hf_token)\n",
-    "\n",
-    "    def get_action(self, observation, step_number, episode_history, few_shot_examples):\n",
-    "        system_prompt = self._build_system_prompt(few_shot_examples)\n",
+    "        self.client = OpenAI(\n",
+    "            base_url=\"https://router.huggingface.co/v1\",\n",
+    "            api_key=hf_token\n",
+    "        )\n",
+    "\n",
+    "    def get_action(\n",
+    "        self,\n",
+    "        observation: dict,\n",
+    "        step_number: int,\n",
+    "        episode_history: list[str],\n",
+    "        few_shot_examples: list[Trajectory],\n",
+    "        task_name: str = \"\"\n",
+    "    ) -> tuple[str, str]:\n",
+    "        \"\"\"\n",
+    "        Returns (action_type, content_json_string).\n",
+    "        action_type: one of ask_clarification | propose_rules | refine_rules\n",
+    "        content: JSON string appropriate for that action\n",
+    "        \"\"\"\n",
+    "        system_prompt = self._build_system_prompt(few_shot_examples, task_name)\n",
     "        user_prompt = self._build_user_prompt(observation, step_number, episode_history)\n",
+    "\n",
     "        try:\n",
     "            response = self.client.chat.completions.create(\n",
     "                model=MODEL,\n",
-    "                messages=[{\"role\": \"system\", \"content\": system_prompt}, {\"role\": \"user\", \"content\": user_prompt}],\n",
-    "                temperature=TEMPERATURE, max_tokens=MAX_TOKENS\n",
+    "                messages=[\n",
+    "                    {\"role\": \"system\", \"content\": system_prompt},\n",
+    "                    {\"role\": \"user\", \"content\": user_prompt}\n",
+    "                ],\n",
+    "                temperature=TEMPERATURE,\n",
+    "                max_tokens=MAX_TOKENS\n",
     "            )\n",
     "            raw = response.choices[0].message.content.strip()\n",
     "            return self._parse_response(raw, observation)\n",
@@ -178,155 +238,489 @@
     "            print(f\"    [LLM ERROR] {e}\")\n",
     "            return \"propose_rules\", json.dumps({\"rules\": [], \"default\": \"DENY\"})\n",
     "\n",
-    "    def _build_system_prompt(self, few_shot_examples):\n",
-    "        base = \"\"\"You are a policy-to-logic agent. Convert natural language policies into executable rules.\n",
+    "    def _build_system_prompt(self, few_shot_examples: list[Trajectory], task_name: str = \"\") -> str:\n",
+    "        base = \"\"\"You are a policy-to-logic agent. Your job is to convert natural language policies into executable rules.\n",
     "\n",
     "AVAILABLE ACTIONS:\n",
     "1. ask_clarification: {\"type\": \"clarification\", \"question\": \"your question\"}\n",
     "2. propose_rules: {\"rules\": [...], \"default\": \"DECISION\"}\n",
     "3. refine_rules: {\"rules\": [...], \"default\": \"DECISION\"}\n",
     "\n",
-    "DSL FORMAT: {\"rules\": [{\"if\": [{\"field\": \"NAME\", \"op\": \"OP\", \"value\": VAL}], \"then\": \"DECISION\"}], \"default\": \"FALLBACK\"}\n",
-    "Operators: >, <, >=, <=, ==, !=. Rules execute top-to-bottom, first match wins.\n",
+    "DSL FORMAT for rules:\n",
+    "{\n",
+    "  \"rules\": [\n",
+    "    {\n",
+    "      \"if\": [\n",
+    "        {\"field\": \"FIELD_NAME\", \"op\": \"OPERATOR\", \"value\": VALUE}\n",
+    "      ],\n",
+    "      \"then\": \"DECISION\"\n",
+    "    }\n",
+    "  ],\n",
+    "  \"default\": \"FALLBACK_DECISION\"\n",
+    "}\n",
+    "\n",
+    "Operators: >, <, >=, <=, ==, !=\n",
+    "Rules execute top-to-bottom. First match wins. Default applies if no rule matches.\n",
+    "\n",
+    "STRATEGY:\n",
+    "- Step 1: Ask 1-2 targeted clarification questions about ambiguous terms\n",
+    "- Step 2: Propose initial rules based on policy + clarifications  \n",
+    "- Step 3+: Refine rules based on failure feedback\n",
+    "\n",
+    "OUTPUT FORMAT: Respond ONLY with valid JSON. No markdown. No explanation.\n",
+    "{\"action_type\": \"propose_rules\", \"content\": \"{...escaped json string...}\"}\n",
+    "\"\"\"\n",
+    "        # Task-specific guidance for complex tasks\n",
+    "        if task_name == \"transaction_approval\":\n",
+    "            base += \"\"\"\n",
+    "IMPORTANT \u00e2\u20ac\u201d TRANSACTION APPROVAL TASK:\n",
+    "This task has 4 possible decisions: APPROVE, REQUIRE_APPROVAL, COMPLIANCE_REVIEW, HOLD\n",
+    "Rules are evaluated TOP-TO-BOTTOM. Order matters critically. You MUST order rules by priority:\n",
+    "  1. FIRST: Check if transfer_type == \"international\" \u00e2\u2020\u2019 then COMPLIANCE_REVIEW (always, overrides everything)\n",
+    "  2. SECOND: Check if amount >= 10000 AND time is outside business hours (hour < 9 or hour >= 17) \u00e2\u2020\u2019 then HOLD\n",
+    "  3. THIRD: Check if amount > 5000 AND initiator_role != \"manager\" \u00e2\u2020\u2019 then REQUIRE_APPROVAL\n",
+    "  4. DEFAULT: APPROVE\n",
+    "\n",
+    "Key details:\n",
+    "- Standard limit is $5,000 (amount > 5000 triggers approval, NOT >=)\n",
+    "- High-value threshold is $10,000 (amount >= 10000)\n",
+    "- Business hours: hour >= 9 AND hour < 17\n",
+    "- Manager exemption ONLY applies to the standard $5,000 limit, NOT to international or high-value HOLD rules\n",
+    "- \"system\" role follows the same rules as \"employee\"\n",
+    "\n",
+    "Here is a working example of valid rules for this task:\n",
+    "{\"rules\": [{\"if\": [{\"field\": \"transfer_type\", \"op\": \"==\", \"value\": \"international\"}], \"then\": \"COMPLIANCE_REVIEW\"}, {\"if\": [{\"field\": \"amount\", \"op\": \">=\", \"value\": 10000}, {\"field\": \"time\", \"op\": \">=\", \"value\": 17}], \"then\": \"HOLD\"}, {\"if\": [{\"field\": \"amount\", \"op\": \">=\", \"value\": 10000}, {\"field\": \"time\", \"op\": \"<\", \"value\": 9}], \"then\": \"HOLD\"}, {\"if\": [{\"field\": \"amount\", \"op\": \">\", \"value\": 5000}, {\"field\": \"initiator_role\", \"op\": \"!=\", \"value\": \"manager\"}], \"then\": \"REQUIRE_APPROVAL\"}], \"default\": \"APPROVE\"}\n",
+    "\"\"\"\n",
+    "        elif task_name == \"resource_access\":\n",
+    "            base += \"\"\"\n",
+    "IMPORTANT \u00e2\u20ac\u201d RESOURCE ACCESS TASK:\n",
+    "This task has roles: junior, senior, contractor. Document types: public, internal, confidential.\n",
+    "- Senior employees: ALLOW everything always\n",
+    "- Contractors: ALLOW only public, DENY everything else\n",
+    "- Junior + confidential: ALWAYS DENY (regardless of time \u00e2\u20ac\u201d the policy is misleading about this)\n",
+    "- Junior + internal: ALLOW only during business hours (hour >= 8 AND hour < 17)\n",
+    "- Junior + public: ALLOW always\n",
+    "- Business hours: hour >= 8 AND hour < 17\n",
+    "\"\"\"\n",
     "\n",
-    "STRATEGY: Ask 1-2 clarifications first, then propose rules, then refine based on failures.\n",
-    "OUTPUT: Respond ONLY with valid JSON: {\"action_type\": \"...\", \"content\": \"...\"}\"\"\"\n",
     "        if few_shot_examples:\n",
-    "            base += \"\\n\\nLEARNED FROM PREVIOUS EPISODES:\\n\"\n",
+    "            base += \"\\n\\nLEARNED FROM PREVIOUS EPISODES (high-reward strategies):\\n\"\n",
     "            for traj in few_shot_examples[-TOP_K_TRAJECTORIES:]:\n",
     "                base += \"\\n\" + traj.to_few_shot_string() + \"\\n\"\n",
     "        return base\n",
     "\n",
-    "    def _build_user_prompt(self, obs, step, history):\n",
-    "        lines = [f\"TASK: {obs.get('task_name', 'unknown')}\", f\"STEP: {step} of {obs.get('max_steps', 7)}\", f\"\\nPOLICY:\\n{obs.get('policy_text', '')}\"]\n",
-    "        if obs.get(\"clarification_response\"): lines.append(f\"\\nLAST CLARIFICATION:\\n{obs['clarification_response']}\")\n",
+    "    def _build_user_prompt(self, obs: dict, step: int, history: list[str]) -> str:\n",
+    "        lines = [\n",
+    "            f\"TASK: {obs.get('task_name', 'unknown')}\",\n",
+    "            f\"STEP: {step} of {obs.get('max_steps', 7)}\",\n",
+    "            f\"\\nPOLICY:\\n{obs.get('policy_text', '')}\",\n",
+    "        ]\n",
+    "        if obs.get(\"clarification_response\"):\n",
+    "            lines.append(f\"\\nLAST CLARIFICATION ANSWER:\\n{obs['clarification_response']}\")\n",
     "        if obs.get(\"test_results\"):\n",
     "            tr = obs[\"test_results\"]\n",
-    "            lines.append(f\"\\nTEST RESULTS: {tr.get('passed',0)}/{tr.get('total',0)} (acc={obs.get('current_accuracy',0):.2f})\")\n",
-    "            if tr.get(\"sample_failures\"): lines.extend([f\"  - {f}\" for f in tr[\"sample_failures\"][:3]])\n",
-    "        if obs.get(\"feedback\"): lines.append(f\"\\nFEEDBACK: {obs['feedback']}\")\n",
-    "        if history: lines.append(f\"\\nHISTORY:\\n\" + \"\\n\".join(history[-3:]))\n",
-    "        lines.append(f\"\\nAVAILABLE: {obs.get('available_actions', [])}\")\n",
-    "        lines.append(\"\\nRespond with JSON only.\")\n",
+    "            lines.append(f\"\\nTEST RESULTS: {tr.get('passed', 0)}/{tr.get('total', 0)} passed (accuracy={obs.get('current_accuracy', 0):.2f})\")\n",
+    "            if tr.get(\"sample_failures\"):\n",
+    "                lines.append(\"SAMPLE FAILURES:\")\n",
+    "                for f in tr[\"sample_failures\"][:3]:\n",
+    "                    lines.append(f\"  - {f}\")\n",
+    "        if obs.get(\"feedback\"):\n",
+    "            lines.append(f\"\\nFEEDBACK: {obs['feedback']}\")\n",
+    "        if history:\n",
+    "            lines.append(f\"\\nACTION HISTORY (last 3):\\n\" + \"\\n\".join(history[-3:]))\n",
+    "        lines.append(f\"\\nAVAILABLE ACTIONS: {obs.get('available_actions', [])}\")\n",
+    "        lines.append(\"\\nRespond with JSON only: {\\\"action_type\\\": \\\"...\\\", \\\"content\\\": \\\"...\\\"}\")\n",
     "        return \"\\n\".join(lines)\n",
     "\n",
-    "    def _parse_response(self, raw, obs):\n",
+    "    def _parse_response(self, raw: str, obs: dict) -> tuple[str, str]:\n",
+    "        # Strip markdown code fences if present\n",
     "        if \"```\" in raw:\n",
     "            raw = raw.split(\"```\")[1]\n",
-    "            if raw.startswith(\"json\"): raw = raw[4:]\n",
+    "            if raw.startswith(\"json\"):\n",
+    "                raw = raw[4:]\n",
     "        raw = raw.strip()\n",
+    "\n",
     "        try:\n",
     "            parsed = json.loads(raw)\n",
     "            action_type = parsed.get(\"action_type\", \"propose_rules\")\n",
     "            content = parsed.get(\"content\", \"{}\")\n",
-    "            valid = obs.get(\"available_actions\", [\"propose_rules\"])\n",
-    "            if action_type not in valid: action_type = valid[0]\n",
-    "            if isinstance(content, dict): content = json.dumps(content)\n",
+    "\n",
+    "            # Validate action_type\n",
+    "            valid_actions = obs.get(\"available_actions\", [\"propose_rules\", \"ask_clarification\"])\n",
+    "            if action_type not in valid_actions:\n",
+    "                action_type = \"propose_rules\" if \"propose_rules\" in valid_actions else valid_actions[0]\n",
+    "\n",
+    "            # Ensure content is a string\n",
+    "            if isinstance(content, dict):\n",
+    "                content = json.dumps(content)\n",
     "            return action_type, content\n",
-    "        except: return \"propose_rules\", json.dumps({\"rules\": [], \"default\": \"DENY\"})\n",
+    "        except Exception:\n",
+    "            return \"propose_rules\", json.dumps({\"rules\": [], \"default\": \"DENY\"})\n",
+    "\n",
+    "# \u00e2\u201d\u20ac\u00e2\u201d\u20ac Trajectory Bank \u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\n",
     "\n",
     "class TrajectoryBank:\n",
-    "    def __init__(self): self.bank = {task: [] for task in TASKS}\n",
-    "    def store(self, t):\n",
-    "        if t.total_reward >= MIN_REWARD_THRESHOLD:\n",
-    "            self.bank[t.task_name].append(t)\n",
-    "            self.bank[t.task_name].sort(key=lambda x: x.total_reward, reverse=True)\n",
-    "            self.bank[t.task_name] = self.bank[t.task_name][:TOP_K_TRAJECTORIES]\n",
-    "    def get_examples(self, task): return self.bank.get(task, [])\n",
-    "    def summary(self): return {t: {\"stored\": len(v), \"best_reward\": max((x.total_reward for x in v), default=0)} for t,v in self.bank.items()}\n",
+    "    \"\"\"Stores and retrieves high-reward trajectories per task.\"\"\"\n",
+    "\n",
+    "    def __init__(self):\n",
+    "        self.bank: dict[str, list[Trajectory]] = {task: [] for task in TASKS}\n",
+    "\n",
+    "    def store(self, trajectory: Trajectory):\n",
+    "        if trajectory.total_reward >= MIN_REWARD_THRESHOLD:\n",
+    "            self.bank[trajectory.task_name].append(trajectory)\n",
+    "            # Keep only top-K by reward\n",
+    "            self.bank[trajectory.task_name].sort(key=lambda t: t.total_reward, reverse=True)\n",
+    "            self.bank[trajectory.task_name] = self.bank[trajectory.task_name][:TOP_K_TRAJECTORIES]\n",
+    "\n",
+    "    def get_examples(self, task_name: str) -> list[Trajectory]:\n",
+    "        return self.bank.get(task_name, [])\n",
+    "\n",
+    "    def summary(self) -> dict:\n",
+    "        return {\n",
+    "            task: {\n",
+    "                \"stored\": len(trajs),\n",
+    "                \"best_reward\": max((t.total_reward for t in trajs), default=0),\n",
+    "                \"best_accuracy\": max((t.final_accuracy for t in trajs), default=0)\n",
+    "            }\n",
+    "            for task, trajs in self.bank.items()\n",
+    "        }\n",
+    "\n",
+    "# \u00e2\u201d\u20ac\u00e2\u201d\u20ac Training Loop \u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\n",
     "\n",
     "class TrainingLoop:\n",
-    "    def __init__(self, env_url, hf_token):\n",
+    "    def __init__(self, env_url: str, hf_token: str):\n",
     "        self.env = EnvClient(env_url)\n",
     "        self.agent = Agent(hf_token)\n",
     "        self.bank = TrajectoryBank()\n",
-    "        self.metrics = []\n",
-    "\n",
-    "    def run_episode(self, task_name, episode_id):\n",
+    "        self.metrics = []  # List of {episode, task, reward, accuracy, success}\n",
+    "\n",
+    "        os.makedirs(\"training/logs\", exist_ok=True)\n",
+    "        log_filename = f\"training/logs/run_{int(time.time())}.log\"\n",
+    "        logging.basicConfig(\n",
+    "            level=logging.INFO,\n",
+    "            format=\"%(asctime)s [%(levelname)s] %(message)s\",\n",
+    "            handlers=[\n",
+    "                logging.FileHandler(log_filename),\n",
+    "                logging.StreamHandler()  # also print to console\n",
+    "            ]\n",
+    "        )\n",
+    "        self.logger = logging.getLogger(\"TrainingLoop\")\n",
+    "        self.log_file = log_filename\n",
+    "\n",
+    "    def run_episode(self, task_name: str, episode_id: int) -> Trajectory:\n",
+    "        \"\"\"Run a single episode and return the trajectory.\"\"\"\n",
     "        few_shots = self.bank.get_examples(task_name)\n",
-    "        traj = Trajectory(task_name=task_name, episode_id=episode_id)\n",
+    "        trajectory = Trajectory(task_name=task_name, episode_id=episode_id)\n",
+    "\n",
+    "        # Reset environment\n",
     "        result = self.env.reset(task_name)\n",
-    "        obs, done, history = result.get(\"observation\", {}), result.get(\"done\", False), []\n",
-    "        print(f\"  [Episode {episode_id}] task={task_name} few_shots={len(few_shots)}\")\n",
+    "        obs = result.get(\"observation\", {})\n",
+    "        done = result.get(\"done\", False)\n",
+    "        history = []\n",
+    "\n",
+    "        self.logger.info(f\"START episode={episode_id} task={task_name} few_shots_available={len(few_shots)}\")\n",
+    "\n",
     "        step_num = 0\n",
     "        while not done and step_num < obs.get(\"max_steps\", 7):\n",
     "            step_num += 1\n",
-    "            action_type, content = self.agent.get_action(obs, step_num, history, few_shots)\n",
+    "\n",
+    "            # Get action from agent\n",
+    "            action_type, content = self.agent.get_action(\n",
+    "                observation=obs,\n",
+    "                step_number=step_num,\n",
+    "                episode_history=history,\n",
+    "                few_shot_examples=few_shots,\n",
+    "                task_name=task_name\n",
+    "            )\n",
+    "\n",
+    "            # Execute action\n",
     "            result = self.env.step(action_type, content)\n",
-    "            reward, done = result.get(\"reward\", 0.0), result.get(\"done\", False)\n",
-    "            obs, info = result.get(\"observation\", {}), result.get(\"info\", {})\n",
-    "            step = Step(step_num, action_type, content[:300], reward, obs.get(\"current_accuracy\", 0.0), obs.get(\"feedback\", \"\") or \"\", obs.get(\"clarification_response\"))\n",
-    "            traj.steps.append(step); traj.total_reward += reward\n",
-    "            history.append(f\"Step {step_num}: {action_type} -> reward={reward:.2f} acc={step.accuracy:.2f}\")\n",
-    "            print(f\"    step={step_num} action={action_type} reward={reward:.3f} acc={step.accuracy:.2f}\")\n",
+    "            reward = result.get(\"reward\", 0.0)\n",
+    "            done = result.get(\"done\", False)\n",
+    "            obs = result.get(\"observation\", {})\n",
+    "            info = result.get(\"info\", {})\n",
+    "\n",
+    "            # Record step\n",
+    "            step = Step(\n",
+    "                step_number=step_num,\n",
+    "                action_type=action_type,\n",
+    "                action_content=content[:300],\n",
+    "                reward=reward,\n",
+    "                accuracy=obs.get(\"current_accuracy\", 0.0),\n",
+    "                feedback=obs.get(\"feedback\", \"\") or \"\",\n",
+    "                clarification_response=obs.get(\"clarification_response\")\n",
+    "            )\n",
+    "            trajectory.steps.append(step)\n",
+    "            trajectory.total_reward += reward\n",
+    "\n",
+    "            # Update history\n",
+    "            history.append(f\"Step {step_num}: {action_type} \u00e2\u2020\u2019 reward={reward:.2f} acc={step.accuracy:.2f}\")\n",
+    "\n",
+    "            self.logger.info(f\"STEP episode={episode_id} step={step_num} action={action_type} reward={reward:.4f} accuracy={step.accuracy:.4f}\")\n",
+    "\n",
     "            if done:\n",
-    "                traj.final_accuracy = info.get(\"episode_score\", obs.get(\"current_accuracy\", 0.0))\n",
-    "                traj.success = obs.get(\"current_accuracy\", 0.0) >= 0.9\n",
+    "                episode_score = info.get(\"episode_score\", obs.get(\"current_accuracy\", 0.0))\n",
+    "                trajectory.final_accuracy = episode_score\n",
+    "                trajectory.success = obs.get(\"current_accuracy\", 0.0) >= 0.9\n",
     "                break\n",
-    "        if not traj.steps: traj.final_accuracy = 0.0\n",
-    "        return traj\n",
+    "\n",
+    "        if not trajectory.steps:\n",
+    "            trajectory.final_accuracy = 0.0\n",
+    "\n",
+    "        self.logger.info(f\"END episode={episode_id} task={task_name} total_reward={trajectory.total_reward:.4f} final_accuracy={trajectory.final_accuracy:.4f} success={trajectory.success} steps={len(trajectory.steps)}\")\n",
+    "\n",
+    "        return trajectory\n",
     "\n",
     "    def run(self):\n",
-    "        print(\"=\" * 60)\n",
-    "        print(\"REWARD-GUIDED TRAJECTORY OPTIMIZATION\")\n",
-    "        print(f\"Tasks: {TASKS}, Episodes/task: {NUM_EPISODES_PER_TASK}\")\n",
-    "        print(\"=\" * 60)\n",
-    "        if not self.env.health(): raise RuntimeError(f\"Env not reachable at {ENV_BASE_URL}\")\n",
-    "        print(f\"Environment: OK\\n\")\n",
-    "        global_ep = 0\n",
+    "        \"\"\"Run full training loop across all tasks.\"\"\"\n",
+    "        self.logger.info(\"=\" * 60)\n",
+    "        self.logger.info(\"REWARD-GUIDED TRAJECTORY OPTIMIZATION\")\n",
+    "        self.logger.info(f\"Tasks: {TASKS}\")\n",
+    "        self.logger.info(f\"Episodes per task: {NUM_EPISODES_PER_TASK}\")\n",
+    "        self.logger.info(f\"Top-K trajectories: {TOP_K_TRAJECTORIES}\")\n",
+    "        self.logger.info(\"=\" * 60)\n",
+    "        self.logger.info(f\"Log file: {self.log_file}\")\n",
+    "\n",
+    "        try:\n",
+    "            wandb.init(\n",
+    "                project=\"policy-to-logic-rl\",\n",
+    "                name=f\"trajectory-opt-{int(time.time())}\",\n",
+    "                config={\n",
+    "                    \"num_episodes_per_task\": NUM_EPISODES_PER_TASK,\n",
+    "                    \"top_k_trajectories\": TOP_K_TRAJECTORIES,\n",
+    "                    \"min_reward_threshold\": MIN_REWARD_THRESHOLD,\n",
+    "                    \"model\": MODEL,\n",
+    "                    \"temperature\": TEMPERATURE,\n",
+    "                    \"tasks\": TASKS,\n",
+    "                    \"env_url\": ENV_BASE_URL,\n",
+    "                }\n",
+    "            )\n",
+    "        except Exception as e:\n",
+    "            self.logger.warning(f\"Wandb init failed: {e}. Continuing without W&B.\")\n",
+    "\n",
+    "        # Health check\n",
+    "        if not self.env.health():\n",
+    "            raise RuntimeError(f\"Environment not reachable at {ENV_BASE_URL}\")\n",
+    "        self.logger.info(f\"Environment: OK ({ENV_BASE_URL})\\n\")\n",
+    "\n",
+    "        global_episode = 0\n",
+    "\n",
     "        for task in TASKS:\n",
-    "            print(f\"\\n--- TASK: {task} ---\")\n",
+    "            self.logger.info(f\"\\n{'\u00e2\u201d\u20ac'*40}\")\n",
+    "            self.logger.info(f\"TASK: {task}\")\n",
+    "            self.logger.info(f\"{'\u00e2\u201d\u20ac'*40}\")\n",
+    "\n",
     "            task_rewards = []\n",
+    "            task_accuracies = []\n",
+    "\n",
     "            for ep in range(1, NUM_EPISODES_PER_TASK + 1):\n",
-    "                global_ep += 1\n",
-    "                traj = self.run_episode(task, ep)\n",
-    "                self.bank.store(traj)\n",
-    "                self.metrics.append({\"global_episode\": global_ep, \"task\": task, \"episode_in_task\": ep, \"total_reward\": traj.total_reward, \"final_accuracy\": traj.final_accuracy, \"success\": traj.success, \"num_steps\": len(traj.steps)})\n",
-    "                task_rewards.append(traj.total_reward)\n",
-    "                print(f\"  -> Ep {ep}: reward={traj.total_reward:.3f} acc={traj.final_accuracy:.2f} success={traj.success}\")\n",
-    "                time.sleep(0.5)\n",
-    "            print(f\"  Improvement: {task_rewards[-1] - task_rewards[0]:+.3f}\")\n",
-    "        print(\"\\n\" + \"=\" * 60 + \"\\nTRAINING COMPLETE\\n\" + \"=\" * 60)\n",
+    "                global_episode += 1\n",
+    "                trajectory = self.run_episode(task, ep)\n",
+    "\n",
+    "                # Store in bank\n",
+    "                self.bank.store(trajectory)\n",
+    "\n",
+    "                try:\n",
+    "                    wandb.log({\n",
+    "                        f\"{task}/total_reward\": trajectory.total_reward,\n",
+    "                        f\"{task}/final_accuracy\": trajectory.final_accuracy,\n",
+    "                        f\"{task}/num_steps\": len(trajectory.steps),\n",
+    "                        f\"{task}/success\": int(trajectory.success),\n",
+    "                        f\"{task}/few_shots_used\": len(self.bank.get_examples(task)),\n",
+    "                        \"global/total_reward\": trajectory.total_reward,\n",
+    "                        \"global/final_accuracy\": trajectory.final_accuracy,\n",
+    "                        \"episode\": global_episode,\n",
+    "                    })\n",
+    "                except Exception:\n",
+    "                    pass\n",
+    "\n",
+    "                # Record metrics\n",
+    "                self.metrics.append({\n",
+    "                    \"global_episode\": global_episode,\n",
+    "                    \"task\": task,\n",
+    "                    \"episode_in_task\": ep,\n",
+    "                    \"total_reward\": trajectory.total_reward,\n",
+    "                    \"final_accuracy\": trajectory.final_accuracy,\n",
+    "                    \"success\": trajectory.success,\n",
+    "                    \"num_steps\": len(trajectory.steps),\n",
+    "                    \"few_shots_used\": len(self.bank.get_examples(task)) - (1 if trajectory.total_reward >= MIN_REWARD_THRESHOLD else 0)\n",
+    "                })\n",
+    "\n",
+    "                task_rewards.append(trajectory.total_reward)\n",
+    "                task_accuracies.append(trajectory.final_accuracy)\n",
+    "\n",
+    "                self.logger.info(f\"  \u00e2\u2020\u2019 Episode {ep} complete: reward={trajectory.total_reward:.3f} accuracy={trajectory.final_accuracy:.2f} success={trajectory.success}\")\n",
+    "                time.sleep(0.5)  # Rate limiting\n",
+    "\n",
+    "            self.logger.info(f\"\\n  Task summary:\")\n",
+    "            self.logger.info(f\"    First episode reward: {task_rewards[0]:.3f}\")\n",
+    "            self.logger.info(f\"    Last episode reward:  {task_rewards[-1]:.3f}\")\n",
+    "            self.logger.info(f\"    Improvement: {task_rewards[-1] - task_rewards[0]:+.3f}\")\n",
+    "\n",
+    "        self.logger.info(\"\\n\" + \"=\" * 60)\n",
+    "        self.logger.info(\"TRAINING COMPLETE\")\n",
+    "        self.logger.info(f\"Bank summary: {self.bank.summary()}\")\n",
+    "        self.logger.info(\"=\" * 60)\n",
+    "\n",
+    "        try:\n",
+    "            wandb.finish()\n",
+    "        except Exception:\n",
+    "            pass\n",
+    "\n",
     "        return self.metrics\n",
     "\n",
-    "def save_plots(metrics):\n",
-    "    import matplotlib; matplotlib.use(\"Agg\")\n",
-    "    import matplotlib.pyplot as plt; import numpy as np\n",
+    "# \u00e2\u201d\u20ac\u00e2\u201d\u20ac Plot Generation \u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\n",
+    "\n",
+    "def save_plots(metrics: list[dict]):\n",
+    "    \"\"\"\n",
+    "    Save reward curve and accuracy curve as PNG files.\n",
+    "    These are REQUIRED for hackathon submission \u00e2\u20ac\u201d must be committed to repo.\n",
+    "    \"\"\"\n",
+    "    try:\n",
+    "        import matplotlib\n",
+    "        matplotlib.use(\"Agg\")  # Non-interactive backend\n",
+    "        import matplotlib.pyplot as plt\n",
+    "        import numpy as np\n",
+    "    except ImportError:\n",
+    "        print(\"matplotlib not installed. Run: pip install matplotlib\")\n",
+    "        return\n",
+    "\n",
     "    os.makedirs(\"training/plots\", exist_ok=True)\n",
+    "\n",
     "    episodes = [m[\"global_episode\"] for m in metrics]\n",
     "    rewards = [m[\"total_reward\"] for m in metrics]\n",
-    "    colors = {\"data_access\": \"#2196F3\", \"resource_access\": \"#FF9800\", \"transaction_approval\": \"#4CAF50\"}\n",
-    "    # Plot 1: Reward\n",
+    "    accuracies = [m[\"final_accuracy\"] for m in metrics]\n",
+    "    tasks = [m[\"task\"] for m in metrics]\n",
+    "\n",
+    "    colors = {\n",
+    "        \"data_access\": \"#2196F3\",\n",
+    "        \"resource_access\": \"#FF9800\",\n",
+    "        \"transaction_approval\": \"#4CAF50\"\n",
+    "    }\n",
+    "\n",
+    "    # \u00e2\u201d\u20ac\u00e2\u201d\u20ac Plot 1: Reward Curve (per-task trend lines) \u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\n",
     "    fig, ax = plt.subplots(figsize=(10, 5))\n",
+    "\n",
     "    for task in TASKS:\n",
-    "        te = [m[\"global_episode\"] for m in metrics if m[\"task\"]==task]\n",
-    "        tr = [m[\"total_reward\"] for m in metrics if m[\"task\"]==task]\n",
-    "        ax.plot(te, tr, marker=\"o\", label=task, color=colors.get(task), linewidth=2, markersize=5)\n",
-    "    z = np.polyfit(episodes, rewards, 1); p = np.poly1d(z)\n",
-    "    ax.plot(episodes, p(episodes), \"--\", color=\"red\", alpha=0.5, label=\"trend\")\n",
-    "    ax.set_xlabel(\"Episode\"); ax.set_ylabel(\"Total Reward\"); ax.set_title(\"Reward Curve\"); ax.legend(); ax.grid(True, alpha=0.3); ax.set_ylim(bottom=0)\n",
-    "    plt.tight_layout(); plt.savefig(\"training/plots/reward_curve.png\", dpi=150); plt.close()\n",
-    "    # Plot 2: Accuracy\n",
+    "        task_eps = [m[\"global_episode\"] for m in metrics if m[\"task\"] == task]\n",
+    "        task_rews = [m[\"total_reward\"] for m in metrics if m[\"task\"] == task]\n",
+    "        ax.plot(task_eps, task_rews, marker=\"o\", label=task,\n",
+    "                color=colors.get(task, \"gray\"), linewidth=2, markersize=5)\n",
+    "        # Per-task trend line\n",
+    "        if len(task_eps) >= 2:\n",
+    "            z = np.polyfit(task_eps, task_rews, 1)\n",
+    "            p = np.poly1d(z)\n",
+    "            ax.plot(task_eps, p(task_eps), \"--\",\n",
+    "                    color=colors.get(task, \"gray\"), alpha=0.4, linewidth=1.5)\n",
+    "\n",
+    "    ax.set_xlabel(\"Episode\")\n",
+    "    ax.set_ylabel(\"Total Reward\")\n",
+    "    ax.set_title(\"Reward Curve \u00e2\u20ac\u201d Reward-Guided Trajectory Optimization\")\n",
+    "    ax.legend()\n",
+    "    ax.grid(True, alpha=0.3)\n",
+    "    ax.set_ylim(bottom=0)\n",
+    "\n",
+    "    plt.tight_layout()\n",
+    "    plt.savefig(\"training/plots/reward_curve.png\", dpi=150, bbox_inches=\"tight\")\n",
+    "    plt.close()\n",
+    "    print(\"Saved: training/plots/reward_curve.png\")\n",
+    "\n",
+    "    # \u00e2\u201d\u20ac\u00e2\u201d\u20ac Plot 2: Accuracy Curve \u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\n",
     "    fig, ax = plt.subplots(figsize=(10, 5))\n",
+    "\n",
     "    for task in TASKS:\n",
-    "        te = [m[\"global_episode\"] for m in metrics if m[\"task\"]==task]\n",
-    "        ta = [m[\"final_accuracy\"] for m in metrics if m[\"task\"]==task]\n",
-    "        ax.plot(te, ta, marker=\"s\", label=task, color=colors.get(task), linewidth=2, markersize=5)\n",
-    "    ax.axhline(y=0.9, color=\"red\", linestyle=\"--\", alpha=0.7, label=\"threshold\")\n",
-    "    ax.set_xlabel(\"Episode\"); ax.set_ylabel(\"Accuracy\"); ax.set_title(\"Accuracy Curve\"); ax.legend(); ax.grid(True, alpha=0.3); ax.set_ylim(0, 1.05)\n",
-    "    plt.tight_layout(); plt.savefig(\"training/plots/accuracy_curve.png\", dpi=150); plt.close()\n",
-    "    # Plot 3: Improvement\n",
-    "    fig, ax = plt.subplots(figsize=(8, 5))\n",
-    "    tnames, imps = [], []\n",
+    "        task_eps = [m[\"global_episode\"] for m in metrics if m[\"task\"] == task]\n",
+    "        task_accs = [m[\"final_accuracy\"] for m in metrics if m[\"task\"] == task]\n",
+    "        ax.plot(task_eps, task_accs, marker=\"s\", label=task,\n",
+    "                color=colors.get(task, \"gray\"), linewidth=2, markersize=5)\n",
+    "\n",
+    "    ax.axhline(y=0.9, color=\"red\", linestyle=\"--\", alpha=0.7, label=\"success threshold (0.9)\")\n",
+    "\n",
+    "    ax.set_xlabel(\"Episode\")\n",
+    "    ax.set_ylabel(\"Final Accuracy\")\n",
+    "    ax.set_title(\"Accuracy Curve \u00e2\u20ac\u201d Policy-to-Logic Agent\")\n",
+    "    ax.legend()\n",
+    "    ax.grid(True, alpha=0.3)\n",
+    "    ax.set_ylim(0, 1.05)\n",
+    "\n",
+    "    plt.tight_layout()\n",
+    "    plt.savefig(\"training/plots/accuracy_curve.png\", dpi=150, bbox_inches=\"tight\")\n",
+    "    plt.close()\n",
+    "    print(\"Saved: training/plots/accuracy_curve.png\")\n",
+    "\n",
+    "    # \u00e2\u201d\u20ac\u00e2\u201d\u20ac Plot 3: Per-Task Summary (Accuracy + Efficiency) \u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\n",
+    "    fig, axes = plt.subplots(1, 2, figsize=(14, 5))\n",
+    "\n",
+    "    task_labels = []\n",
+    "    acc_improvements = []\n",
+    "    eff_improvements = []\n",
+    "    best_accuracies = []\n",
+    "\n",
     "    for task in TASKS:\n",
-    "        accs = [m[\"final_accuracy\"] for m in metrics if m[\"task\"]==task]\n",
-    "        if len(accs) >= 2: tnames.append(task.replace(\"_\",\"\\n\")); imps.append(accs[-1]-accs[0])\n",
-    "    bars = ax.bar(tnames, imps, color=[\"#2196F3\",\"#FF9800\",\"#4CAF50\"][:len(tnames)])\n",
-    "    ax.axhline(y=0, color=\"black\"); ax.set_ylabel(\"Improvement\"); ax.set_title(\"Per-Task Improvement\"); ax.grid(True, axis=\"y\", alpha=0.3)\n",
-    "    for bar, val in zip(bars, imps): ax.text(bar.get_x()+bar.get_width()/2, bar.get_height()+0.01, f\"{val:+.2f}\", ha=\"center\", fontweight=\"bold\")\n",
-    "    plt.tight_layout(); plt.savefig(\"training/plots/improvement_chart.png\", dpi=150); plt.close()\n",
-    "    with open(\"training/plots/metrics.json\", \"w\") as f: json.dump(metrics, f, indent=2)\n",
-    "    print(\"All plots saved to training/plots/\")"
+    "        task_data = [m for m in metrics if m[\"task\"] == task]\n",
+    "        if len(task_data) >= 2:\n",
+    "            task_labels.append(task.replace(\"_\", \"\\n\"))\n",
+    "            acc_improvements.append(task_data[-1][\"final_accuracy\"] - task_data[0][\"final_accuracy\"])\n",
+    "            # Efficiency: steps saved (first vs best)\n",
+    "            first_steps = task_data[0][\"num_steps\"]\n",
+    "            best_steps = min(m[\"num_steps\"] for m in task_data)\n",
+    "            eff_pct = ((first_steps - best_steps) / first_steps * 100) if first_steps > 0 else 0\n",
+    "            eff_improvements.append(eff_pct)\n",
+    "            best_accuracies.append(max(m[\"final_accuracy\"] for m in task_data))\n",
+    "\n",
+    "    # Left: Best accuracy per task\n",
+    "    bars1 = axes[0].bar(task_labels, best_accuracies,\n",
+    "                        color=[\"#2196F3\", \"#FF9800\", \"#4CAF50\"][:len(task_labels)],\n",
+    "                        edgecolor=\"white\", linewidth=1.5)\n",
+    "    axes[0].axhline(y=0.9, color=\"red\", linestyle=\"--\", alpha=0.7, label=\"success threshold\")\n",
+    "    axes[0].set_ylabel(\"Best Accuracy Achieved\")\n",
+    "    axes[0].set_title(\"Best Accuracy Per Task\")\n",
+    "    axes[0].set_ylim(0, 1.1)\n",
+    "    axes[0].grid(True, axis=\"y\", alpha=0.3)\n",
+    "    axes[0].legend()\n",
+    "    for bar, val in zip(bars1, best_accuracies):\n",
+    "        axes[0].text(bar.get_x() + bar.get_width() / 2, bar.get_height() + 0.02,\n",
+    "                     f\"{val:.0%}\", ha=\"center\", va=\"bottom\", fontweight=\"bold\")\n",
+    "\n",
+    "    # Right: Efficiency improvement (% steps saved)\n",
+    "    bars2 = axes[1].bar(task_labels, eff_improvements,\n",
+    "                        color=[\"#2196F3\", \"#FF9800\", \"#4CAF50\"][:len(task_labels)],\n",
+    "                        edgecolor=\"white\", linewidth=1.5)\n",
+    "    axes[1].axhline(y=0, color=\"black\", linewidth=0.8)\n",
+    "    axes[1].set_ylabel(\"Steps Saved (%)\")\n",
+    "    axes[1].set_title(\"Efficiency Improvement (First \u00e2\u2020\u2019 Best Episode)\")\n",
+    "    axes[1].grid(True, axis=\"y\", alpha=0.3)\n",
+    "    for bar, val in zip(bars2, eff_improvements):\n",
+    "        y_pos = max(bar.get_height() + 1, 2)\n",
+    "        axes[1].text(bar.get_x() + bar.get_width() / 2, y_pos,\n",
+    "                     f\"{val:.0f}%\", ha=\"center\", va=\"bottom\", fontweight=\"bold\")\n",
+    "\n",
+    "    plt.tight_layout()\n",
+    "    plt.savefig(\"training/plots/improvement_chart.png\", dpi=150, bbox_inches=\"tight\")\n",
+    "    plt.close()\n",
+    "    print(\"Saved: training/plots/improvement_chart.png\")\n",
+    "\n",
+    "    # \u00e2\u201d\u20ac\u00e2\u201d\u20ac Save raw metrics as JSON \u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\n",
+    "    timestamp = int(time.time())\n",
+    "    with open(f\"training/plots/metrics_{timestamp}.json\", \"w\") as f:\n",
+    "        json.dump(metrics, f, indent=2)\n",
+    "    with open(\"training/plots/metrics_latest.json\", \"w\") as f:\n",
+    "        json.dump(metrics, f, indent=2)\n",
+    "    print(f\"Saved: training/plots/metrics_{timestamp}.json and metrics_latest.json\")\n",
+    "\n",
+    "# \u00e2\u201d\u20ac\u00e2\u201d\u20ac Entry Point \u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\n",
+    "\n",
+    "if __name__ == \"__main__\":\n",
+    "    hf_token = os.getenv(\"HF_TOKEN\", \"\")\n",
+    "    if not hf_token:\n",
+    "        raise ValueError(\"HF_TOKEN environment variable not set\")\n",
+    "\n",
+    "    loop = TrainingLoop(ENV_BASE_URL, hf_token)\n",
+    "    metrics = loop.run()\n",
+    "    save_plots(metrics)\n",
+    "\n",
+    "    print(\"\\nNext step: commit training/plots/*.png to repo for submission.\")\n",
+    ""
    ]
   },
   {
@@ -335,7 +729,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Cell 5: Run training loop\n",
+    "# Cell 6: Run training loop\n",
     "loop = TrainingLoop(ENV_URL, HF_TOKEN)\n",
     "metrics = loop.run()\n",
     "print(f\"\\nTotal episodes run: {len(metrics)}\")"
@@ -347,7 +741,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Cell 6: Generate plots and display inline\n",
+    "# Cell 7: Generate plots and display inline\n",
     "save_plots(metrics)\n",
     "\n",
     "from IPython.display import Image, display\n",
@@ -362,17 +756,27 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Cell 7: Download plots to commit to repo\n",
-    "# After running this, download the files and commit them to your GitHub repo\n",
+    "# Cell 8: Display wandb run link\n",
+    "print(f\"Wandb run: https://wandb.ai/YOUR_USERNAME/policy-to-logic-rl\")\n",
+    "print(\"Add this link to your README under Deliverables.\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Cell 9: Download plots to commit to repo\n",
     "from google.colab import files\n",
     "\n",
     "files.download(\"training/plots/reward_curve.png\")\n",
     "files.download(\"training/plots/accuracy_curve.png\")\n",
     "files.download(\"training/plots/improvement_chart.png\")\n",
-    "files.download(\"training/plots/metrics.json\")\n",
+    "files.download(\"training/plots/metrics_latest.json\")\n",
     "\n",
     "print(\"Downloaded. Now commit these files to: training/plots/ in your repo.\")"
    ]
   }
  ]
-}
+}

training/results-iteration1/metrics (1).json

@@ -0,0 +1,218 @@
+[
+  {
+    "global_episode": 1,
+    "task": "data_access",
+    "episode_in_task": 1,
+    "total_reward": 1.1988333333333334,
+    "final_accuracy": 0.92,
+    "success": true,
+    "num_steps": 4
+  },
+  {
+    "global_episode": 2,
+    "task": "data_access",
+    "episode_in_task": 2,
+    "total_reward": 0.7585,
+    "final_accuracy": 0.9600000000000001,
+    "success": true,
+    "num_steps": 2
+  },
+  {
+    "global_episode": 3,
+    "task": "data_access",
+    "episode_in_task": 3,
+    "total_reward": 0.7585,
+    "final_accuracy": 0.9600000000000001,
+    "success": true,
+    "num_steps": 2
+  },
+  {
+    "global_episode": 4,
+    "task": "data_access",
+    "episode_in_task": 4,
+    "total_reward": 0.7585,
+    "final_accuracy": 0.9600000000000001,
+    "success": true,
+    "num_steps": 2
+  },
+  {
+    "global_episode": 5,
+    "task": "data_access",
+    "episode_in_task": 5,
+    "total_reward": 0.7585,
+    "final_accuracy": 0.9600000000000001,
+    "success": true,
+    "num_steps": 2
+  },
+  {
+    "global_episode": 6,
+    "task": "data_access",
+    "episode_in_task": 6,
+    "total_reward": 0.7585,
+    "final_accuracy": 0.9600000000000001,
+    "success": true,
+    "num_steps": 2
+  },
+  {
+    "global_episode": 7,
+    "task": "data_access",
+    "episode_in_task": 7,
+    "total_reward": 0.7585,
+    "final_accuracy": 0.9600000000000001,
+    "success": true,
+    "num_steps": 2
+  },
+  {
+    "global_episode": 8,
+    "task": "data_access",
+    "episode_in_task": 8,
+    "total_reward": 0.7585,
+    "final_accuracy": 0.9600000000000001,
+    "success": true,
+    "num_steps": 2
+  },
+  {
+    "global_episode": 9,
+    "task": "resource_access",
+    "episode_in_task": 1,
+    "total_reward": 1.041,
+    "final_accuracy": 0.8931428571428572,
+    "success": true,
+    "num_steps": 3
+  },
+  {
+    "global_episode": 10,
+    "task": "resource_access",
+    "episode_in_task": 2,
+    "total_reward": 1.041,
+    "final_accuracy": 0.8931428571428572,
+    "success": true,
+    "num_steps": 3
+  },
+  {
+    "global_episode": 11,
+    "task": "resource_access",
+    "episode_in_task": 3,
+    "total_reward": 0.7335,
+    "final_accuracy": 0.9074285714285715,
+    "success": true,
+    "num_steps": 2
+  },
+  {
+    "global_episode": 12,
+    "task": "resource_access",
+    "episode_in_task": 4,
+    "total_reward": 0.7335,
+    "final_accuracy": 0.9074285714285715,
+    "success": true,
+    "num_steps": 2
+  },
+  {
+    "global_episode": 13,
+    "task": "resource_access",
+    "episode_in_task": 5,
+    "total_reward": 0.7434999999999999,
+    "final_accuracy": 0.9234285714285714,
+    "success": true,
+    "num_steps": 2
+  },
+  {
+    "global_episode": 14,
+    "task": "resource_access",
+    "episode_in_task": 6,
+    "total_reward": 0.7434999999999999,
+    "final_accuracy": 0.9234285714285714,
+    "success": true,
+    "num_steps": 2
+  },
+  {
+    "global_episode": 15,
+    "task": "resource_access",
+    "episode_in_task": 7,
+    "total_reward": 1.929,
+    "final_accuracy": 0.8645714285714287,
+    "success": true,
+    "num_steps": 5
+  },
+  {
+    "global_episode": 16,
+    "task": "resource_access",
+    "episode_in_task": 8,
+    "total_reward": 0.7335,
+    "final_accuracy": 0.9074285714285715,
+    "success": true,
+    "num_steps": 2
+  },
+  {
+    "global_episode": 17,
+    "task": "transaction_approval",
+    "episode_in_task": 1,
+    "total_reward": 0.10799999999999998,
+    "final_accuracy": 0.0,
+    "success": false,
+    "num_steps": 7
+  },
+  {
+    "global_episode": 18,
+    "task": "transaction_approval",
+    "episode_in_task": 2,
+    "total_reward": 0.10799999999999998,
+    "final_accuracy": 0.0,
+    "success": false,
+    "num_steps": 7
+  },
+  {
+    "global_episode": 19,
+    "task": "transaction_approval",
+    "episode_in_task": 3,
+    "total_reward": 0.10799999999999998,
+    "final_accuracy": 0.0,
+    "success": false,
+    "num_steps": 7
+  },
+  {
+    "global_episode": 20,
+    "task": "transaction_approval",
+    "episode_in_task": 4,
+    "total_reward": 0.10799999999999998,
+    "final_accuracy": 0.0,
+    "success": false,
+    "num_steps": 7
+  },
+  {
+    "global_episode": 21,
+    "task": "transaction_approval",
+    "episode_in_task": 5,
+    "total_reward": 0.10799999999999998,
+    "final_accuracy": 0.0,
+    "success": false,
+    "num_steps": 7
+  },
+  {
+    "global_episode": 22,
+    "task": "transaction_approval",
+    "episode_in_task": 6,
+    "total_reward": 0.10799999999999998,
+    "final_accuracy": 0.0,
+    "success": false,
+    "num_steps": 7
+  },
+  {
+    "global_episode": 23,
+    "task": "transaction_approval",
+    "episode_in_task": 7,
+    "total_reward": 0.10799999999999998,
+    "final_accuracy": 0.0,
+    "success": false,
+    "num_steps": 7
+  },
+  {
+    "global_episode": 24,
+    "task": "transaction_approval",
+    "episode_in_task": 8,
+    "total_reward": 0.10799999999999998,
+    "final_accuracy": 0.0,
+    "success": false,
+    "num_steps": 7
+  }
+]

training/trajectory_optimizer.py

@@ -13,6 +13,8 @@ import json
 import os
 import time
 import requests
+import logging
+import wandb
 from dataclasses import dataclass, field
 from typing import Optional
 from openai import OpenAI
@@ -106,14 +108,15 @@ class Agent:
         observation: dict,
         step_number: int,
         episode_history: list[str],
-        few_shot_examples: list[Trajectory]
+        few_shot_examples: list[Trajectory],
+        task_name: str = ""
     ) -> tuple[str, str]:
         """
         Returns (action_type, content_json_string).
         action_type: one of ask_clarification | propose_rules | refine_rules
         content: JSON string appropriate for that action
         """
-        system_prompt = self._build_system_prompt(few_shot_examples)
+        system_prompt = self._build_system_prompt(few_shot_examples, task_name)
         user_prompt = self._build_user_prompt(observation, step_number, episode_history)
 
         try:
@@ -132,7 +135,7 @@ class Agent:
             print(f"    [LLM ERROR] {e}")
             return "propose_rules", json.dumps({"rules": [], "default": "DENY"})
 
-    def _build_system_prompt(self, few_shot_examples: list[Trajectory]) -> str:
+    def _build_system_prompt(self, few_shot_examples: list[Trajectory], task_name: str = "") -> str:
         base = """You are a policy-to-logic agent. Your job is to convert natural language policies into executable rules.
 
 AVAILABLE ACTIONS:
@@ -164,6 +167,39 @@ STRATEGY:
 OUTPUT FORMAT: Respond ONLY with valid JSON. No markdown. No explanation.
 {"action_type": "propose_rules", "content": "{...escaped json string...}"}
 """
+        # Task-specific guidance for complex tasks
+        if task_name == "transaction_approval":
+            base += """
+IMPORTANT — TRANSACTION APPROVAL TASK:
+This task has 4 possible decisions: APPROVE, REQUIRE_APPROVAL, COMPLIANCE_REVIEW, HOLD
+Rules are evaluated TOP-TO-BOTTOM. Order matters critically. You MUST order rules by priority:
+  1. FIRST: Check if transfer_type == "international" → then COMPLIANCE_REVIEW (always, overrides everything)
+  2. SECOND: Check if amount >= 10000 AND time is outside business hours (hour < 9 or hour >= 17) → then HOLD
+  3. THIRD: Check if amount > 5000 AND initiator_role != "manager" → then REQUIRE_APPROVAL
+  4. DEFAULT: APPROVE
+
+Key details:
+- Standard limit is $5,000 (amount > 5000 triggers approval, NOT >=)
+- High-value threshold is $10,000 (amount >= 10000)
+- Business hours: hour >= 9 AND hour < 17
+- Manager exemption ONLY applies to the standard $5,000 limit, NOT to international or high-value HOLD rules
+- "system" role follows the same rules as "employee"
+
+Here is a working example of valid rules for this task:
+{"rules": [{"if": [{"field": "transfer_type", "op": "==", "value": "international"}], "then": "COMPLIANCE_REVIEW"}, {"if": [{"field": "amount", "op": ">=", "value": 10000}, {"field": "time", "op": ">=", "value": 17}], "then": "HOLD"}, {"if": [{"field": "amount", "op": ">=", "value": 10000}, {"field": "time", "op": "<", "value": 9}], "then": "HOLD"}, {"if": [{"field": "amount", "op": ">", "value": 5000}, {"field": "initiator_role", "op": "!=", "value": "manager"}], "then": "REQUIRE_APPROVAL"}], "default": "APPROVE"}
+"""
+        elif task_name == "resource_access":
+            base += """
+IMPORTANT — RESOURCE ACCESS TASK:
+This task has roles: junior, senior, contractor. Document types: public, internal, confidential.
+- Senior employees: ALLOW everything always
+- Contractors: ALLOW only public, DENY everything else
+- Junior + confidential: ALWAYS DENY (regardless of time — the policy is misleading about this)
+- Junior + internal: ALLOW only during business hours (hour >= 8 AND hour < 17)
+- Junior + public: ALLOW always
+- Business hours: hour >= 8 AND hour < 17
+"""
+
         if few_shot_examples:
             base += "\n\nLEARNED FROM PREVIOUS EPISODES (high-reward strategies):\n"
             for traj in few_shot_examples[-TOP_K_TRAJECTORIES:]:
@@ -255,6 +291,19 @@ class TrainingLoop:
         self.bank = TrajectoryBank()
         self.metrics = []  # List of {episode, task, reward, accuracy, success}
 
+        os.makedirs("training/logs", exist_ok=True)
+        log_filename = f"training/logs/run_{int(time.time())}.log"
+        logging.basicConfig(
+            level=logging.INFO,
+            format="%(asctime)s [%(levelname)s] %(message)s",
+            handlers=[
+                logging.FileHandler(log_filename),
+                logging.StreamHandler()  # also print to console
+            ]
+        )
+        self.logger = logging.getLogger("TrainingLoop")
+        self.log_file = log_filename
+
     def run_episode(self, task_name: str, episode_id: int) -> Trajectory:
         """Run a single episode and return the trajectory."""
         few_shots = self.bank.get_examples(task_name)
@@ -266,7 +315,7 @@ class TrainingLoop:
         done = result.get("done", False)
         history = []
 
-        print(f"  [Episode {episode_id}] task={task_name} few_shots={len(few_shots)}")
+        self.logger.info(f"START episode={episode_id} task={task_name} few_shots_available={len(few_shots)}")
 
         step_num = 0
         while not done and step_num < obs.get("max_steps", 7):
@@ -277,7 +326,8 @@ class TrainingLoop:
                 observation=obs,
                 step_number=step_num,
                 episode_history=history,
-                few_shot_examples=few_shots
+                few_shot_examples=few_shots,
+                task_name=task_name
             )
 
             # Execute action
@@ -303,7 +353,7 @@ class TrainingLoop:
             # Update history
             history.append(f"Step {step_num}: {action_type} → reward={reward:.2f} acc={step.accuracy:.2f}")
 
-            print(f"    step={step_num} action={action_type} reward={reward:.3f} acc={step.accuracy:.2f}")
+            self.logger.info(f"STEP episode={episode_id} step={step_num} action={action_type} reward={reward:.4f} accuracy={step.accuracy:.4f}")
 
             if done:
                 episode_score = info.get("episode_score", obs.get("current_accuracy", 0.0))
@@ -314,28 +364,48 @@ class TrainingLoop:
         if not trajectory.steps:
             trajectory.final_accuracy = 0.0
 
+        self.logger.info(f"END episode={episode_id} task={task_name} total_reward={trajectory.total_reward:.4f} final_accuracy={trajectory.final_accuracy:.4f} success={trajectory.success} steps={len(trajectory.steps)}")
+
         return trajectory
 
     def run(self):
         """Run full training loop across all tasks."""
-        print("=" * 60)
-        print("REWARD-GUIDED TRAJECTORY OPTIMIZATION")
-        print(f"Tasks: {TASKS}")
-        print(f"Episodes per task: {NUM_EPISODES_PER_TASK}")
-        print(f"Top-K trajectories: {TOP_K_TRAJECTORIES}")
-        print("=" * 60)
+        self.logger.info("=" * 60)
+        self.logger.info("REWARD-GUIDED TRAJECTORY OPTIMIZATION")
+        self.logger.info(f"Tasks: {TASKS}")
+        self.logger.info(f"Episodes per task: {NUM_EPISODES_PER_TASK}")
+        self.logger.info(f"Top-K trajectories: {TOP_K_TRAJECTORIES}")
+        self.logger.info("=" * 60)
+        self.logger.info(f"Log file: {self.log_file}")
+
+        try:
+            wandb.init(
+                project="policy-to-logic-rl",
+                name=f"trajectory-opt-{int(time.time())}",
+                config={
+                    "num_episodes_per_task": NUM_EPISODES_PER_TASK,
+                    "top_k_trajectories": TOP_K_TRAJECTORIES,
+                    "min_reward_threshold": MIN_REWARD_THRESHOLD,
+                    "model": MODEL,
+                    "temperature": TEMPERATURE,
+                    "tasks": TASKS,
+                    "env_url": ENV_BASE_URL,
+                }
+            )
+        except Exception as e:
+            self.logger.warning(f"Wandb init failed: {e}. Continuing without W&B.")
 
         # Health check
         if not self.env.health():
             raise RuntimeError(f"Environment not reachable at {ENV_BASE_URL}")
-        print(f"Environment: OK ({ENV_BASE_URL})\n")
+        self.logger.info(f"Environment: OK ({ENV_BASE_URL})\n")
 
         global_episode = 0
 
         for task in TASKS:
-            print(f"\n{'─'*40}")
-            print(f"TASK: {task}")
-            print(f"{'─'*40}")
+            self.logger.info(f"\n{'─'*40}")
+            self.logger.info(f"TASK: {task}")
+            self.logger.info(f"{'─'*40}")
 
             task_rewards = []
             task_accuracies = []
@@ -347,6 +417,20 @@ class TrainingLoop:
                 # Store in bank
                 self.bank.store(trajectory)
 
+                try:
+                    wandb.log({
+                        f"{task}/total_reward": trajectory.total_reward,
+                        f"{task}/final_accuracy": trajectory.final_accuracy,
+                        f"{task}/num_steps": len(trajectory.steps),
+                        f"{task}/success": int(trajectory.success),
+                        f"{task}/few_shots_used": len(self.bank.get_examples(task)),
+                        "global/total_reward": trajectory.total_reward,
+                        "global/final_accuracy": trajectory.final_accuracy,
+                        "episode": global_episode,
+                    })
+                except Exception:
+                    pass
+
                 # Record metrics
                 self.metrics.append({
                     "global_episode": global_episode,
@@ -362,18 +446,23 @@ class TrainingLoop:
                 task_rewards.append(trajectory.total_reward)
                 task_accuracies.append(trajectory.final_accuracy)
 
-                print(f"  → Episode {ep} complete: reward={trajectory.total_reward:.3f} accuracy={trajectory.final_accuracy:.2f} success={trajectory.success}")
+                self.logger.info(f"  → Episode {ep} complete: reward={trajectory.total_reward:.3f} accuracy={trajectory.final_accuracy:.2f} success={trajectory.success}")
                 time.sleep(0.5)  # Rate limiting
 
-            print(f"\n  Task summary:")
-            print(f"    First episode reward: {task_rewards[0]:.3f}")
-            print(f"    Last episode reward:  {task_rewards[-1]:.3f}")
-            print(f"    Improvement: {task_rewards[-1] - task_rewards[0]:+.3f}")
+            self.logger.info(f"\n  Task summary:")
+            self.logger.info(f"    First episode reward: {task_rewards[0]:.3f}")
+            self.logger.info(f"    Last episode reward:  {task_rewards[-1]:.3f}")
+            self.logger.info(f"    Improvement: {task_rewards[-1] - task_rewards[0]:+.3f}")
 
-        print("\n" + "=" * 60)
-        print("TRAINING COMPLETE")
-        print(f"Bank summary: {self.bank.summary()}")
-        print("=" * 60)
+        self.logger.info("\n" + "=" * 60)
+        self.logger.info("TRAINING COMPLETE")
+        self.logger.info(f"Bank summary: {self.bank.summary()}")
+        self.logger.info("=" * 60)
+
+        try:
+            wandb.finish()
+        except Exception:
+            pass
 
         return self.metrics
 
@@ -406,7 +495,7 @@ def save_plots(metrics: list[dict]):
         "transaction_approval": "#4CAF50"
     }
 
-    # ── Plot 1: Reward Curve ──────────────────────────────────────────────────
+    # ── Plot 1: Reward Curve (per-task trend lines) ────────────────────────────
     fig, ax = plt.subplots(figsize=(10, 5))
 
     for task in TASKS:
@@ -414,11 +503,12 @@ def save_plots(metrics: list[dict]):
         task_rews = [m["total_reward"] for m in metrics if m["task"] == task]
         ax.plot(task_eps, task_rews, marker="o", label=task,
                 color=colors.get(task, "gray"), linewidth=2, markersize=5)
-
-    # Trend line
-    z = np.polyfit(episodes, rewards, 1)
-    p = np.poly1d(z)
-    ax.plot(episodes, p(episodes), "--", color="red", alpha=0.5, linewidth=1.5, label="overall trend")
+        # Per-task trend line
+        if len(task_eps) >= 2:
+            z = np.polyfit(task_eps, task_rews, 1)
+            p = np.poly1d(z)
+            ax.plot(task_eps, p(task_eps), "--",
+                    color=colors.get(task, "gray"), alpha=0.4, linewidth=1.5)
 
     ax.set_xlabel("Episode")
     ax.set_ylabel("Total Reward")
@@ -455,32 +545,52 @@ def save_plots(metrics: list[dict]):
     plt.close()
     print("Saved: training/plots/accuracy_curve.png")
 
-    # ── Plot 3: Per-Task Improvement Bar Chart ────────────────────────────────
-    fig, ax = plt.subplots(figsize=(8, 5))
+    # ── Plot 3: Per-Task Summary (Accuracy + Efficiency) ──────────────────────
+    fig, axes = plt.subplots(1, 2, figsize=(14, 5))
 
-    task_names = []
-    improvements = []
+    task_labels = []
+    acc_improvements = []
+    eff_improvements = []
+    best_accuracies = []
 
     for task in TASKS:
-        task_accs = [m["final_accuracy"] for m in metrics if m["task"] == task]
-        if len(task_accs) >= 2:
-            first = task_accs[0]
-            last = task_accs[-1]
-            task_names.append(task.replace("_", "\n"))
-            improvements.append(last - first)
-
-    bars = ax.bar(task_names, improvements,
-                  color=["#2196F3", "#FF9800", "#4CAF50"][:len(task_names)],
-                  edgecolor="white", linewidth=1.5)
-
-    ax.axhline(y=0, color="black", linewidth=0.8)
-    ax.set_ylabel("Accuracy Improvement (last - first episode)")
-    ax.set_title("Per-Task Improvement from Trajectory Accumulation")
-    ax.grid(True, axis="y", alpha=0.3)
-
-    for bar, val in zip(bars, improvements):
-        ax.text(bar.get_x() + bar.get_width() / 2, bar.get_height() + 0.01,
-                f"{val:+.2f}", ha="center", va="bottom", fontweight="bold")
+        task_data = [m for m in metrics if m["task"] == task]
+        if len(task_data) >= 2:
+            task_labels.append(task.replace("_", "\n"))
+            acc_improvements.append(task_data[-1]["final_accuracy"] - task_data[0]["final_accuracy"])
+            # Efficiency: steps saved (first vs best)
+            first_steps = task_data[0]["num_steps"]
+            best_steps = min(m["num_steps"] for m in task_data)
+            eff_pct = ((first_steps - best_steps) / first_steps * 100) if first_steps > 0 else 0
+            eff_improvements.append(eff_pct)
+            best_accuracies.append(max(m["final_accuracy"] for m in task_data))
+
+    # Left: Best accuracy per task
+    bars1 = axes[0].bar(task_labels, best_accuracies,
+                        color=["#2196F3", "#FF9800", "#4CAF50"][:len(task_labels)],
+                        edgecolor="white", linewidth=1.5)
+    axes[0].axhline(y=0.9, color="red", linestyle="--", alpha=0.7, label="success threshold")
+    axes[0].set_ylabel("Best Accuracy Achieved")
+    axes[0].set_title("Best Accuracy Per Task")
+    axes[0].set_ylim(0, 1.1)
+    axes[0].grid(True, axis="y", alpha=0.3)
+    axes[0].legend()
+    for bar, val in zip(bars1, best_accuracies):
+        axes[0].text(bar.get_x() + bar.get_width() / 2, bar.get_height() + 0.02,
+                     f"{val:.0%}", ha="center", va="bottom", fontweight="bold")
+
+    # Right: Efficiency improvement (% steps saved)
+    bars2 = axes[1].bar(task_labels, eff_improvements,
+                        color=["#2196F3", "#FF9800", "#4CAF50"][:len(task_labels)],
+                        edgecolor="white", linewidth=1.5)
+    axes[1].axhline(y=0, color="black", linewidth=0.8)
+    axes[1].set_ylabel("Steps Saved (%)")
+    axes[1].set_title("Efficiency Improvement (First → Best Episode)")
+    axes[1].grid(True, axis="y", alpha=0.3)
+    for bar, val in zip(bars2, eff_improvements):
+        y_pos = max(bar.get_height() + 1, 2)
+        axes[1].text(bar.get_x() + bar.get_width() / 2, y_pos,
+                     f"{val:.0f}%", ha="center", va="bottom", fontweight="bold")
 
     plt.tight_layout()
     plt.savefig("training/plots/improvement_chart.png", dpi=150, bbox_inches="tight")
@@ -488,9 +598,12 @@ def save_plots(metrics: list[dict]):
     print("Saved: training/plots/improvement_chart.png")
 
     # ── Save raw metrics as JSON ──────────────────────────────────────────────
-    with open("training/plots/metrics.json", "w") as f:
+    timestamp = int(time.time())
+    with open(f"training/plots/metrics_{timestamp}.json", "w") as f:
+        json.dump(metrics, f, indent=2)
+    with open("training/plots/metrics_latest.json", "w") as f:
         json.dump(metrics, f, indent=2)
-    print("Saved: training/plots/metrics.json")
+    print(f"Saved: training/plots/metrics_{timestamp}.json and metrics_latest.json")
 
 # ── Entry Point ───────────────────────────────────────────────────────────────
 

training/update_colab.py

@@ -0,0 +1,173 @@
+import json
+import os
+
+with open("training/trajectory_optimizer.py", "r") as f:
+    traj_code = f.read()
+
+notebook = {
+ "nbformat": 4,
+ "nbformat_minor": 0,
+ "metadata": {
+  "colab": {
+   "provenance": [],
+   "name": "Policy-to-Logic Training"
+  },
+  "kernelspec": {
+   "name": "python3",
+   "display_name": "Python 3"
+  },
+  "language_info": {
+   "name": "python"
+  }
+ },
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Policy-to-Logic RL Environment \u2014 Training Notebook\n",
+    "\n",
+    "This notebook runs the **reward-guided trajectory optimization loop** against the deployed environment.\n",
+    "\n",
+    "**What it does:**\n",
+    "1. Connects to the live HF Spaces environment\n",
+    "2. Runs 8 episodes per task (3 tasks = 24 total episodes)\n",
+    "3. Accumulates high-reward trajectories as few-shot examples\n",
+    "4. Generates training evidence plots (reward curve, accuracy curve, improvement chart)\n",
+    "5. Logs everything to Weights & Biases"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": None,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Cell 1: Install dependencies\n",
+    "!pip install openai requests matplotlib numpy wandb"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": None,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Cell 2: Configuration\n",
+    "import os\n",
+    "\n",
+    "# SET THESE BEFORE RUNNING\n",
+    "HF_TOKEN = \"\"  # Your Hugging Face token with inference access\n",
+    "ENV_URL = \"https://godreign-policy2logic.hf.space\"  # Your deployed environment URL\n",
+    "WANDB_API_KEY = \"\" # Your Wandb API key\n",
+    "\n",
+    "os.environ[\"HF_TOKEN\"] = HF_TOKEN\n",
+    "os.environ[\"ENV_BASE_URL\"] = ENV_URL\n",
+    "if WANDB_API_KEY:\n",
+    "    os.environ[\"WANDB_API_KEY\"] = WANDB_API_KEY\n",
+    "\n",
+    "print(f\"Environment URL: {ENV_URL}\")\n",
+    "print(f\"HF Token set: {'Yes' if HF_TOKEN else 'NO - MUST SET THIS'}\")\n",
+    "print(f\"Wandb Token set: {'Yes' if WANDB_API_KEY else 'NO - WILL PROMPT'}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": None,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Cell 3: Verify environment is reachable\n",
+    "import requests\n",
+    "\n",
+    "r = requests.get(f\"{ENV_URL}/health\")\n",
+    "print(f\"Status: {r.status_code}\")\n",
+    "print(f\"Response: {r.json()}\")\n",
+    "\n",
+    "r2 = requests.get(f\"{ENV_URL}/tasks\")\n",
+    "tasks = r2.json()\n",
+    "print(f\"\\nAvailable tasks: {list(tasks['tasks'].keys())}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": None,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Cell 4: Wandb login — run this before training\n",
+    "import wandb\n",
+    "wandb.login()  # Will prompt for API key if WANDB_API_KEY is not set"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": None,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Cell 5: Training loop implementation (full trajectory_optimizer.py)\n"
+   ] + [line + "\n" for line in traj_code.split("\n")[:-1]] + [traj_code.split("\n")[-1]]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": None,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Cell 6: Run training loop\n",
+    "loop = TrainingLoop(ENV_URL, HF_TOKEN)\n",
+    "metrics = loop.run()\n",
+    "print(f\"\\nTotal episodes run: {len(metrics)}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": None,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Cell 7: Generate plots and display inline\n",
+    "save_plots(metrics)\n",
+    "\n",
+    "from IPython.display import Image, display\n",
+    "display(Image(\"training/plots/reward_curve.png\"))\n",
+    "display(Image(\"training/plots/accuracy_curve.png\"))\n",
+    "display(Image(\"training/plots/improvement_chart.png\"))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": None,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Cell 8: Display wandb run link\n",
+    "print(f\"Wandb run: https://wandb.ai/YOUR_USERNAME/policy-to-logic-rl\")\n",
+    "print(\"Add this link to your README under Deliverables.\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": None,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Cell 9: Download plots to commit to repo\n",
+    "from google.colab import files\n",
+    "\n",
+    "files.download(\"training/plots/reward_curve.png\")\n",
+    "files.download(\"training/plots/accuracy_curve.png\")\n",
+    "files.download(\"training/plots/improvement_chart.png\")\n",
+    "files.download(\"training/plots/metrics_latest.json\")\n",
+    "\n",
+    "print(\"Downloaded. Now commit these files to: training/plots/ in your repo.\")"
+   ]
+  }
+ ]
+}
+
+with open("training/colab_training.ipynb", "w") as f:
+    json.dump(notebook, f, indent=1)
+
+print("Colab Notebook updated successfully")

uv.lock

@@ -449,6 +449,30 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/d5/1f/5f4a3cd9e4440e9d9bc78ad0a91a1c8d46b4d429d5239ebe6793c9fe5c41/fsspec-2026.3.0-py3-none-any.whl", hash = "sha256:d2ceafaad1b3457968ed14efa28798162f1638dbb5d2a6868a2db002a5ee39a4", size = 202595, upload-time = "2026-03-27T19:11:13.595Z" },
 ]
 
+[[package]]
+name = "gitdb"
+version = "4.0.12"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "smmap" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/72/94/63b0fc47eb32792c7ba1fe1b694daec9a63620db1e313033d18140c2320a/gitdb-4.0.12.tar.gz", hash = "sha256:5ef71f855d191a3326fcfbc0d5da835f26b13fbcba60c32c21091c349ffdb571", size = 394684, upload-time = "2025-01-02T07:20:46.413Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/a0/61/5c78b91c3143ed5c14207f463aecfc8f9dbb5092fb2869baf37c273b2705/gitdb-4.0.12-py3-none-any.whl", hash = "sha256:67073e15955400952c6565cc3e707c554a4eea2e428946f7a4c162fab9bd9bcf", size = 62794, upload-time = "2025-01-02T07:20:43.624Z" },
+]
+
+[[package]]
+name = "gitpython"
+version = "3.1.47"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "gitdb" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/c1/bd/50db468e9b1310529a19fce651b3b0e753b5c07954d486cba31bbee9a5d5/gitpython-3.1.47.tar.gz", hash = "sha256:dba27f922bd2b42cb54c87a8ab3cb6beb6bf07f3d564e21ac848913a05a8a3cd", size = 216978, upload-time = "2026-04-22T02:44:44.059Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/f2/c5/a1bc0996af85757903cf2bf444a7824e68e0035ce63fb41d6f76f9def68b/gitpython-3.1.47-py3-none-any.whl", hash = "sha256:489f590edfd6d20571b2c0e72c6a6ac6915ee8b8cd04572330e3842207a78905", size = 209547, upload-time = "2026-04-22T02:44:41.271Z" },
+]
+
 [[package]]
 name = "h11"
 version = "0.16.0"
@@ -1161,6 +1185,15 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/bc/60/5382c03e1970de634027cee8e1b7d39776b778b81812aaf45b694dfe9e28/pillow-12.2.0-pp311-pypy311_pp73-win_amd64.whl", hash = "sha256:bfa9c230d2fe991bed5318a5f119bd6780cda2915cca595393649fc118ab895e", size = 7080946, upload-time = "2026-04-01T14:46:11.734Z" },
 ]
 
+[[package]]
+name = "platformdirs"
+version = "4.9.6"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/9f/4a/0883b8e3802965322523f0b200ecf33d31f10991d0401162f4b23c698b42/platformdirs-4.9.6.tar.gz", hash = "sha256:3bfa75b0ad0db84096ae777218481852c0ebc6c727b3168c1b9e0118e458cf0a", size = 29400, upload-time = "2026-04-09T00:04:10.812Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/75/a6/a0a304dc33b49145b21f4808d763822111e67d1c3a32b524a1baf947b6e1/platformdirs-4.9.6-py3-none-any.whl", hash = "sha256:e61adb1d5e5cb3441b4b7710bea7e4c12250ca49439228cc1021c00dcfac0917", size = 21348, upload-time = "2026-04-09T00:04:09.463Z" },
+]
+
 [[package]]
 name = "pluggy"
 version = "1.6.0"
@@ -1185,6 +1218,7 @@ dependencies = [
     { name = "pydantic" },
     { name = "requests" },
     { name = "uvicorn" },
+    { name = "wandb" },
 ]
 
 [package.optional-dependencies]
@@ -1206,9 +1240,25 @@ requires-dist = [
     { name = "pytest", marker = "extra == 'dev'", specifier = ">=7.0" },
     { name = "requests", specifier = ">=2.25.0" },
     { name = "uvicorn", specifier = ">=0.24.0" },
+    { name = "wandb", specifier = ">=0.16.0" },
 ]
 provides-extras = ["dev"]
 
+[[package]]
+name = "protobuf"
+version = "7.34.1"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/6b/6b/a0e95cad1ad7cc3f2c6821fcab91671bd5b78bd42afb357bb4765f29bc41/protobuf-7.34.1.tar.gz", hash = "sha256:9ce42245e704cc5027be797c1db1eb93184d44d1cdd71811fb2d9b25ad541280", size = 454708, upload-time = "2026-03-20T17:34:47.036Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/ec/11/3325d41e6ee15bf1125654301211247b042563bcc898784351252549a8ad/protobuf-7.34.1-cp310-abi3-macosx_10_9_universal2.whl", hash = "sha256:d8b2cc79c4d8f62b293ad9b11ec3aebce9af481fa73e64556969f7345ebf9fc7", size = 429247, upload-time = "2026-03-20T17:34:37.024Z" },
+    { url = "https://files.pythonhosted.org/packages/eb/9d/aa69df2724ff63efa6f72307b483ce0827f4347cc6d6df24b59e26659fef/protobuf-7.34.1-cp310-abi3-manylinux2014_aarch64.whl", hash = "sha256:5185e0e948d07abe94bb76ec9b8416b604cfe5da6f871d67aad30cbf24c3110b", size = 325753, upload-time = "2026-03-20T17:34:38.751Z" },
+    { url = "https://files.pythonhosted.org/packages/92/e8/d174c91fd48e50101943f042b09af9029064810b734e4160bbe282fa1caa/protobuf-7.34.1-cp310-abi3-manylinux2014_s390x.whl", hash = "sha256:403b093a6e28a960372b44e5eb081775c9b056e816a8029c61231743d63f881a", size = 340198, upload-time = "2026-03-20T17:34:39.871Z" },
+    { url = "https://files.pythonhosted.org/packages/53/1b/3b431694a4dc6d37b9f653f0c64b0a0d9ec074ee810710c0c3da21d67ba7/protobuf-7.34.1-cp310-abi3-manylinux2014_x86_64.whl", hash = "sha256:8ff40ce8cd688f7265326b38d5a1bed9bfdf5e6723d49961432f83e21d5713e4", size = 324267, upload-time = "2026-03-20T17:34:41.1Z" },
+    { url = "https://files.pythonhosted.org/packages/85/29/64de04a0ac142fb685fd09999bc3d337943fb386f3a0ec57f92fd8203f97/protobuf-7.34.1-cp310-abi3-win32.whl", hash = "sha256:34b84ce27680df7cca9f231043ada0daa55d0c44a2ddfaa58ec1d0d89d8bf60a", size = 426628, upload-time = "2026-03-20T17:34:42.536Z" },
+    { url = "https://files.pythonhosted.org/packages/4d/87/cb5e585192a22b8bd457df5a2c16a75ea0db9674c3a0a39fc9347d84e075/protobuf-7.34.1-cp310-abi3-win_amd64.whl", hash = "sha256:e97b55646e6ce5cbb0954a8c28cd39a5869b59090dfaa7df4598a7fba869468c", size = 437901, upload-time = "2026-03-20T17:34:44.112Z" },
+    { url = "https://files.pythonhosted.org/packages/88/95/608f665226bca68b736b79e457fded9a2a38c4f4379a4a7614303d9db3bc/protobuf-7.34.1-py3-none-any.whl", hash = "sha256:bb3812cd53aefea2b028ef42bd780f5b96407247f20c6ef7c679807e9d188f11", size = 170715, upload-time = "2026-03-20T17:34:45.384Z" },
+]
+
 [[package]]
 name = "pydantic"
 version = "2.13.3"
@@ -1480,6 +1530,19 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/82/3b/64d4899d73f91ba49a8c18a8ff3f0ea8f1c1d75481760df8c68ef5235bf5/rich-15.0.0-py3-none-any.whl", hash = "sha256:33bd4ef74232fb73fe9279a257718407f169c09b78a87ad3d296f548e27de0bb", size = 310654, upload-time = "2026-04-12T08:24:02.83Z" },
 ]
 
+[[package]]
+name = "sentry-sdk"
+version = "2.58.0"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "certifi" },
+    { name = "urllib3" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/26/b3/fb8291170d0e844173164709fc0fa0c221ed75a5da740c8746f2a83b4eb1/sentry_sdk-2.58.0.tar.gz", hash = "sha256:c1144d947352d54e5b7daa63596d9f848adf684989c06c4f5a659f0c85a18f6f", size = 438764, upload-time = "2026-04-13T17:23:26.265Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/fa/eb/d875669993b762556ae8b2efd86219943b4c0864d22204d622a9aee3052b/sentry_sdk-2.58.0-py2.py3-none-any.whl", hash = "sha256:688d1c704ddecf382ea3326f21a67453d4caa95592d722b7c780a36a9d23109e", size = 460919, upload-time = "2026-04-13T17:23:24.675Z" },
+]
+
 [[package]]
 name = "shellingham"
 version = "1.5.4"
@@ -1498,6 +1561,15 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/b7/ce/149a00dd41f10bc29e5921b496af8b574d8413afcd5e30dfa0ed46c2cc5e/six-1.17.0-py2.py3-none-any.whl", hash = "sha256:4721f391ed90541fddacab5acf947aa0d3dc7d27b2e1e8eda2be8970586c3274", size = 11050, upload-time = "2024-12-04T17:35:26.475Z" },
 ]
 
+[[package]]
+name = "smmap"
+version = "5.0.3"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/1f/ea/49c993d6dfdd7338c9b1000a0f36817ed7ec84577ae2e52f890d1a4ff909/smmap-5.0.3.tar.gz", hash = "sha256:4d9debb8b99007ae47165abc08670bd74cb74b5227dda7f643eccc4e9eb5642c", size = 22506, upload-time = "2026-03-09T03:43:26.1Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/c1/d4/59e74daffcb57a07668852eeeb6035af9f32cbfd7a1d2511f17d2fe6a738/smmap-5.0.3-py3-none-any.whl", hash = "sha256:c106e05d5a61449cf6ba9a1e650227ecfb141590d2a98412103ff35d89fc7b2f", size = 24390, upload-time = "2026-03-09T03:43:24.361Z" },
+]
+
 [[package]]
 name = "sniffio"
 version = "1.3.1"
@@ -1644,3 +1716,32 @@ sdist = { url = "https://files.pythonhosted.org/packages/1f/93/041fca8274050e40e
 wheels = [
     { url = "https://files.pythonhosted.org/packages/31/a3/5b1562db76a5a488274b2332a97199b32d0442aca0ed193697fd47786316/uvicorn-0.46.0-py3-none-any.whl", hash = "sha256:bbebbcbed972d162afca128605223022bedd345b7bc7855ce66deb31487a9048", size = 70926, upload-time = "2026-04-23T07:15:58.355Z" },
 ]
+
+[[package]]
+name = "wandb"
+version = "0.26.1"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "click" },
+    { name = "gitpython" },
+    { name = "packaging" },
+    { name = "platformdirs" },
+    { name = "protobuf" },
+    { name = "pydantic" },
+    { name = "pyyaml" },
+    { name = "requests" },
+    { name = "sentry-sdk" },
+    { name = "typing-extensions" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/6a/a4/72a6640e1f566e81f184a426e3e45298d4c6672664de41adb7eb6f64370a/wandb-0.26.1.tar.gz", hash = "sha256:eef2dbaea06f0b1c0cdc5d76f544ae4c2b8848fc512442a00bd59f0502fc8aa1", size = 42159814, upload-time = "2026-04-23T16:27:34.033Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/8c/09/3296235f3906e904f06f2df29eed4d672fb23c0932c9486e2af64f2f2a66/wandb-0.26.1-py3-none-macosx_12_0_arm64.whl", hash = "sha256:2955fe190c005fb83ee6d73f066c8a33f09f3212a1f2eb53faa6581440e456be", size = 24857204, upload-time = "2026-04-23T16:26:58.576Z" },
+    { url = "https://files.pythonhosted.org/packages/a1/ad/e39ca3086534129e42208ba00ed2c6247ce425f890219eeec33b4f162864/wandb-0.26.1-py3-none-macosx_12_0_x86_64.whl", hash = "sha256:55d91cabde98162d7116a5e19ddd052bd9848556243f1da4cbb9ffb7ad435bfc", size = 26014649, upload-time = "2026-04-23T16:27:02.559Z" },
+    { url = "https://files.pythonhosted.org/packages/56/af/400d84a3bdce0b062b4baa70acb6becd2c8018697f4fbf5af9a9e1e406e5/wandb-0.26.1-py3-none-manylinux_2_28_aarch64.whl", hash = "sha256:7c78bc2454cfe1ffa1c3a256060a387356eed8a4488e024d9d2eba8f2b5bd51d", size = 25421317, upload-time = "2026-04-23T16:27:06.411Z" },
+    { url = "https://files.pythonhosted.org/packages/7b/e9/b4bf8f3509dcea1cec52233a38991459654635b5a8e6a494eb912e1b9cfb/wandb-0.26.1-py3-none-manylinux_2_28_x86_64.whl", hash = "sha256:a2c8eeec8706dcd2872e69c3b4d20ec523082fdb4440295491556e219ad2aa67", size = 27192831, upload-time = "2026-04-23T16:27:10.308Z" },
+    { url = "https://files.pythonhosted.org/packages/62/cf/4a6dce0c782223ef0eeea7139daee73418a7322befcf083512c31cebaa18/wandb-0.26.1-py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:2fa768ee0636a569afb7541cf996e56309c47070566a38916823f94e02afe586", size = 25593326, upload-time = "2026-04-23T16:27:14.259Z" },
+    { url = "https://files.pythonhosted.org/packages/df/99/58c3d8c36ae8e2b7d70bf6493eb5daa1cca0231a04b025717b4cd1a78f1e/wandb-0.26.1-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:5854928725cfeff1f284d5c043cd353f810e5da02eead2c120ef5056ad026fea", size = 27535542, upload-time = "2026-04-23T16:27:18.473Z" },
+    { url = "https://files.pythonhosted.org/packages/7c/d0/4e846ffc1d0cc435518dfa581ce73ac82cfd0ebbf35f3853c9277f632e5f/wandb-0.26.1-py3-none-win32.whl", hash = "sha256:5c2bd44e575ae9944e2764d1aaa031461178276bf2636d5558399c2816ef5cfe", size = 24968151, upload-time = "2026-04-23T16:27:22.086Z" },
+    { url = "https://files.pythonhosted.org/packages/e3/9b/487413eaccefdb58799a226726e24b486e9192d2671c75a4550c160aba23/wandb-0.26.1-py3-none-win_amd64.whl", hash = "sha256:5817785467d3f1676f1812ec19a89f77f6e56dfe67d9f47080075af95f705d3e", size = 24968155, upload-time = "2026-04-23T16:27:25.731Z" },
+    { url = "https://files.pythonhosted.org/packages/04/dc/5baf3e99b3eeb709d6f75124b5bec8cb73d4b38d2b10df7fdcfde4966200/wandb-0.26.1-py3-none-win_arm64.whl", hash = "sha256:f848b7744f896bc04cabbb28360a2814d1551a91fa2c456243e06435729c8a2e", size = 22912416, upload-time = "2026-04-23T16:27:29.456Z" },
+]