Match eval sampling to training, log all I/O, single round

- Drop greedy at eval; always sample with temperature=1.0, top_p=0.95 so
eval reflects the same distribution training optimized
- Add IO logging: every prompt/response written to plots/io_log.jsonl
with phase tag (before / train_roundN / after) for inspection
- NUM_ROUNDS = 1 to iterate quickly while debugging the no-transfer issue

Made-with: Cursor

training/train_grpo.ipynb

@@ -25,7 +25,9 @@
     },
     {
       "cell_type": "code",
+      "execution_count": null,
       "metadata": {},
+      "outputs": [],
       "source": [
         "# Cell 1: Install dependencies (quote versions — zsh treats `>` as redirect otherwise)\n",
         "!pip install -q torch torchvision torchaudio\n",
@@ -34,13 +36,13 @@
         "!pip install -q \"typing_extensions>=4.13.0\" pydantic httpx\n",
         "!pip install -q \"openenv-core[core]>=0.2.2\"\n",
         "!pip install -q flash-attn --no-build-isolation || echo \"flash-attn install skipped; will use sdpa\""
-      ],
-      "execution_count": null,
-      "outputs": []
+      ]
     },
     {
       "cell_type": "code",
+      "execution_count": null,
       "metadata": {},
+      "outputs": [],
       "source": [
         "# Cell 2: Resolve repo path (Colab: fresh clone. Local: auto-detect project root)\n",
         "import os\n",
@@ -116,13 +118,13 @@
         "print(f\"Branch: {REPO_BRANCH}\")\n",
         "print(f\"Commit: {commit}\")\n",
         "print(f\"Plots dir: {PLOTS_DIR}\")"
-      ],
-      "execution_count": null,
-      "outputs": []
+      ]
     },
     {
       "cell_type": "code",
+      "execution_count": null,
       "metadata": {},
+      "outputs": [],
       "source": [
         "# Cell 3: Imports (with runtime validation)\n",
         "import json, random, time, textwrap, copy, os, sys\n",
@@ -176,9 +178,7 @@
         "import ast\n",
         "ast.parse(\"def _t(x: int) -> str: return f'{x}'\")\n",
         "print(\"OK: ast.parse (syntax check)\")"
-      ],
-      "execution_count": null,
-      "outputs": []
+      ]
     },
     {
       "cell_type": "markdown",
@@ -191,7 +191,9 @@
     },
     {
       "cell_type": "code",
+      "execution_count": null,
       "metadata": {},
+      "outputs": [],
       "source": [
         "# Cell 4: Define heuristic agents + episode runner\n",
         "_rng = random.Random(42)\n",
@@ -267,13 +269,13 @@
         "            \"rewards\": rewards, \"energies\": energies}\n",
         "\n",
         "print(\"Agents and episode runner defined.\")"
-      ],
-      "execution_count": null,
-      "outputs": []
+      ]
     },
     {
       "cell_type": "code",
+      "execution_count": null,
       "metadata": {},
+      "outputs": [],
       "source": [
         "# Cell 5: Run baselines (safe)\n",
         "print(\"Running heuristic baselines (5 agents × 3 tasks)...\")\n",
@@ -308,13 +310,13 @@
         "for name in BASELINE_AGENTS:\n",
         "    scores = [baseline_results[name][t][\"grader_score\"] for t in TASKS]\n",
         "    print(f\"{name:<14s} {scores[0]:>10.4f} {scores[1]:>12.4f} {scores[2]:>14.4f} {sum(scores)/3:>8.4f}\")"
-      ],
-      "execution_count": null,
-      "outputs": []
+      ]
     },
     {
       "cell_type": "code",
+      "execution_count": null,
       "metadata": {},
+      "outputs": [],
       "source": [
         "# Cell 6: Baseline plots\n",
         "fig, axes = plt.subplots(1, 3, figsize=(16, 5), sharey=True)\n",
@@ -332,9 +334,7 @@
         "fig.tight_layout()\n",
         "fig.savefig(f\"{PLOTS_DIR}/baseline_leaderboard.png\", dpi=150, bbox_inches='tight')\n",
         "plt.show()"
-      ],
-      "execution_count": null,
-      "outputs": []
+      ]
     },
     {
       "cell_type": "markdown",
@@ -347,7 +347,9 @@
     },
     {
       "cell_type": "code",
+      "execution_count": null,
       "metadata": {},
+      "outputs": [],
       "source": [
         "# Cell 7: Load model (Qwen2.5-3B bf16 on CUDA + flash-attn-2; fp16/fp32 fallback)\n",
         "from transformers import AutoTokenizer, AutoModelForCausalLM\n",
@@ -391,13 +393,13 @@
         "print(f\"Model loaded. dtype={next(model.parameters()).dtype} device={next(model.parameters()).device}\")\n",
         "if torch.cuda.is_available():\n",
         "    print(f\"CUDA memory: {torch.cuda.memory_allocated()/1e9:.2f} GB\")"
-      ],
-      "execution_count": null,
-      "outputs": []
+      ]
     },
     {
       "cell_type": "code",
+      "execution_count": null,
       "metadata": {},
+      "outputs": [],
       "source": [
         "# Cell 8: LLM agent functions\n",
         "_SYSTEM_BASE = textwrap.dedent(\"\"\"\\\n",
@@ -547,18 +549,29 @@
         "\n",
         "def _batched_generate(mdl, tok, prompts, eval=False, max_new_tokens=512):\n",
         "    enc = tok(prompts, return_tensors=\"pt\", padding=True, truncation=False).to(_infer_model_device(mdl))\n",
-        "    gen_kwargs = dict(max_new_tokens=max_new_tokens, pad_token_id=tok.pad_token_id)\n",
-        "    if eval:\n",
-        "        gen_kwargs.update(do_sample=False)\n",
-        "    else:\n",
-        "        gen_kwargs.update(do_sample=True, temperature=1.0, top_p=0.95)\n",
+        "    gen_kwargs = dict(\n",
+        "        max_new_tokens=max_new_tokens,\n",
+        "        pad_token_id=tok.pad_token_id,\n",
+        "        do_sample=True, temperature=1.0, top_p=0.95,\n",
+        "    )\n",
         "    with torch.no_grad():\n",
         "        out = mdl.generate(**enc, **gen_kwargs)\n",
         "    resps = tok.batch_decode(out[:, enc[\"input_ids\"].shape[1]:], skip_special_tokens=True)\n",
         "    return resps, enc[\"input_ids\"].shape[1]\n",
         "\n",
         "\n",
-        "def run_llm_episodes_batched(mdl, tok, tasks_seeds, verbose=True, eval=False, system=None):\n",
+        "IO_LOG_PATH = os.path.join(PLOTS_DIR, \"io_log.jsonl\")\n",
+        "open(IO_LOG_PATH, \"w\").close()  # truncate\n",
+        "\n",
+        "\n",
+        "def _log_io(tag, ep_idx, day, task, seed, prompt, response):\n",
+        "    rec = {\"tag\": tag, \"ep\": ep_idx, \"day\": day, \"task\": task, \"seed\": seed,\n",
+        "           \"prompt\": prompt, \"response\": response}\n",
+        "    with open(IO_LOG_PATH, \"a\") as f:\n",
+        "        f.write(json.dumps(rec) + \"\\n\")\n",
+        "\n",
+        "\n",
+        "def run_llm_episodes_batched(mdl, tok, tasks_seeds, verbose=True, eval=False, system=None, log_tag=None):\n",
         "    \"\"\"Run N episodes in parallel. tasks_seeds: list of (task, seed). One batched generate per day.\"\"\"\n",
         "    sys_prompt = system or (SYSTEM_PROMPT_EVAL if eval else SYSTEM_PROMPT_TRAIN)\n",
         "    n = len(tasks_seeds)\n",
@@ -588,6 +601,9 @@
         "                actions_by_idx[i] = parse_model_output(resps[j])\n",
         "                pairs[i].append({\"prompt\": prompts[j], \"response\": resps[j],\n",
         "                                 \"step\": len(rewards[i])})\n",
+        "                if log_tag is not None:\n",
+        "                    t, s = tasks_seeds[i]\n",
+        "                    _log_io(log_tag, i, day, t, s, prompts[j], resps[j])\n",
         "\n",
         "        for i in range(n):\n",
         "            if done_mask[i] or i not in actions_by_idx:\n",
@@ -626,9 +642,7 @@
         "\n",
         "\n",
         "print(\"LLM agent functions defined (batched).\")"
-      ],
-      "execution_count": null,
-      "outputs": []
+      ]
     },
     {
       "cell_type": "markdown",
@@ -641,23 +655,23 @@
     },
     {
       "cell_type": "code",
+      "execution_count": null,
       "metadata": {},
+      "outputs": [],
       "source": [
         "# Cell 9: Run untrained model (batched: all 3 tasks in parallel envs)\n",
         "print(\"Running UNTRAINED base model on all tasks (batched)...\")\n",
         "print(\"=\" * 60)\n",
         "\n",
         "t0 = time.time()\n",
-        "results = run_llm_episodes_batched(model, tokenizer, [(t, 42) for t in TASKS], verbose=True, eval=True)\n",
+        "results = run_llm_episodes_batched(model, tokenizer, [(t, 42) for t in TASKS], verbose=True, eval=True, log_tag=\"before\")\n",
         "before_results = {r[\"task\"]: r for r in results}\n",
         "\n",
         "print(\"\\n\" + \"=\" * 60)\n",
         "print(f\"BEFORE TRAINING (took {time.time()-t0:.1f}s):\")\n",
         "for t in TASKS:\n",
         "    print(f\"  {t}: grader={before_results[t]['grader_score']:.4f}\")"
-      ],
-      "execution_count": null,
-      "outputs": []
+      ]
     },
     {
       "cell_type": "markdown",
@@ -676,7 +690,9 @@
     },
     {
       "cell_type": "code",
+      "execution_count": null,
       "metadata": {},
+      "outputs": [],
       "source": [
         "# Cell 10: Attach LoRA adapter\n",
         "from peft import LoraConfig, get_peft_model, TaskType\n",
@@ -690,19 +706,19 @@
         "model.enable_input_require_grads()\n",
         "peft_model = get_peft_model(model, lora_config)\n",
         "peft_model.print_trainable_parameters()"
-      ],
-      "execution_count": null,
-      "outputs": []
+      ]
     },
     {
       "cell_type": "code",
+      "execution_count": null,
       "metadata": {},
+      "outputs": [],
       "source": [
         "# Cell 11: Training loop\n",
         "from trl import SFTTrainer, SFTConfig\n",
         "from datasets import Dataset\n",
         "\n",
-        "NUM_ROUNDS = 4\n",
+        "NUM_ROUNDS = 1\n",
         "EPISODES_PER_ROUND = 6\n",
         "QUALITY_FLOOR = 0.40  # skip SFT for the round if no episode beats this grader score\n",
         "\n",
@@ -723,7 +739,8 @@
         "    tasks_seeds = [(TASKS[ep % len(TASKS)], 42 + (round_idx - 1) * 100 + ep) for ep in range(EPISODES_PER_ROUND)]\n",
         "    t_roll = time.time()\n",
         "    results = run_llm_episodes_batched(peft_model, tokenizer, tasks_seeds, verbose=False,\n",
-        "                                       eval=False, system=SYSTEM_PROMPT_TRAIN)\n",
+        "                                       eval=False, system=SYSTEM_PROMPT_TRAIN,\n",
+        "                                       log_tag=f\"train_round{round_idx}\")\n",
         "    print(f\"  Rollouts: {len(results)} eps × {TASK_HORIZON} days in {time.time()-t_roll:.1f}s\")\n",
         "\n",
         "    all_pairs, episode_rewards, episode_graders = [], [], []\n",
@@ -800,9 +817,7 @@
         "elapsed = time.time() - t_start\n",
         "print(f\"\\nTraining complete in {elapsed/60:.1f} min\")\n",
         "print(pd.DataFrame(training_log).to_string(index=False))"
-      ],
-      "execution_count": null,
-      "outputs": []
+      ]
     },
     {
       "cell_type": "markdown",
@@ -815,7 +830,9 @@
     },
     {
       "cell_type": "code",
+      "execution_count": null,
       "metadata": {},
+      "outputs": [],
       "source": [
         "# Cell 12: Run trained model (batched)\n",
         "print(\"Running TRAINED model on all tasks (batched)...\")\n",
@@ -823,16 +840,14 @@
         "\n",
         "peft_model.eval()\n",
         "t0 = time.time()\n",
-        "results = run_llm_episodes_batched(peft_model, tokenizer, [(t, 42) for t in TASKS], verbose=True, eval=True)\n",
-        "after_results = {r[\"task\"]: r for r in results}\n",
+    "results = run_llm_episodes_batched(peft_model, tokenizer, [(t, 42) for t in TASKS], verbose=True, eval=True, log_tag=\"after\")\n",
+    "after_results = {r[\"task\"]: r for r in results}\n",
         "\n",
         "print(\"\\n\" + \"=\" * 60)\n",
         "print(f\"AFTER TRAINING (took {time.time()-t0:.1f}s):\")\n",
         "for t in TASKS:\n",
         "    print(f\"  {t}: grader={after_results[t]['grader_score']:.4f}\")"
-      ],
-      "execution_count": null,
-      "outputs": []
+      ]
     },
     {
       "cell_type": "markdown",
@@ -843,7 +858,9 @@
     },
     {
       "cell_type": "code",
+      "execution_count": null,
       "metadata": {},
+      "outputs": [],
       "source": [
         "# Cell 13: Training curves\n",
         "fig, axes = plt.subplots(1, 2, figsize=(14, 5))\n",
@@ -865,13 +882,13 @@
         "fig.tight_layout()\n",
         "fig.savefig(f'{PLOTS_DIR}/reward_curve.png', dpi=150, bbox_inches='tight')\n",
         "plt.show()"
-      ],
-      "execution_count": null,
-      "outputs": []
+      ]
     },
     {
       "cell_type": "code",
+      "execution_count": null,
       "metadata": {},
+      "outputs": [],
       "source": [
         "# Cell 14: Before vs After\n",
         "task_labels = [t.replace('monthly_', '').title() for t in TASKS]\n",
@@ -901,13 +918,13 @@
         "fig.tight_layout()\n",
         "fig.savefig(f'{PLOTS_DIR}/before_after.png', dpi=150, bbox_inches='tight')\n",
         "plt.show()"
-      ],
-      "execution_count": null,
-      "outputs": []
+      ]
     },
     {
       "cell_type": "code",
+      "execution_count": null,
       "metadata": {},
+      "outputs": [],
       "source": [
         "# Cell 15: Trajectory comparison\n",
         "fig, axes = plt.subplots(2, 3, figsize=(16, 8))\n",
@@ -931,9 +948,7 @@
         "fig.tight_layout()\n",
         "fig.savefig(f'{PLOTS_DIR}/training_trajectories.png', dpi=150, bbox_inches='tight')\n",
         "plt.show()"
-      ],
-      "execution_count": null,
-      "outputs": []
+      ]
     },
     {
       "cell_type": "markdown",
@@ -944,7 +959,9 @@
     },
     {
       "cell_type": "code",
+      "execution_count": null,
       "metadata": {},
+      "outputs": [],
       "source": [
         "# Cell 16: Final summary\n",
         "print(\"=\" * 67)\n",
@@ -981,13 +998,13 @@
         "\n",
         "print(f\"\\nSaved to {PLOTS_DIR}/\")\n",
         "print(\"All results are from real LoRA weight updates on real environment runs.\")"
-      ],
-      "execution_count": null,
-      "outputs": []
+      ]
     },
     {
       "cell_type": "code",
+      "execution_count": null,
       "metadata": {},
+      "outputs": [],
       "source": [
         "# Cell 17: Save adapter\n",
         "save_path = \"./viraltest_trained_adapter\"\n",
@@ -995,9 +1012,7 @@
         "tokenizer.save_pretrained(save_path)\n",
         "print(f\"LoRA adapter saved to {save_path}\")\n",
         "print(\"Load with: PeftModel.from_pretrained(base_model, save_path)\")"
-      ],
-      "execution_count": null,
-      "outputs": []
+      ]
     }
   ],
   "metadata": {
@@ -1023,4 +1038,4 @@
   },
   "nbformat": 4,
   "nbformat_minor": 4
-}
+}