Refresh training notebook clone step

SupplyMind_Training_Run.ipynb

@@ -1,295 +1,297 @@
-{
- "cells": [
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "# SupplyMind Training Run\n",
-    "\n",
-    "This notebook is the compact, judge-runnable training path for SupplyMind: environment smoke test → SFT warm-start → GRPO from SFT → held-out evaluation.\n",
-    "\n",
-    "Default settings train the **center** role on the easy task so the notebook can run quickly. Change `ROLE` to `\"warehouse\"` or `TASK_ID` to `\"v2_train_medium\"` / `\"v2_train_hard\"` for a larger run."
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## 1. Setup"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-                "!pip -q uninstall -y torchao\n",
-                "!pip -q install \"torch\" \"transformers>=4.45.0\" \"trl>=0.12.0\" \"peft>=0.13.0\" accelerate datasets bitsandbytes huggingface_hub pydantic pyyaml matplotlib"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import os\n",
-    "from pathlib import Path\n",
-    "\n",
-    "if not Path(\"supplymind\").exists():\n",
-    "    !git clone -q https://huggingface.co/spaces/rishavutk/supplymind supplymind\n",
-    "\n",
-    "%cd /content/supplymind\n",
-    "!pip -q install -e ."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from huggingface_hub import notebook_login, whoami\n",
-    "\n",
-    "notebook_login()\n",
-    "HF_NAMESPACE = whoami()[\"name\"]\n",
-    "print(\"Using HF namespace:\", HF_NAMESPACE)"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## 2. Controls\n",
-    "\n",
-    "Change only these values for quick variants. `ROLE` controls which policy is trained; `TASK_ID` controls easy/medium/hard world generation."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "ROLE = \"center\"  # \"center\" or \"warehouse\"\n",
-    "TASK_ID = \"v2_train_easy\"  # also: \"v2_train_medium\", \"v2_train_hard\"\n",
-    "TRAIN_SEEDS = \"101,113,127\"\n",
-    "EVAL_SEEDS = \"131,149,163\"\n",
-    "SFT_STEPS = 20\n",
-    "GRPO_STEPS = 20\n",
-    "MAX_COMPLETION_LENGTH = 256\n",
-    "\n",
-    "SFT_ADAPTER_ID = f\"{HF_NAMESPACE}/supplymind-{ROLE}-qwen-0.5b-sft-notebook\"\n",
-    "GRPO_ADAPTER_ID = f\"{HF_NAMESPACE}/supplymind-{ROLE}-qwen-0.5b-grpo-notebook\"\n",
-    "\n",
-    "print({\n",
-    "    \"role\": ROLE,\n",
-    "    \"task_id\": TASK_ID,\n",
-    "    \"train_seeds\": TRAIN_SEEDS,\n",
-    "    \"eval_seeds\": EVAL_SEEDS,\n",
-    "    \"sft_adapter\": SFT_ADAPTER_ID,\n",
-    "    \"grpo_adapter\": GRPO_ADAPTER_ID,\n",
-    "})"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## 3. Environment Smoke Test"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import json\n",
-    "import os\n",
-    "import sys\n",
-    "\n",
-    "sys.path.insert(0, \"/content/supplymind/src\")\n",
-    "os.environ[\"SUPPLYMIND_REWARD_CONFIG\"] = \"/content/supplymind/configs/supplymind_v2_rewards.yaml\"\n",
-    "\n",
-                "from supplymind_env_v2.environment import V2SupplyMindEnv\n",
-                "from supplymind_env_v2.models import V2JointAction\n",
-    "\n",
-    "env = V2SupplyMindEnv(default_task_id=TASK_ID)\n",
-    "obs = env.reset_internal(TASK_ID, 131)\n",
-    "print(\"observation keys:\", sorted(obs.model_dump(mode=\"json\").keys()))\n",
-    "print(\"round:\", obs.round_index, \"warehouses:\", list(obs.warehouses.keys()))\n",
-    "\n",
-    "empty_action = {\n",
-    "    \"warehouse_actions\": {},\n",
-    "    \"central_action\": {\n",
-    "        \"central_procurements\": [],\n",
-    "        \"central_liquidations\": [],\n",
-    "        \"central_replenishments\": [],\n",
-    "        \"inventory_transfer_proposals\": [],\n",
-    "        \"offer_matches\": [],\n",
-    "    },\n",
-    "}\n",
-                "result = env.step(V2JointAction.model_validate(empty_action))\n",
-    "print(\"step reward:\", result.reward.step_reward)\n",
-    "print(\"done:\", result.done)\n",
-    "print(\"info keys:\", sorted(result.info.keys()))"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## 4. SFT Warm-Start\n",
-    "\n",
-    "SFT teaches the model the action JSON shape and a reasonable heuristic policy. For the warehouse role, the notebook enables the conservative SFT flag to reduce invalid or overactive actions."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "warehouse_flags = \"--warehouse-conservative-sft --warehouse-signal-limit 2\" if ROLE == \"warehouse\" else \"\"\n",
-    "\n",
-    "!python scripts/hf_sft_supplymind_roles.py \\\n",
-    "  --role {ROLE} \\\n",
-    "  --task-id {TASK_ID} \\\n",
-    "  --seeds {TRAIN_SEEDS} \\\n",
-    "  --max-steps {SFT_STEPS} \\\n",
-    "  --hub-model-id {SFT_ADAPTER_ID} \\\n",
-    "  --output-dir outputs/{ROLE}-sft-notebook \\\n",
-    "  {warehouse_flags}"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## 5. GRPO From SFT\n",
-    "\n",
-    "The GRPO script routes reward by role: center updates from center reward deltas, warehouse updates from warehouse reward deltas, while global reward is logged for audit."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "!python scripts/hf_train_supplymind_roles.py \\\n",
-    "  --role {ROLE} \\\n",
-    "  --task-id {TASK_ID} \\\n",
-    "  --seeds {TRAIN_SEEDS} \\\n",
-    "  --max-steps {GRPO_STEPS} \\\n",
-    "  --max-completion-length {MAX_COMPLETION_LENGTH} \\\n",
-    "  --init-adapter-id {SFT_ADAPTER_ID} \\\n",
-    "  --hub-model-id {GRPO_ADAPTER_ID} \\\n",
-    "  --output-dir outputs/{ROLE}-grpo-notebook"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## 6. Held-Out Evaluation\n",
-    "\n",
-    "Evaluate base Qwen, SFT, and GRPO on held-out seeds. The role score is the training target; global score is the environment-level audit metric."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "!python scripts/hf_eval_supplymind_adapters.py \\\n",
-    "  --role {ROLE} \\\n",
-    "  --task-id {TASK_ID} \\\n",
-    "  --seeds {EVAL_SEEDS} \\\n",
-    "  --sft-adapter-id {SFT_ADAPTER_ID} \\\n",
-    "  --grpo-adapter-id {GRPO_ADAPTER_ID} \\\n",
-    "  --max-new-tokens {MAX_COMPLETION_LENGTH} | tee outputs/{ROLE}-eval-notebook.log"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import json\n",
-    "import re\n",
-    "from pathlib import Path\n",
-    "\n",
-    "import matplotlib.pyplot as plt\n",
-    "import pandas as pd\n",
-    "\n",
-    "log_path = Path(f\"outputs/{ROLE}-eval-notebook.log\")\n",
-    "rows = []\n",
-    "for line in log_path.read_text(encoding=\"utf-8\", errors=\"ignore\").splitlines():\n",
-    "    line = line.strip()\n",
-    "    if not line.startswith(\"{\"):\n",
-    "        continue\n",
-    "    try:\n",
-    "        payload = json.loads(line)\n",
-    "    except json.JSONDecodeError:\n",
-    "        continue\n",
-    "    if payload.get(\"message\") == \"eval_done\":\n",
-    "        evaluations = {key: payload[key] for key in (\"base\", \"sft\", \"grpo\") if key in payload}\n",
-    "        for label, item in evaluations.items():\n",
-    "            role_score_key = \"mean_center_role_score\" if ROLE == \"center\" else \"mean_warehouse_role_score\"\n",
-    "            rows.append({\n",
-    "                \"policy\": label,\n",
-    "                \"global_score\": item.get(\"mean_global_score\"),\n",
-    "                \"role_score\": item.get(role_score_key),\n",
-    "                \"raw_reward\": item.get(\"mean_raw_reward\"),\n",
-    "                \"invalid_payloads\": item.get(\"invalid_payloads\"),\n",
-    "                \"invalid_actions\": item.get(\"invalid_actions\"),\n",
-    "            })\n",
-    "\n",
-    "df = pd.DataFrame(rows)\n",
-    "display(df)\n",
-    "\n",
-    "if not df.empty:\n",
-    "    fig, axes = plt.subplots(1, 2, figsize=(10, 4))\n",
-    "    df.plot.bar(x=\"policy\", y=\"role_score\", ax=axes[0], legend=False, color=\"#2563eb\")\n",
-    "    axes[0].set_title(f\"{ROLE} role score\")\n",
-    "    axes[0].set_ylim(0, 1)\n",
-    "    df.plot.bar(x=\"policy\", y=[\"invalid_payloads\", \"invalid_actions\"], ax=axes[1], color=[\"#dc2626\", \"#f59e0b\"])\n",
-    "    axes[1].set_title(\"Invalid outputs\")\n",
-    "    plt.tight_layout()\n",
-    "    plt.show()"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## Rerun For The Other Role\n",
-    "\n",
-    "To train warehouses instead of center, change `ROLE = \"warehouse\"` in the controls cell and rerun sections 4-6. To try larger worlds, change `TASK_ID` to `v2_train_medium` or `v2_train_hard`."
-   ]
-  }
- ],
- "metadata": {
-  "accelerator": "GPU",
-  "colab": {
-   "gpuType": "T4",
-   "provenance": []
-  },
-  "kernelspec": {
-   "display_name": "Python 3",
-   "language": "python",
-   "name": "python3"
-  },
-  "language_info": {
-   "name": "python",
-   "version": "3.11"
-  }
- },
- "nbformat": 4,
- "nbformat_minor": 5
-}
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# SupplyMind Training Run\n",
+    "\n",
+    "This notebook is the compact, judge-runnable training path for SupplyMind: environment smoke test → SFT warm-start → GRPO from SFT → held-out evaluation.\n",
+    "\n",
+    "Default settings train the **center** role on the easy task so the notebook can run quickly. Change `ROLE` to `\"warehouse\"` or `TASK_ID` to `\"v2_train_medium\"` / `\"v2_train_hard\"` for a larger run."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 1. Setup"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "!pip -q uninstall -y torchao\n",
+    "!pip -q install \"torch\" \"transformers>=4.45.0\" \"trl>=0.12.0\" \"peft>=0.13.0\" accelerate datasets bitsandbytes huggingface_hub pydantic pyyaml matplotlib"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "from pathlib import Path\n",
+    "\n",
+    "if not Path(\"supplymind\").exists():\n",
+    "    !git clone -q https://huggingface.co/spaces/rishavutk/supplymind supplymind\n",
+    "else:\n",
+    "    !git -C supplymind pull -q\n",
+    "\n",
+    "%cd /content/supplymind\n",
+    "!pip -q install -e .\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from huggingface_hub import notebook_login, whoami\n",
+    "\n",
+    "notebook_login()\n",
+    "HF_NAMESPACE = whoami()[\"name\"]\n",
+    "print(\"Using HF namespace:\", HF_NAMESPACE)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 2. Controls\n",
+    "\n",
+    "Change only these values for quick variants. `ROLE` controls which policy is trained; `TASK_ID` controls easy/medium/hard world generation."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "ROLE = \"center\"  # \"center\" or \"warehouse\"\n",
+    "TASK_ID = \"v2_train_easy\"  # also: \"v2_train_medium\", \"v2_train_hard\"\n",
+    "TRAIN_SEEDS = \"101,113,127\"\n",
+    "EVAL_SEEDS = \"131,149,163\"\n",
+    "SFT_STEPS = 20\n",
+    "GRPO_STEPS = 20\n",
+    "MAX_COMPLETION_LENGTH = 256\n",
+    "\n",
+    "SFT_ADAPTER_ID = f\"{HF_NAMESPACE}/supplymind-{ROLE}-qwen-0.5b-sft-notebook\"\n",
+    "GRPO_ADAPTER_ID = f\"{HF_NAMESPACE}/supplymind-{ROLE}-qwen-0.5b-grpo-notebook\"\n",
+    "\n",
+    "print({\n",
+    "    \"role\": ROLE,\n",
+    "    \"task_id\": TASK_ID,\n",
+    "    \"train_seeds\": TRAIN_SEEDS,\n",
+    "    \"eval_seeds\": EVAL_SEEDS,\n",
+    "    \"sft_adapter\": SFT_ADAPTER_ID,\n",
+    "    \"grpo_adapter\": GRPO_ADAPTER_ID,\n",
+    "})"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 3. Environment Smoke Test"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import json\n",
+    "import os\n",
+    "import sys\n",
+    "\n",
+    "sys.path.insert(0, \"/content/supplymind/src\")\n",
+    "os.environ[\"SUPPLYMIND_REWARD_CONFIG\"] = \"/content/supplymind/configs/supplymind_v2_rewards.yaml\"\n",
+    "\n",
+    "from supplymind_env_v2.environment import V2SupplyMindEnv\n",
+    "from supplymind_env_v2.models import V2JointAction\n",
+    "\n",
+    "env = V2SupplyMindEnv(default_task_id=TASK_ID)\n",
+    "obs = env.reset_internal(TASK_ID, 131)\n",
+    "print(\"observation keys:\", sorted(obs.model_dump(mode=\"json\").keys()))\n",
+    "print(\"round:\", obs.round_index, \"warehouses:\", list(obs.warehouses.keys()))\n",
+    "\n",
+    "empty_action = {\n",
+    "    \"warehouse_actions\": {},\n",
+    "    \"central_action\": {\n",
+    "        \"central_procurements\": [],\n",
+    "        \"central_liquidations\": [],\n",
+    "        \"central_replenishments\": [],\n",
+    "        \"inventory_transfer_proposals\": [],\n",
+    "        \"offer_matches\": [],\n",
+    "    },\n",
+    "}\n",
+    "result = env.step(V2JointAction.model_validate(empty_action))\n",
+    "print(\"step reward:\", result.reward.step_reward)\n",
+    "print(\"done:\", result.done)\n",
+    "print(\"info keys:\", sorted(result.info.keys()))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 4. SFT Warm-Start\n",
+    "\n",
+    "SFT teaches the model the action JSON shape and a reasonable heuristic policy. For the warehouse role, the notebook enables the conservative SFT flag to reduce invalid or overactive actions."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "warehouse_flags = \"--warehouse-conservative-sft --warehouse-signal-limit 2\" if ROLE == \"warehouse\" else \"\"\n",
+    "\n",
+    "!python scripts/hf_sft_supplymind_roles.py \\\n",
+    "  --role {ROLE} \\\n",
+    "  --task-id {TASK_ID} \\\n",
+    "  --seeds {TRAIN_SEEDS} \\\n",
+    "  --max-steps {SFT_STEPS} \\\n",
+    "  --hub-model-id {SFT_ADAPTER_ID} \\\n",
+    "  --output-dir outputs/{ROLE}-sft-notebook \\\n",
+    "  {warehouse_flags}"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 5. GRPO From SFT\n",
+    "\n",
+    "The GRPO script routes reward by role: center updates from center reward deltas, warehouse updates from warehouse reward deltas, while global reward is logged for audit."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "!python scripts/hf_train_supplymind_roles.py \\\n",
+    "  --role {ROLE} \\\n",
+    "  --task-id {TASK_ID} \\\n",
+    "  --seeds {TRAIN_SEEDS} \\\n",
+    "  --max-steps {GRPO_STEPS} \\\n",
+    "  --max-completion-length {MAX_COMPLETION_LENGTH} \\\n",
+    "  --init-adapter-id {SFT_ADAPTER_ID} \\\n",
+    "  --hub-model-id {GRPO_ADAPTER_ID} \\\n",
+    "  --output-dir outputs/{ROLE}-grpo-notebook"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 6. Held-Out Evaluation\n",
+    "\n",
+    "Evaluate base Qwen, SFT, and GRPO on held-out seeds. The role score is the training target; global score is the environment-level audit metric."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "!python scripts/hf_eval_supplymind_adapters.py \\\n",
+    "  --role {ROLE} \\\n",
+    "  --task-id {TASK_ID} \\\n",
+    "  --seeds {EVAL_SEEDS} \\\n",
+    "  --sft-adapter-id {SFT_ADAPTER_ID} \\\n",
+    "  --grpo-adapter-id {GRPO_ADAPTER_ID} \\\n",
+    "  --max-new-tokens {MAX_COMPLETION_LENGTH} | tee outputs/{ROLE}-eval-notebook.log"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import json\n",
+    "import re\n",
+    "from pathlib import Path\n",
+    "\n",
+    "import matplotlib.pyplot as plt\n",
+    "import pandas as pd\n",
+    "\n",
+    "log_path = Path(f\"outputs/{ROLE}-eval-notebook.log\")\n",
+    "rows = []\n",
+    "for line in log_path.read_text(encoding=\"utf-8\", errors=\"ignore\").splitlines():\n",
+    "    line = line.strip()\n",
+    "    if not line.startswith(\"{\"):\n",
+    "        continue\n",
+    "    try:\n",
+    "        payload = json.loads(line)\n",
+    "    except json.JSONDecodeError:\n",
+    "        continue\n",
+    "    if payload.get(\"message\") == \"eval_done\":\n",
+    "        evaluations = {key: payload[key] for key in (\"base\", \"sft\", \"grpo\") if key in payload}\n",
+    "        for label, item in evaluations.items():\n",
+    "            role_score_key = \"mean_center_role_score\" if ROLE == \"center\" else \"mean_warehouse_role_score\"\n",
+    "            rows.append({\n",
+    "                \"policy\": label,\n",
+    "                \"global_score\": item.get(\"mean_global_score\"),\n",
+    "                \"role_score\": item.get(role_score_key),\n",
+    "                \"raw_reward\": item.get(\"mean_raw_reward\"),\n",
+    "                \"invalid_payloads\": item.get(\"invalid_payloads\"),\n",
+    "                \"invalid_actions\": item.get(\"invalid_actions\"),\n",
+    "            })\n",
+    "\n",
+    "df = pd.DataFrame(rows)\n",
+    "display(df)\n",
+    "\n",
+    "if not df.empty:\n",
+    "    fig, axes = plt.subplots(1, 2, figsize=(10, 4))\n",
+    "    df.plot.bar(x=\"policy\", y=\"role_score\", ax=axes[0], legend=False, color=\"#2563eb\")\n",
+    "    axes[0].set_title(f\"{ROLE} role score\")\n",
+    "    axes[0].set_ylim(0, 1)\n",
+    "    df.plot.bar(x=\"policy\", y=[\"invalid_payloads\", \"invalid_actions\"], ax=axes[1], color=[\"#dc2626\", \"#f59e0b\"])\n",
+    "    axes[1].set_title(\"Invalid outputs\")\n",
+    "    plt.tight_layout()\n",
+    "    plt.show()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Rerun For The Other Role\n",
+    "\n",
+    "To train warehouses instead of center, change `ROLE = \"warehouse\"` in the controls cell and rerun sections 4-6. To try larger worlds, change `TASK_ID` to `v2_train_medium` or `v2_train_hard`."
+   ]
+  }
+ ],
+ "metadata": {
+  "accelerator": "GPU",
+  "colab": {
+   "gpuType": "T4",
+   "provenance": []
+  },
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3.11"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}