Clean final training notebook: no demo, no Drive, all steps including plots + eval + download

plot_rewards.py

@@ -98,10 +98,16 @@ def load_log(log_path: str) -> pd.DataFrame:
 
     df.sort_values("episode", inplace=True)
     df.reset_index(drop=True, inplace=True)
-    print(f"[plot_rewards] Loaded {len(df)} episodes from {log_path}")
+    n = len(df)
+    print(f"[plot_rewards] Loaded {n} episodes from {log_path}")
     return df
 
 
+def _adaptive_window(df: pd.DataFrame, requested: int) -> int:
+    """Use min(requested, len/3) so plots are never flat lines with few data points."""
+    return max(3, min(requested, len(df) // 3))
+
+
 # ── Plots ─────────────────────────────────────────────────────────────────────
 
 def plot_reward_curve(df: pd.DataFrame, out_dir: Path, window: int = 20):
@@ -235,16 +241,33 @@ def plot_all(log_path: str = "logs/training.jsonl", out_dir: str = "plots/",
     out = Path(out_dir)
     out.mkdir(parents=True, exist_ok=True)
 
-    plot_reward_curve(df, out, window)
-    plot_success_rate(df, out, window)
-    plot_accuracy_gain(df, out, window)
+    # Adaptive window — avoid flat lines with small datasets
+    w = _adaptive_window(df, window)
+
+    plot_reward_curve(df, out, w)
+    plot_success_rate(df, out, w)
+    plot_accuracy_gain(df, out, w)
     plot_curriculum(df, out)
 
-    print(f"\n[plot_rewards] All plots saved to {out}/")
-    print(f"  Episodes: {len(df)} | "
-          f"Avg reward: {df['reward'].mean():.3f} | "
-          f"Success rate: {df['success'].mean():.1%} | "
-          f"Max level reached: {int(df['level'].max())}")
+    # ── Print summary stats ───────────────────────────────────────────────────
+    n = len(df)
+    avg_r = df["reward"].mean()
+    max_r = df["reward"].max()
+    min_r = df["reward"].min()
+    succ  = df["success"].mean()
+    max_lvl = int(df["level"].max())
+    lvl_names = {0: "tutorial", 1: "easy", 2: "medium", 3: "hard"}
+
+    print(f"\n{'='*50}")
+    print(f"  TRAINING SUMMARY ({n} episodes)")
+    print(f"{'='*50}")
+    print(f"  Avg reward   : {avg_r:+.3f}")
+    print(f"  Min / Max    : {min_r:+.3f} / {max_r:+.3f}")
+    print(f"  Success rate : {succ:.1%}")
+    print(f"  Max level    : {max_lvl} ({lvl_names.get(max_lvl, '?')})")
+    print(f"  Window used  : {w} episodes")
+    print(f"{'='*50}")
+    print(f"\n  Plots saved to: {out}/")
 
 
 if __name__ == "__main__":

train_colab.ipynb

@@ -3,8 +3,7 @@
  "nbformat_minor": 5,
  "metadata": {
   "colab": {
-   "provenance": [],
-   "gpuType": "T4"
+   "provenance": []
   },
   "kernelspec": {
    "display_name": "Python 3",
@@ -18,306 +17,153 @@
  "cells": [
   {
    "cell_type": "markdown",
-   "id": "a0",
+   "id": "hdr",
    "metadata": {},
    "source": [
-    "# \ud83e\udde0 Data-Centric AI Agent \u2014 Training Notebook\n",
-    "**OpenEnv Hackathon 2026** \u00b7 [Space](https://huggingface.co/spaces/Aswinis-Kumar/data-centric-env) \u00b7 [GitHub](https://github.com/CelestialWorthyOfHeavenAndEarth/data-centric-env)\n\n",
-    "Trains **Qwen2.5-3B-Instruct** (4-bit, LoRA) to orchestrate data-cleaning specialist agents via RL.\n\n",
-    "| Phase | What happens | Time |\n",
-    "|-------|-------------|------|\n",
-    "| SFT warmup | Teaches valid command grammar on ~9k examples | ~20 min |\n",
-    "| GRPO training | Learns strategy via live environment rollouts | ~3\u20135 hrs |\n",
-    "| Eval + plots | Reward curves, baseline comparison | ~5 min |\n\n",
-    "**Runtime required:** T4 GPU \u2192 `Runtime \u2192 Change runtime type \u2192 T4 GPU`\n\n",
-    "> **Quick demo mode:** Set `DEMO_MODE = True` in cell 1 to run a tiny 10-step version in ~10 min."
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "a1",
-   "metadata": {},
-   "source": [
-    "## \u2699\ufe0f Step 1 \u2014 Configuration"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "b0_config",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# \u2500\u2500\u2500 Set DEMO_MODE=True for a quick 10-step judge verification run (~5 min)\n",
-    "DEMO_MODE = False  # \u2190 Change to True for quick verification\n",
+    "# 🧠 Data-Centric AI Agent — Training Notebook\n",
     "\n",
-    "REPO     = 'https://github.com/CelestialWorthyOfHeavenAndEarth/data-centric-env.git'\n",
-    "HF_SPACE = 'https://aswinis-kumar-data-centric-env.hf.space'\n",
-    "\n",
-    "# Model: 1.5B trains 3x faster than 3B on T4, still very capable\n",
-    "# Change to '3B-Instruct' if you want more capacity but slower runs\n",
-    "MODEL_NAME = 'Qwen/Qwen2.5-1.5B-Instruct'\n",
-    "\n",
-    "if DEMO_MODE:\n",
-    "    print('\ud83d\ude80 DEMO MODE \u2014 quick 10-step run for judge verification')\n",
-    "    SFT_EPOCHS   = 1\n",
-    "    GRPO_STEPS   = 10\n",
-    "    GRPO_BATCH   = 2\n",
-    "    GRPO_GENS    = 2\n",
-    "else:\n",
-    "    print('\ud83d\udd2c FULL TRAINING MODE')\n",
-    "    SFT_EPOCHS   = 1\n",
-    "    GRPO_STEPS   = 150   # ~45-60 min on T4 with 1.5B. Run 2-3x for best results!\n",
-    "    GRPO_BATCH   = 2\n",
-    "    GRPO_GENS    = 2\n",
+    "**OpenEnv Hackathon 2026**\n",
     "\n",
-    "import os\n",
-    "os.environ['MODEL_NAME'] = MODEL_NAME\n",
-    "\n",
-    "print(f'  Model      : {MODEL_NAME}')\n",
-    "print(f'  SFT epochs : {SFT_EPOCHS}')\n",
-    "print(f'  GRPO steps : {GRPO_STEPS}')\n",
-    "print(f'  Batch/Gens : {GRPO_BATCH}/{GRPO_GENS}')\n",
-    "print()\n",
-    "print('TIP: Run 2-3 training runs of 150 steps each rather than one long run!')\n",
-    "print('     Multiple short runs = better learning curves for judges to see.')"
+    "Trains **Qwen2.5-1.5B-Instruct** (4-bit QLoRA) to orchestrate data-cleaning specialists\n",
+    "via GRPO reinforcement learning on the Data-Centric AI environment.\n",
+    "\n",
+    "| | |\n",
+    "|---|---|\n",
+    "| **Space** | https://huggingface.co/spaces/Aswinis-Kumar/data-centric-env |\n",
+    "| **Repo** | https://github.com/CelestialWorthyOfHeavenAndEarth/data-centric-env |\n",
+    "| **Algorithm** | SFT warmup → GRPO (TRL) |\n",
+    "| **Model** | Qwen/Qwen2.5-1.5B-Instruct, 4-bit QLoRA r=8 |\n",
+    "\n",
+    "**Pipeline:**\n",
+    "1. Install deps\n",
+    "2. Clone repo + start env server\n",
+    "3. Generate SFT warmup data\n",
+    "4. Phase 1 — SFT warmup (~15 min)\n",
+    "5. Phase 2 — GRPO training (~45-90 min depending on hardware)\n",
+    "6. Plot reward curves\n",
+    "7. Evaluate vs baselines\n",
+    "8. Save model\n"
    ]
   },
   {
    "cell_type": "markdown",
-   "id": "a1b",
+   "id": "s1_md",
    "metadata": {},
    "source": [
-    "## \ud83d\udce6 Step 2 \u2014 Install Dependencies (~5 min)"
+    "## 📦 Step 1 — Install Dependencies"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "b1",
+   "id": "s1",
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Step 1: Pin transformers to a version compatible with Colab's torchao\n",
-    "# (transformers >= 4.52 pulls torchao >= 0.9 which requires torch.int1 \u2014 not in Colab yet)\n",
+    "# Pin transformers first to avoid torchao conflict\n",
+    "# (transformers ≥ 4.52 requires torchao ≥ 0.9 which needs torch.int1 — not in Colab yet)\n",
     "!pip install \"transformers>=4.44.0,<4.52.0\" --quiet\n",
     "\n",
-    "# Step 2: Install Unsloth for the current Colab torch version\n",
+    "# Unsloth — detects the installed torch version automatically\n",
     "!pip install \"unsloth[colab-new] @ git+https://github.com/unslothai/unsloth.git\" --quiet\n",
     "\n",
-    "# Step 3: Install training + env + tracking dependencies\n",
+    "# Training stack\n",
     "!pip install trl>=0.15.0 datasets>=2.0.0 accelerate>=0.30.0 --quiet\n",
     "!pip install scikit-learn>=1.3.0 pandas>=2.0.0 numpy matplotlib --quiet\n",
     "!pip install \"openenv-core[core]>=0.2.1\" --quiet\n",
-    "# TensorBoard is already included in PyTorch/HF Trainer \u2014 no extra install needed\n",
-    "\n",
-    "# Step 4: Verify Unsloth loads correctly\n",
-    "try:\n",
-    "    from unsloth import FastLanguageModel\n",
-    "    print('\u2705 Unsloth loaded successfully')\n",
-    "except Exception as e:\n",
-    "    print(f'\u274c Unsloth import failed: {e}')\n",
-    "    print('   Try: Runtime \u2192 Factory reset runtime, then re-run')\n",
-    "    raise"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "a2",
-   "metadata": {},
-   "source": [
-    "## \ud83d\udcbe Step 3 \u2014 Mount Google Drive (optional, recommended)\n",
-    "Checkpoints save every 50 GRPO steps to Drive. If Colab disconnects, you can resume.\n\n",
-    "**Skip this cell if you don't want Drive integration.**"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "b2",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from google.colab import drive\n",
-    "drive.mount('/content/drive')\n",
-    "print('Drive mounted')"
+    "\n",
+    "# Verify Unsloth\n",
+    "from unsloth import FastLanguageModel\n",
+    "import torch\n",
+    "print(f'✅ Unsloth ready | torch {torch.__version__} | GPU: {torch.cuda.get_device_name(0) if torch.cuda.is_available() else \"CPU\"}')\n"
    ]
   },
   {
    "cell_type": "markdown",
-   "id": "a3",
+   "id": "s2_md",
    "metadata": {},
    "source": [
-    "## \ud83d\udd01 Step 4 \u2014 Clone / Update Repo + Setup"
+    "## 🔁 Step 2 — Clone Repo & Start Environment Server"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "b3",
+   "id": "s2",
    "metadata": {},
    "outputs": [],
    "source": [
-    "import os, sys, shutil\n",
+    "import os, sys, shutil, subprocess, time, requests\n",
     "\n",
+    "REPO     = 'https://github.com/CelestialWorthyOfHeavenAndEarth/data-centric-env.git'\n",
     "WORK_DIR = '/content/data-centric-env'\n",
     "\n",
-    "# Clone or pull latest\n",
+    "# Clone or update\n",
     "if os.path.exists(f'{WORK_DIR}/pyproject.toml'):\n",
-    "    print('Repo exists, pulling latest...')\n",
-    "    !git -C {WORK_DIR} pull origin main\n",
+    "    print('Repo exists — pulling latest...')\n",
+    "    os.system(f'git -C {WORK_DIR} pull origin main')\n",
     "else:\n",
-    "    if os.path.exists(WORK_DIR):\n",
-    "        shutil.rmtree(WORK_DIR)\n",
-    "    !git clone {REPO} {WORK_DIR}\n",
+    "    if os.path.exists(WORK_DIR): shutil.rmtree(WORK_DIR)\n",
+    "    os.system(f'git clone {REPO} {WORK_DIR}')\n",
     "\n",
     "os.chdir(WORK_DIR)\n",
     "sys.path.insert(0, WORK_DIR)\n",
+    "os.makedirs('logs', exist_ok=True)\n",
+    "os.makedirs('plots', exist_ok=True)\n",
     "print('CWD:', os.getcwd())\n",
-    "!git log --oneline -3\n",
-    "\n",
-    "# Install as editable package (fixes relative imports)\n",
-    "!pip install -e . --quiet 2>/dev/null || echo 'pip install -e . skipped'\n",
-    "\n",
-    "# Symlink output dirs to Drive if mounted\n",
-    "DRIVE_DIR = '/content/drive/MyDrive/data-centric-training'\n",
-    "if os.path.exists('/content/drive/MyDrive'):\n",
-    "    os.makedirs(DRIVE_DIR, exist_ok=True)\n",
-    "    for d in ['data-centric-checkpoints', 'data-centric-adapter', 'logs', 'plots']:\n",
-    "        local  = os.path.join(WORK_DIR, d)\n",
-    "        remote = os.path.join(DRIVE_DIR, d)\n",
-    "        os.makedirs(remote, exist_ok=True)\n",
-    "        if not os.path.exists(local):\n",
-    "            os.symlink(remote, local)\n",
-    "            print(f'  Drive link: {d}')\n",
-    "        else:\n",
-    "            print(f'  Exists: {d}')\n",
-    "    print('Checkpoints will save to Drive')\n",
-    "else:\n",
-    "    for d in ['logs', 'plots']:\n",
-    "        os.makedirs(os.path.join(WORK_DIR, d), exist_ok=True)\n",
-    "    print('Drive not mounted \u2014 checkpoints save locally only')\n",
-    "\n",
-    "print('Setup complete')"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "a35",
-   "metadata": {},
-   "source": [
-    "## \u2705 Step 5 \u2014 Verify All Features Work\n",
-    "Quick smoke test: confirms the environment, all specialist agents, and the rubric system are connected."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "b35",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import os, sys\n",
-    "os.chdir('/content/data-centric-env')\n",
-    "sys.path.insert(0, '/content/data-centric-env')\n",
-    "\n",
-    "from models import DataCentricAction\n",
-    "from server.data_centric_environment import DataCentricEnvironment\n",
-    "\n",
-    "env = DataCentricEnvironment()\n",
-    "obs = env.reset(task='task_1_easy', seed=42)\n",
-    "\n",
-    "checks = {}\n",
-    "\n",
-    "# 1. query_analyst\n",
-    "obs = env.step(DataCentricAction(message='query_analyst'))\n",
-    "checks['query_analyst'] = 'DIAGNOSIS' in obs.response and 'RECOMMENDED PLAN' in obs.response\n",
-    "\n",
-    "# 2. Smarter specialists\n",
-    "obs = env.step(DataCentricAction(message='query_cleaner'))\n",
-    "checks['smarter_specialists'] = any(k in obs.response for k in ['skew', 'Risk:', 'Reason:', '%'])\n",
-    "\n",
-    "# 3. Drift detection (apply then check)\n",
-    "obs = env.step(DataCentricAction(message='apply 1'))\n",
-    "checks['drift_detection'] = 'drift' in obs.response.lower() or obs.response != ''\n",
-    "\n",
-    "# 4+5. Dual classifier + feature importance\n",
-    "obs = env.step(DataCentricAction(message='validate'))\n",
-    "checks['dual_classifier']    = 'RF Accuracy' in obs.response or 'accuracy' in obs.response.lower()\n",
-    "checks['feature_importance'] = 'Feature importance' in obs.response or 'feature' in obs.response.lower()\n",
-    "\n",
-    "print('\\n=== Feature Smoke Test ===')\n",
-    "all_pass = True\n",
-    "for name, passed in checks.items():\n",
-    "    status = '\u2705 PASS' if passed else '\u274c FAIL'\n",
-    "    if not passed: all_pass = False\n",
-    "    print(f'  {status}: {name}')\n",
-    "\n",
-    "if all_pass:\n",
-    "    print('\\n\u2705 All features verified \u2014 ready to train!')\n",
-    "else:\n",
-    "    print('\\n\u26a0\ufe0f  Some checks failed \u2014 check git pull ran correctly')"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "a4",
-   "metadata": {},
-   "source": [
-    "## \ud83d\udda5\ufe0f Step 6 \u2014 Start Environment Server"
+    "os.system('git log --oneline -3')\n",
+    "os.system('pip install -e . --quiet 2>/dev/null || true')\n",
+    "print('✅ Repo ready')\n"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "b4",
+   "id": "s2b",
    "metadata": {},
    "outputs": [],
    "source": [
-    "import subprocess, time, requests, os, sys\n",
-    "os.chdir('/content/data-centric-env')\n",
-    "sys.path.insert(0, '/content/data-centric-env')\n",
+    "import subprocess, time, requests, os\n",
     "\n",
-    "# Kill any existing server on port 8000\n",
-    "!fuser -k 8000/tcp 2>/dev/null || true\n",
+    "# Kill any leftover server\n",
+    "os.system('fuser -k 8000/tcp 2>/dev/null || true')\n",
     "time.sleep(2)\n",
     "\n",
-    "server_proc = subprocess.Popen(\n",
+    "server = subprocess.Popen(\n",
     "    ['python', '-m', 'uvicorn', 'server.app:app',\n",
     "     '--host', '0.0.0.0', '--port', '8000', '--log-level', 'warning'],\n",
     "    stdout=subprocess.PIPE, stderr=subprocess.PIPE\n",
     ")\n",
     "\n",
-    "for i in range(30):\n",
+    "for i in range(40):\n",
     "    try:\n",
-    "        if requests.get('http://localhost:8000/health', timeout=2).status_code == 200:\n",
-    "            print(f'\u2705 Server ready after {i+1}s')\n",
+    "        r = requests.get('http://localhost:8000/health', timeout=2)\n",
+    "        if r.status_code == 200:\n",
+    "            print(f'✅ Server ready after {i+1}s — {r.json()}')\n",
     "            break\n",
-    "    except:\n",
-    "        time.sleep(1)\n",
+    "    except: time.sleep(1)\n",
     "else:\n",
-    "    server_proc.terminate()\n",
-    "    out, err = server_proc.communicate()\n",
+    "    out, err = server.communicate(timeout=5)\n",
     "    print('STDOUT:', out.decode()[-2000:])\n",
     "    print('STDERR:', err.decode()[-2000:])\n",
-    "    raise RuntimeError('Server failed to start in 30s \u2014 see output above')\n",
+    "    raise RuntimeError('Server failed to start')\n",
     "\n",
-    "print('ENV_URL: http://localhost:8000')\n",
-    "os.environ['ENV_URL'] = 'http://localhost:8000'"
+    "os.environ['ENV_URL'] = 'http://localhost:8000'\n",
+    "print('ENV_URL set to http://localhost:8000')\n"
    ]
   },
   {
    "cell_type": "markdown",
-   "id": "a5",
+   "id": "s3_md",
    "metadata": {},
    "source": [
-    "## \ud83d\udcca Step 7 \u2014 Generate SFT Warmup Data"
+    "## 📊 Step 3 — Generate SFT Warmup Data"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "b5",
+   "id": "s3",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -328,26 +174,26 @@
     "if os.path.exists(sft_path):\n",
     "    count = sum(1 for _ in open(sft_path))\n",
     "    age   = (time.time() - os.path.getmtime(sft_path)) / 60\n",
-    "    print(f'sft_data.jsonl: {count} examples, {age:.0f} min old \u2014 reusing')\n",
+    "    print(f'✅ sft_data.jsonl: {count} examples ({age:.0f} min old) — reusing')\n",
     "else:\n",
-    "    print('Generating SFT data...')\n",
-    "    !python sft_generator.py\n",
+    "    print('Generating SFT warmup data (~2 min)...')\n",
+    "    os.system('python sft_generator.py')\n",
     "    count = sum(1 for _ in open(sft_path))\n",
-    "    print(f'Generated {count} examples')"
+    "    print(f'✅ Generated {count} SFT examples')\n"
    ]
   },
   {
    "cell_type": "markdown",
-   "id": "a6",
+   "id": "s4_md",
    "metadata": {},
    "source": [
-    "## \ud83e\udd16 Step 8 \u2014 Load Model (Qwen2.5-3B-Instruct, 4-bit)"
+    "## 🤖 Step 4 — Load Qwen2.5-1.5B (4-bit QLoRA)"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "b6",
+   "id": "s4",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -356,47 +202,33 @@
     "sys.path.insert(0, '/content/data-centric-env')\n",
     "\n",
     "from unsloth import FastLanguageModel\n",
+    "from train_data_centric import load_model\n",
     "\n",
-    "MAX_SEQ_LENGTH = 512  # Short commands; saves significant VRAM\n",
-    "\n",
-    "model, tokenizer = FastLanguageModel.from_pretrained(\n",
-    "    model_name=MODEL_NAME,\n",
-    "    max_seq_length=MAX_SEQ_LENGTH,\n",
-    "    load_in_4bit=True,   # QLoRA \u2014 mandatory for T4\n",
-    "    dtype=None,\n",
-    ")\n",
-    "\n",
-    "model = FastLanguageModel.get_peft_model(\n",
-    "    model,\n",
-    "    r=8,                 # r=8 is sufficient for short command vocab\n",
-    "    target_modules=['q_proj', 'v_proj', 'k_proj', 'o_proj'],\n",
-    "    lora_alpha=16,       # = r, standard rule-of-thumb\n",
-    "    lora_dropout=0,\n",
-    "    bias='none',\n",
-    "    use_gradient_checkpointing='unsloth',\n",
-    "    random_state=42,\n",
-    ")\n",
+    "model, tokenizer = load_model()\n",
     "\n",
-    "vram = torch.cuda.memory_allocated()/1e9\n",
-    "total = torch.cuda.get_device_properties(0).total_memory/1e9\n",
-    "print(f'\u2705 Model loaded: {MODEL_NAME}')\n",
-    "print(f'   VRAM: {vram:.1f} GB / {total:.1f} GB ({100*vram/total:.0f}% used)')\n",
-    "print(f'   LoRA trainable params: ~1.8M (r=8)')"
+    "vram  = torch.cuda.memory_allocated() / 1e9\n",
+    "total = torch.cuda.get_device_properties(0).total_memory / 1e9\n",
+    "print(f'\\n📊 VRAM: {vram:.1f} GB used / {total:.1f} GB total ({100*vram/total:.0f}%)')\n"
    ]
   },
   {
    "cell_type": "markdown",
-   "id": "a7",
+   "id": "s5_md",
    "metadata": {},
    "source": [
-    "## \ud83c\udf93 Step 9 \u2014 Phase 1: SFT Warmup\n",
-    "Teaches the model to output valid commands before RL starts. Without this, the model outputs free-form text and gets zero reward."
+    "## 🎓 Step 5 — Phase 1: SFT Warmup\n",
+    "\n",
+    "One epoch of supervised fine-tuning on ~9,480 heuristic trajectories.\n",
+    "Teaches the model valid command syntax before RL starts.\n",
+    "Without this, the model outputs gibberish and gets zero reward.\n",
+    "\n",
+    "**TensorBoard:** Run Step 6 to open the live dashboard.\n"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "b7",
+   "id": "s5",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -406,232 +238,211 @@
     "\n",
     "from train_data_centric import run_sft_warmup\n",
     "\n",
-    "# In demo mode: reduce to 1 pass of 200 examples for speed\n",
-    "if DEMO_MODE:\n",
-    "    import json\n",
-    "    raw = [json.loads(l) for l in open('sft_data.jsonl')][:200]\n",
-    "    with open('sft_data_demo.jsonl', 'w') as f:\n",
-    "        for r in raw: f.write(json.dumps(r) + '\\n')\n",
-    "    import shutil\n",
-    "    shutil.copy('sft_data.jsonl', 'sft_data_full.jsonl')\n",
-    "    shutil.copy('sft_data_demo.jsonl', 'sft_data.jsonl')\n",
-    "    print('Demo mode: using 200 SFT examples')\n",
-    "\n",
     "model = run_sft_warmup(model, tokenizer)\n",
-    "print('\u2705 SFT warmup complete')\n",
-    "\n",
-    "# Restore full data if demo\n",
-    "if DEMO_MODE and os.path.exists('sft_data_full.jsonl'):\n",
-    "    import shutil\n",
-    "    shutil.copy('sft_data_full.jsonl', 'sft_data.jsonl')"
+    "print('✅ SFT warmup complete — model knows command syntax')\n"
    ]
   },
   {
    "cell_type": "markdown",
-   "id": "a8",
+   "id": "s6_md",
    "metadata": {},
    "source": [
-    "## \ud83c\udfcb\ufe0f Step 10 \u2014 Phase 2: GRPO Training\n",
-    "Agent learns *when* to use each command via reinforcement learning against the live environment.\n\n",
-    "- **Full mode:** ~300 steps, 3\u20135 hrs on T4. Checkpoints every 50 steps.\n",
-    "- **Demo mode:** 10 steps, ~10 min. Enough to verify the training loop works.\n\n",
-    "If Colab disconnects, re-run all cells \u2014 it auto-resumes from the latest checkpoint."
+    "## 📡 Step 6 — TensorBoard (Experiment Tracking)\n",
+    "\n",
+    "Launch this now — it will show SFT logs immediately and GRPO logs as training progresses.\n",
+    "Refresh the TensorBoard tab every few minutes during Step 7.\n"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "b8",
+   "id": "s6",
    "metadata": {},
    "outputs": [],
    "source": [
-    "import os, sys, glob\n",
-    "os.chdir('/content/data-centric-env')\n",
-    "sys.path.insert(0, '/content/data-centric-env')\n",
-    "\n",
-    "from train_data_centric import run_grpo_training\n",
-    "\n",
-    "# Auto-detect checkpoint to resume from\n",
-    "ckpt_dir = './data-centric-checkpoints'\n",
-    "resume_from = None\n",
-    "if os.path.exists(ckpt_dir):\n",
-    "    checkpoints = sorted(glob.glob(f'{ckpt_dir}/checkpoint-*'))\n",
-    "    if checkpoints:\n",
-    "        resume_from = checkpoints[-1]\n",
-    "        print(f'Resuming from: {resume_from}')\n",
-    "\n",
-    "# max_steps > 0 limits the run (used for demo mode)\n",
-    "# max_steps = -1 means run full num_train_epochs\n",
-    "demo_steps = GRPO_STEPS if DEMO_MODE else -1\n",
-    "\n",
-    "model = run_grpo_training(\n",
-    "    model, tokenizer,\n",
-    "    resume_from_checkpoint=resume_from,\n",
-    "    max_steps=demo_steps,\n",
-    ")\n",
-    "print('\u2705 GRPO training complete')"
+    "%load_ext tensorboard\n",
+    "%tensorboard --logdir /content/data-centric-env/logs\n"
    ]
   },
   {
    "cell_type": "markdown",
-   "id": "a9",
+   "id": "s7_md",
    "metadata": {},
    "source": [
-    "## \ud83d\udcbd Step 11 \u2014 Save Trained Model"
+    "## 🏋️ Step 7 — Phase 2: GRPO Training\n",
+    "\n",
+    "Agent learns *when* to use each command via reinforcement learning.\n",
+    "Each GRPO step: model generates a command → environment executes it → reward flows back.\n",
+    "\n",
+    "**Reward discrimination:** The reward is now strictly discriminating:\n",
+    "- Hitting target efficiently (few steps): **+0.7 → +0.9**\n",
+    "- Hitting target slowly (many steps): **+0.3 → +0.5**\n",
+    "- Failing to hit target: **−0.1 → −0.5**\n",
+    "- Blind apply / blind submit: **large penalties**\n",
+    "\n",
+    "Checkpoints saved every 25 steps to `./data-centric-checkpoints/`.\n",
+    "If the runtime disconnects, re-run all cells — it auto-resumes from the latest checkpoint.\n"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "b9",
+   "id": "s7",
    "metadata": {},
    "outputs": [],
    "source": [
-    "import os, sys\n",
+    "import os, sys, glob\n",
     "os.chdir('/content/data-centric-env')\n",
     "sys.path.insert(0, '/content/data-centric-env')\n",
     "\n",
-    "from train_data_centric import save_model\n",
-    "save_model(model, tokenizer)\n",
-    "print('\u2705 Model saved to ./data-centric-adapter/')"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "a_tb",
-   "metadata": {},
-   "source": [
-    "## \ud83d\udcca Step 10.5 \u2014 Launch TensorBoard (Experiment Tracking)\n",
-    "View live loss and reward curves while training is running.\n",
-    "Both SFT and GRPO write logs automatically \u2014 no API key required."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "b_tb",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# Load TensorBoard extension (built into Colab/PyTorch)\n",
-    "%load_ext tensorboard\n",
-    "\n",
-    "# View SFT logs (run after Step 9)\n",
-    "# %tensorboard --logdir /content/data-centric-env/logs/sft\n",
+    "from train_data_centric import run_grpo_training\n",
     "\n",
-    "# View GRPO logs (run after Step 10 starts)\n",
-    "%tensorboard --logdir /content/data-centric-env/logs/grpo\n",
+    "# Auto-resume from latest checkpoint\n",
+    "ckpt_dir = './data-centric-checkpoints'\n",
+    "resume_from = None\n",
+    "if os.path.exists(ckpt_dir):\n",
+    "    checkpoints = sorted(glob.glob(f'{ckpt_dir}/checkpoint-*'))\n",
+    "    if checkpoints:\n",
+    "        resume_from = checkpoints[-1]\n",
+    "        print(f'Resuming from: {resume_from}')\n",
     "\n",
-    "print(\"TensorBoard launched above \u2014 refresh as training progresses\")"
+    "model = run_grpo_training(model, tokenizer, resume_from_checkpoint=resume_from)\n",
+    "print('✅ GRPO training complete')\n"
    ]
   },
   {
    "cell_type": "markdown",
-   "id": "a10",
+   "id": "s8_md",
    "metadata": {},
    "source": [
-    "## \ud83d\udcc8 Step 12 \u2014 Plot Reward Curves"
+    "## 📈 Step 8 — Plot Reward Curves\n",
+    "\n",
+    "Generates 4 plots from the training log:\n",
+    "- **reward_curve.png** — Episode reward over time with rolling mean\n",
+    "- **success_rate.png** — Success rate per curriculum level\n",
+    "- **accuracy_gain.png** — Accuracy gain vs baseline per episode\n",
+    "- **curriculum.png** — Curriculum level progression\n"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "b10",
+   "id": "s8",
    "metadata": {},
    "outputs": [],
    "source": [
     "import os\n",
     "os.chdir('/content/data-centric-env')\n",
     "\n",
-    "if os.path.exists('logs/training.jsonl'):\n",
-    "    lines = sum(1 for _ in open('logs/training.jsonl'))\n",
-    "    print(f'Training log: {lines} episodes recorded')\n",
-    "    !python plot_rewards.py --log logs/training.jsonl --out plots/\n",
+    "log_path = 'logs/training.jsonl'\n",
+    "\n",
+    "if not os.path.exists(log_path):\n",
+    "    print('No training log yet — run Step 7 first')\n",
+    "else:\n",
+    "    n_eps = sum(1 for _ in open(log_path))\n",
+    "    print(f'Training log: {n_eps} episodes')\n",
+    "\n",
+    "    os.system(f'python plot_rewards.py --log {log_path} --out plots/')\n",
+    "\n",
     "    from IPython.display import Image, display\n",
-    "    for f in ['reward_curve.png', 'success_rate.png', 'accuracy_gain.png', 'curriculum.png']:\n",
-    "        path = f'plots/{f}'\n",
+    "    for fname in ['reward_curve.png', 'success_rate.png', 'accuracy_gain.png', 'curriculum.png']:\n",
+    "        path = f'plots/{fname}'\n",
     "        if os.path.exists(path):\n",
-    "            print(f'\\n--- {f} ---')\n",
-    "            display(Image(path))\n",
-    "else:\n",
-    "    print('No training log yet \u2014 run Step 10 first')"
+    "            print(f'\\n--- {fname} ---')\n",
+    "            display(Image(path, width=900))\n"
    ]
   },
   {
    "cell_type": "markdown",
-   "id": "a11",
+   "id": "s9_md",
    "metadata": {},
    "source": [
-    "## \ud83c\udfc6 Step 13 \u2014 Evaluate: Trained Agent vs Random vs Heuristic"
+    "## 🏆 Step 9 — Evaluate: Trained Agent vs Baselines\n",
+    "\n",
+    "Runs the trained agent against:\n",
+    "- **Random baseline** — random command each step\n",
+    "- **Heuristic baseline** — fixed rule-based sequence\n",
+    "- **Trained agent** — your GRPO-trained model\n"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "b11",
+   "id": "s9",
    "metadata": {},
    "outputs": [],
    "source": [
     "import os, json\n",
     "os.chdir('/content/data-centric-env')\n",
     "\n",
-    "!python eval_data_centric.py\n",
+    "os.system('python eval_data_centric.py')\n",
     "\n",
     "if os.path.exists('eval_results.json'):\n",
     "    with open('eval_results.json') as f:\n",
     "        results = json.load(f)\n",
-    "    print('\\n=== Eval Results ===')\n",
+    "    print('\\n=== Evaluation Results ===')\n",
     "    for k, v in results.items():\n",
-    "        print(f'  {k}: {v}')"
+    "        print(f'  {k}: {v}')\n"
    ]
   },
   {
    "cell_type": "markdown",
-   "id": "a12",
+   "id": "s10_md",
    "metadata": {},
    "source": [
-    "## \ud83d\udce4 Step 14 \u2014 Push Results to GitHub\n",
-    "Run this **only after training + eval are done**."
+    "## 💾 Step 10 — Save Trained LoRA Adapter"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "b12",
+   "id": "s10",
    "metadata": {},
    "outputs": [],
    "source": [
-    "import os\n",
+    "import os, sys\n",
     "os.chdir('/content/data-centric-env')\n",
+    "sys.path.insert(0, '/content/data-centric-env')\n",
     "\n",
-    "from getpass import getpass\n",
-    "token = getpass('GitHub token (repo write access): ')\n",
-    "repo_url = f'https://{token}@github.com/CelestialWorthyOfHeavenAndEarth/data-centric-env.git'\n",
-    "\n",
-    "!git config user.email 'training@colab.run'\n",
-    "!git config user.name 'Colab Training'\n",
-    "!git remote set-url origin {repo_url}\n",
-    "!git add plots/ logs/ eval_results.json 2>/dev/null || true\n",
-    "!git commit -m 'Add training results: reward curves + eval [Colab]'\n",
-    "!git push origin main\n",
-    "print('\u2705 Results pushed to GitHub')"
+    "from train_data_centric import save_model\n",
+    "save_model(model, tokenizer)\n",
+    "print('✅ Adapter saved to ./data-centric-adapter/')\n",
+    "\n",
+    "# List outputs\n",
+    "import glob\n",
+    "print('\\n📁 Training outputs:')\n",
+    "for f in sorted(glob.glob('plots/*.png') + glob.glob('logs/*.jsonl') + ['eval_results.json']):\n",
+    "    if os.path.exists(f):\n",
+    "        size = os.path.getsize(f) / 1024\n",
+    "        print(f'  {f} ({size:.1f} KB)')\n"
    ]
   },
   {
    "cell_type": "markdown",
-   "id": "a13",
+   "id": "s11_md",
    "metadata": {},
    "source": [
-    "---\n",
-    "## \u2705 Done!\n\n",
-    "| Output | Location |\n",
-    "|--------|----------|\n",
-    "| Trained LoRA adapter | `./data-centric-adapter/` |\n",
-    "| Training log | `logs/training.jsonl` |\n",
-    "| Reward curves | `plots/*.png` |\n",
-    "| Eval results | `eval_results.json` |\n\n",
-    "**Links:**  \n",
-    "\ud83e\udd17 [HF Space](https://huggingface.co/spaces/Aswinis-Kumar/data-centric-env) \u00b7 [Blog Post](https://huggingface.co/blog/Aswinis-Kumar/data-centric-ai-agent) \u00b7 [GitHub](https://github.com/CelestialWorthyOfHeavenAndEarth/data-centric-env)"
+    "## 📤 Step 11 — Download Results\n",
+    "\n",
+    "Download plots and training log to your local machine.\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "s11",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from google.colab import files\n",
+    "import glob, os\n",
+    "\n",
+    "to_download = glob.glob('plots/*.png') + ['logs/training.jsonl', 'eval_results.json']\n",
+    "\n",
+    "for f in to_download:\n",
+    "    if os.path.exists(f):\n",
+    "        print(f'Downloading: {f}')\n",
+    "        files.download(f)\n"
    ]
   }
  ]