feat: add GRPO training pipeline for GridMind-RL environment via Unsloth and TRL

scripts/gridmind_grpo_colab.ipynb

@@ -343,14 +343,16 @@
     "torch.cuda.empty_cache()\n",
     "\n",
     "MODEL_NAME = \"Qwen/Qwen2.5-1.5B-Instruct\"\n",
-    "print(f\"Loading {MODEL_NAME} with 4-bit quantization for T4 16GB...\")\n",
+    "gpu_name = torch.cuda.get_device_name(0) if torch.cuda.is_available() else \"CPU\"\n",
+    "gpu_total_gb = torch.cuda.get_device_properties(0).total_memory / 1e9 if torch.cuda.is_available() else 0\n",
+    "print(f\"Loading {MODEL_NAME} with 4-bit quantization on {gpu_name} ({gpu_total_gb:.2f} GB VRAM)...\")\n",
     "\n",
     "tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)\n",
     "if tokenizer.pad_token is None:\n",
     "    tokenizer.pad_token = tokenizer.eos_token\n",
     "tokenizer.padding_side = \"left\"  # required for GRPO\n",
     "\n",
-    "# 4-bit quantization - fits safely on T4 16GB\n",
+    "# 4-bit quantization for memory-efficient QLoRA training\n",
     "bnb_config = BitsAndBytesConfig(\n",
     "    load_in_4bit=True,\n",
     "    bnb_4bit_compute_dtype=torch.float16,\n",
@@ -366,8 +368,8 @@
     ")\n",
     "\n",
     "print(f\"Model loaded on: {next(model.parameters()).device}\")\n",
-    "print(f\"Memory allocated: {torch.cuda.memory_allocated()/1e9:.2f} GB / 16 GB\")\n",
-    "print(f\"Memory reserved:  {torch.cuda.memory_reserved()/1e9:.2f} GB / 16 GB\")"
+    "print(f\"Memory allocated: {torch.cuda.memory_allocated()/1e9:.2f} GB / {gpu_total_gb:.2f} GB\")\n",
+    "print(f\"Memory reserved:  {torch.cuda.memory_reserved()/1e9:.2f} GB / {gpu_total_gb:.2f} GB\")"
    ]
   },
   {
@@ -417,8 +419,15 @@
     "        reset_payload = {\"task_id\": task_id, \"seed\": batch_seed}\n",
     "        reset_r = _requests.post(f\"{ENV_URL}/reset\", json=reset_payload, timeout=10)\n",
     "        reset_ok = reset_r.status_code == 200\n",
+    "        reset_data = reset_r.json() if reset_ok else {}\n",
+    "        reset_obs = reset_data.get(\"observations\", [reset_data.get(\"observation\", {})])\n",
+    "        if isinstance(reset_obs, list):\n",
+    "            base_obs = reset_obs[0] if reset_obs else {}\n",
+    "        else:\n",
+    "            base_obs = reset_obs or {}\n",
     "    except Exception:\n",
     "        reset_ok = False\n",
+    "        base_obs = {}\n",
     "\n",
     "    if not reset_ok:\n",
     "        return [-0.1] * len(completions)\n",
@@ -450,25 +459,33 @@
     "            cleaned_action = {}\n",
     "\n",
     "            try:\n",
-    "                cleaned_action[\"hvac_power_level\"] = max(0.0, min(1.0, float(action.get(\"hvac_power_level\", 0.5))))\n",
+    "                if \"hvac_power_level\" not in action:\n",
+    "                    raise KeyError(\"hvac_power_level\")\n",
+    "                cleaned_action[\"hvac_power_level\"] = max(0.0, min(1.0, float(action[\"hvac_power_level\"])))\n",
     "                valid_fields += 1\n",
     "            except Exception:\n",
     "                cleaned_action[\"hvac_power_level\"] = 0.5\n",
     "\n",
     "            try:\n",
-    "                cleaned_action[\"thermal_charge_rate\"] = max(-1.0, min(1.0, float(action.get(\"thermal_charge_rate\", 0.0))))\n",
+    "                if \"thermal_charge_rate\" not in action:\n",
+    "                    raise KeyError(\"thermal_charge_rate\")\n",
+    "                cleaned_action[\"thermal_charge_rate\"] = max(-1.0, min(1.0, float(action[\"thermal_charge_rate\"])))\n",
     "                valid_fields += 1\n",
     "            except Exception:\n",
     "                cleaned_action[\"thermal_charge_rate\"] = 0.0\n",
     "\n",
     "            try:\n",
-    "                cleaned_action[\"batch_job_slot\"] = max(0, min(4, int(action.get(\"batch_job_slot\", 0))))\n",
+    "                if \"batch_job_slot\" not in action:\n",
+    "                    raise KeyError(\"batch_job_slot\")\n",
+    "                cleaned_action[\"batch_job_slot\"] = max(0, min(4, int(action[\"batch_job_slot\"])))\n",
     "                valid_fields += 1\n",
     "            except Exception:\n",
     "                cleaned_action[\"batch_job_slot\"] = 0\n",
     "\n",
     "            try:\n",
-    "                cleaned_action[\"load_shed_fraction\"] = max(0.0, min(0.5, float(action.get(\"load_shed_fraction\", 0.0))))\n",
+    "                if \"load_shed_fraction\" not in action:\n",
+    "                    raise KeyError(\"load_shed_fraction\")\n",
+    "                cleaned_action[\"load_shed_fraction\"] = max(0.0, min(0.5, float(action[\"load_shed_fraction\"])))\n",
     "                valid_fields += 1\n",
     "            except Exception:\n",
     "                cleaned_action[\"load_shed_fraction\"] = 0.0\n",
@@ -496,18 +513,32 @@
     "            grid_r = float(comps.get(\"grid_response\", 0.0))\n",
     "            task_r = float(comps.get(\"task_satisfaction\", 0.0))\n",
     "\n",
+    "            price = float(base_obs.get(\"current_price\", base_obs.get(\"price\", 0.10)))\n",
+    "            stress = float(base_obs.get(\"grid_stress_signal\", base_obs.get(\"grid_stress\", 0.0)))\n",
+    "            temp = float(base_obs.get(\"indoor_temperature\", 21.0))\n",
+    "            charge = cleaned_action[\"thermal_charge_rate\"]\n",
+    "            hvac = cleaned_action[\"hvac_power_level\"]\n",
+    "            shed = cleaned_action[\"load_shed_fraction\"]\n",
+    "\n",
+    "            price_signal = 0.0\n",
+    "            if price < 0.08:\n",
+    "                price_signal += 0.08 * charge\n",
+    "            elif price > 0.15:\n",
+    "                price_signal += 0.08 * (-charge)\n",
+    "            price_signal -= 0.03 * abs(charge)\n",
+    "\n",
+    "            stress_signal = 0.12 * shed if stress > 0.65 else -0.04 * shed\n",
+    "            comfort_signal = -0.04 * abs(temp - 21.0) * abs(hvac - 0.5)\n",
+    "            action_signal = price_signal + stress_signal + comfort_signal\n",
+    "\n",
     "            if comps:\n",
-    "                composite = (\n",
-    "                    cost_r * 0.40 +\n",
-    "                    comfort_r * 0.25 +\n",
-    "                    grid_r * 0.15 +\n",
-    "                    task_r * 0.20 +\n",
-    "                    completeness_bonus\n",
-    "                )\n",
+    "                component_signal = 0.04 * cost_r + 0.03 * comfort_r + 0.03 * grid_r + 0.03 * task_r\n",
     "            else:\n",
-    "                composite = env_reward * 0.5 + completeness_bonus\n",
+    "                component_signal = 0.0\n",
     "\n",
-    "            composite = max(-0.6, min(0.6, composite))\n",
+    "            # Center raw env reward to avoid saturating all valid JSON at the clip boundary.\n",
+    "            composite = (env_reward - 0.5) * 0.35 + component_signal + action_signal + completeness_bonus\n",
+    "            composite = max(-0.45, min(0.45, composite))\n",
     "\n",
     "            rewards.append(composite)\n",
     "            batch_raw.append(composite)\n",
@@ -641,25 +672,37 @@
     "    task_type=\"CAUSAL_LM\",\n",
     ")\n",
     "\n",
-    "# GRPOConfig - trl==0.23.0 compatible. Pass this as args=, not config=.\n",
-    "# generation_kwargs is not a GRPOTrainer init parameter in trl 0.23.0.\n",
-    "grpo_config = GRPOConfig(\n",
-    "    output_dir=\"./gridmind-grpo-output\",\n",
-    "    num_train_epochs=1,\n",
-    "    max_steps=60,\n",
-    "    per_device_train_batch_size=1,\n",
-    "    gradient_accumulation_steps=4,\n",
-    "    max_prompt_length=400,\n",
-    "    max_completion_length=80,\n",
-    "    num_generations=4,\n",
-    "    learning_rate=5e-5,\n",
-    "    fp16=True,\n",
-    "    logging_steps=1,\n",
-    "    save_steps=60,\n",
-    "    report_to=\"none\",\n",
-    "    dataloader_num_workers=0,\n",
-    "    remove_unused_columns=False,\n",
-    ")\n",
+    "# GRPOConfig compatibility shim. HF/Colab images can have TRL builds whose\n",
+    "# GRPOConfig fields differ, so only pass arguments accepted by this runtime.\n",
+    "grpo_config_requested = {\n",
+    "    \"output_dir\": \"./gridmind-grpo-output\",\n",
+    "    \"num_train_epochs\": 1,\n",
+    "    \"max_steps\": 60,\n",
+    "    \"per_device_train_batch_size\": 1,\n",
+    "    \"gradient_accumulation_steps\": 4,\n",
+    "    \"max_prompt_length\": 400,\n",
+    "    \"max_completion_length\": 80,\n",
+    "    \"max_new_tokens\": 80,\n",
+    "    \"num_generations\": 4,\n",
+    "    \"learning_rate\": 5e-5,\n",
+    "    \"fp16\": True,\n",
+    "    \"logging_steps\": 1,\n",
+    "    \"save_steps\": 60,\n",
+    "    \"report_to\": \"none\",\n",
+    "    \"dataloader_num_workers\": 0,\n",
+    "    \"remove_unused_columns\": False,\n",
+    "}\n",
+    "\n",
+    "grpo_config_sig = inspect.signature(GRPOConfig.__init__)\n",
+    "grpo_config_params = set(grpo_config_sig.parameters.keys()) - {\"self\"}\n",
+    "grpo_config_kwargs = {k: v for k, v in grpo_config_requested.items() if k in grpo_config_params}\n",
+    "if \"max_completion_length\" in grpo_config_kwargs and \"max_new_tokens\" in grpo_config_kwargs:\n",
+    "    grpo_config_kwargs.pop(\"max_new_tokens\")\n",
+    "skipped_config_keys = [k for k in grpo_config_requested if k not in grpo_config_params]\n",
+    "print(f\"GRPOConfig accepted keys: {sorted(grpo_config_kwargs.keys())}\")\n",
+    "print(f\"GRPOConfig skipped unsupported keys: {skipped_config_keys}\")\n",
+    "\n",
+    "grpo_config = GRPOConfig(**grpo_config_kwargs)\n",
     "\n",
     "# Confirm the installed TRL API before constructing the trainer.\n",
     "import trl\n",
@@ -671,8 +714,10 @@
     "print(f\"Uses 'args=':   {'args' in params}\")\n",
     "print(f\"Uses 'config=': {'config' in params}\")\n",
     "\n",
-    "print(f\"\\nGPU memory: {torch.cuda.memory_allocated()/1e9:.2f} GB used / 16 GB total\")\n",
-    "print(f\"Free: {(16 - torch.cuda.memory_allocated()/1e9):.2f} GB\")\n",
+    "gpu_total_gb = torch.cuda.get_device_properties(0).total_memory / 1e9 if torch.cuda.is_available() else 0\n",
+    "gpu_used_gb = torch.cuda.memory_allocated() / 1e9 if torch.cuda.is_available() else 0\n",
+    "print(f\"\\nGPU memory: {gpu_used_gb:.2f} GB used / {gpu_total_gb:.2f} GB total\")\n",
+    "print(f\"Free: {max(0, gpu_total_gb - gpu_used_gb):.2f} GB\")\n",
     "\n",
     "# Custom callback to capture loss at every step for graphing.\n",
     "from transformers import TrainerCallback\n",
@@ -680,12 +725,16 @@
     "step_losses = []\n",
     "step_numbers = []\n",
     "step_reward_means = []\n",
+    "training_log_history = []\n",
     "\n",
     "class LossCaptureCallback(TrainerCallback):\n",
     "    def on_log(self, args, state, control, logs=None, **kwargs):\n",
     "        if not logs:\n",
     "            return\n",
     "        step = state.global_step\n",
+    "        row = {\"step\": step}\n",
+    "        row.update({k: float(v) if isinstance(v, (int, float)) else v for k, v in logs.items()})\n",
+    "        training_log_history.append(row)\n",
     "        loss = logs.get(\"loss\", logs.get(\"train_loss\", None))\n",
     "        if loss is not None:\n",
     "            step_losses.append(float(loss))\n",
@@ -714,7 +763,7 @@
     "\n",
     "print(\"\\nStarting GRPO training with QLoRA...\")\n",
     "print(\"Watch for non-zero loss values. If all zeros, reward variance is still too low.\\n\")\n",
-    "print(f\"Steps: {grpo_config.max_steps} | Batch: {grpo_config.per_device_train_batch_size} | Generations: {grpo_config.num_generations}\")\n",
+    "print(f\"Steps: {getattr(grpo_config, 'max_steps', 60)} | Batch: {getattr(grpo_config, 'per_device_train_batch_size', 1)} | Generations: {getattr(grpo_config, 'num_generations', 4)}\")\n",
     "print(\"Estimated time: ~25-35 min on T4\\n\")\n",
     "\n",
     "train_result = trainer.train()\n",
@@ -732,6 +781,38 @@
     "else:\n",
     "    print(f\"\\nTraining produced gradient signal on {len(non_zero_losses)} steps.\")\n",
     "\n",
+    "# Preserve the exact tabular statistics that TRL prints during training.\n",
+    "try:\n",
+    "    import pandas as pd\n",
+    "    trainer_log_rows = [r for r in trainer.state.log_history if \"loss\" in r or \"reward\" in r or \"rewards / reward_func / mean\" in r]\n",
+    "    if trainer_log_rows:\n",
+    "        training_metrics_df = pd.DataFrame(trainer_log_rows)\n",
+    "        if \"step\" not in training_metrics_df.columns:\n",
+    "            training_metrics_df.insert(0, \"step\", range(1, len(training_metrics_df) + 1))\n",
+    "    elif training_log_history:\n",
+    "        training_metrics_df = pd.DataFrame(training_log_history)\n",
+    "    else:\n",
+    "        training_metrics_df = pd.DataFrame({\"step\": step_numbers, \"loss\": step_losses, \"reward\": step_reward_means[:len(step_numbers)]})\n",
+    "\n",
+    "    os.makedirs(\"results\", exist_ok=True)\n",
+    "    training_metrics_path = \"results/gridmind_training_metrics.csv\"\n",
+    "    training_metrics_df.to_csv(training_metrics_path, index=False)\n",
+    "    print(f\"\\nSaved TRL training metrics table to {training_metrics_path}\")\n",
+    "\n",
+    "    preferred_cols = [\n",
+    "        \"step\", \"loss\", \"reward\", \"reward_std\",\n",
+    "        \"completions / mean_length\", \"completions / min_length\", \"completions / max_length\",\n",
+    "        \"completions / clipped_ratio\", \"kl\", \"rewards / reward_func / mean\", \"rewards / reward_func / std\",\n",
+    "    ]\n",
+    "    display_cols = [c for c in preferred_cols if c in training_metrics_df.columns]\n",
+    "    if display_cols:\n",
+    "        print(\"\\nTraining metrics preview:\")\n",
+    "        display(training_metrics_df[display_cols].tail(10))\n",
+    "except Exception as e:\n",
+    "    training_metrics_df = None\n",
+    "    training_metrics_path = None\n",
+    "    print(f\"Could not build training metrics table: {e}\")\n",
+    "\n",
     "print(f\"\\nMemory after training: {torch.cuda.memory_allocated()/1e9:.2f} GB\")\n",
     "\n",
     "# Save LoRA adapter (much smaller than full model)\n",
@@ -863,9 +944,35 @@
     "import matplotlib.pyplot as plt\n",
     "import matplotlib.gridspec as gridspec\n",
     "import numpy as np\n",
+    "import pandas as pd\n",
     "import os\n",
     "\n",
     "os.makedirs(\"results\", exist_ok=True)\n",
+    "os.makedirs(\"plots\", exist_ok=True)\n",
+    "\n",
+    "# Build a TRL-style metrics table from trainer logs. This matches the tabular\n",
+    "# output with columns like reward, reward_std, completion lengths, tools, and KL.\n",
+    "if 'training_metrics_df' not in globals() or training_metrics_df is None:\n",
+    "    trainer_log_rows = [r for r in trainer.state.log_history if \"loss\" in r or \"reward\" in r or \"rewards / reward_func / mean\" in r]\n",
+    "    training_metrics_df = pd.DataFrame(trainer_log_rows if trainer_log_rows else training_log_history)\n",
+    "    if not training_metrics_df.empty and \"step\" not in training_metrics_df.columns:\n",
+    "        training_metrics_df.insert(0, \"step\", range(1, len(training_metrics_df) + 1))\n",
+    "\n",
+    "training_metrics_path = \"results/gridmind_training_metrics.csv\"\n",
+    "if not training_metrics_df.empty:\n",
+    "    training_metrics_df.to_csv(training_metrics_path, index=False)\n",
+    "    print(f\"Saved TRL metrics table to {training_metrics_path}\")\n",
+    "    preferred_cols = [\n",
+    "        \"step\", \"loss\", \"reward\", \"reward_std\",\n",
+    "        \"completions / mean_length\", \"completions / min_length\", \"completions / max_length\",\n",
+    "        \"completions / clipped_ratio\", \"completions / mean_terminated_length\",\n",
+    "        \"completions / min_terminated_length\", \"completions / max_terminated_length\",\n",
+    "        \"tools / call_frequency\", \"tools / failure_frequency\", \"kl\",\n",
+    "        \"rewards / reward_func / mean\", \"rewards / reward_func / std\",\n",
+    "    ]\n",
+    "    display_cols = [c for c in preferred_cols if c in training_metrics_df.columns]\n",
+    "    if display_cols:\n",
+    "        display(training_metrics_df[display_cols].tail(20))\n",
     "\n",
     "tasks = [1, 2, 3, 4]\n",
     "task_labels = [\n",
@@ -963,23 +1070,39 @@
     "# Panel B: reward signal over time.\n",
     "ax_rew = fig.add_subplot(gs[1, 0])\n",
     "style_ax(ax_rew, 'GRPO Training: Reward Signal per Step')\n",
-    "if training_rewards and len(training_rewards) >= 4:\n",
-    "    raw = training_rewards\n",
-    "    steps_r = list(range(1, len(raw) + 1))\n",
-    "    ax_rew.plot(steps_r, raw, alpha=0.20, color=C['reward'], linewidth=1)\n",
+    "if not training_metrics_df.empty and (\"reward\" in training_metrics_df.columns or \"rewards / reward_func / mean\" in training_metrics_df.columns):\n",
+    "    reward_col = \"reward\" if \"reward\" in training_metrics_df.columns else \"rewards / reward_func / mean\"\n",
+    "    std_col = \"reward_std\" if \"reward_std\" in training_metrics_df.columns else \"rewards / reward_func / std\"\n",
+    "    reward_df = training_metrics_df[[\"step\", reward_col] + ([std_col] if std_col in training_metrics_df.columns else [])].dropna(subset=[reward_col])\n",
+    "    steps_r = reward_df[\"step\"].astype(float).tolist()\n",
+    "    raw = reward_df[reward_col].astype(float).tolist()\n",
+    "    ax_rew.plot(steps_r, raw, alpha=0.28, color=C['reward'], linewidth=1.2, marker='o', markersize=2, label='Logged reward')\n",
     "    ax_rew.plot(steps_r, smooth(raw, window=6), color=C['reward'], linewidth=2.5, label='Smoothed reward')\n",
+    "    if std_col in reward_df.columns:\n",
+    "        std = reward_df[std_col].fillna(0).astype(float).to_numpy()\n",
+    "        raw_np = np.array(raw)\n",
+    "        steps_np = np.array(steps_r)\n",
+    "        ax_rew.fill_between(steps_np, raw_np - std, raw_np + std, color=C['reward'], alpha=0.12, label='Reward std')\n",
     "    if len(steps_r) > 8:\n",
     "        z = np.polyfit(steps_r, raw, 1)\n",
     "        p = np.poly1d(z)\n",
     "        ax_rew.plot(steps_r, p(steps_r), '--', color='white', alpha=0.35, linewidth=1.5,\n",
     "                    label=f'Trend ({z[0]:+.5f}/step)')\n",
-    "    ax_rew.set_xlabel('Reward Function Call')\n",
+    "    ax_rew.set_xlabel('Training Step')\n",
     "    ax_rew.set_ylabel('Reward Value')\n",
     "    ax_rew.legend(fontsize=9, facecolor=C['grid'], labelcolor=C['text'], framealpha=0.9, edgecolor=C['border'])\n",
     "    if np.var(raw) < 0.01:\n",
     "        ax_rew.text(0.5, 0.5, 'Low reward variance detected.\\nThis graph exposes weak learning signal.',\n",
     "                    transform=ax_rew.transAxes, ha='center', va='center', color=C['random'], fontsize=10,\n",
     "                    bbox=dict(boxstyle='round', facecolor=C['panel'], alpha=0.8))\n",
+    "elif training_rewards and len(training_rewards) >= 4:\n",
+    "    raw = training_rewards\n",
+    "    steps_r = list(range(1, len(raw) + 1))\n",
+    "    ax_rew.plot(steps_r, raw, alpha=0.20, color=C['reward'], linewidth=1)\n",
+    "    ax_rew.plot(steps_r, smooth(raw, window=6), color=C['reward'], linewidth=2.5, label='Smoothed reward')\n",
+    "    ax_rew.set_xlabel('Reward Function Call')\n",
+    "    ax_rew.set_ylabel('Reward Value')\n",
+    "    ax_rew.legend(fontsize=9, facecolor=C['grid'], labelcolor=C['text'], framealpha=0.9, edgecolor=C['border'])\n",
     "else:\n",
     "    ax_rew.text(0.5, 0.5, 'No training rewards captured.\\nRe-run with fixed reward function.',\n",
     "                transform=ax_rew.transAxes, ha='center', va='center', color=C['subtext'], fontsize=11)\n",
@@ -1027,6 +1150,92 @@
     "fig.savefig(dashboard_path, dpi=180, facecolor=fig.get_facecolor(), bbox_inches='tight')\n",
     "plt.close(fig)\n",
     "\n",
+    "# Standalone training reward curve for reports/slides.\n",
+    "reward_curve_path = 'results/gridmind_training_reward_curve.png'\n",
+    "fig_reward, ax_reward = plt.subplots(figsize=(11, 6))\n",
+    "fig_reward.patch.set_facecolor(C['bg'])\n",
+    "style_ax(ax_reward, 'Training Reward Curve')\n",
+    "if not training_metrics_df.empty and (\"reward\" in training_metrics_df.columns or \"rewards / reward_func / mean\" in training_metrics_df.columns):\n",
+    "    reward_col = \"reward\" if \"reward\" in training_metrics_df.columns else \"rewards / reward_func / mean\"\n",
+    "    std_col = \"reward_std\" if \"reward_std\" in training_metrics_df.columns else \"rewards / reward_func / std\"\n",
+    "    reward_df = training_metrics_df[[\"step\", reward_col] + ([std_col] if std_col in training_metrics_df.columns else [])].dropna(subset=[reward_col])\n",
+    "    xs = reward_df[\"step\"].astype(float).to_numpy()\n",
+    "    ys = reward_df[reward_col].astype(float).to_numpy()\n",
+    "    ax_reward.plot(xs, ys, color=C['reward'], alpha=0.35, linewidth=1.2, marker='o', markersize=2, label='Reward')\n",
+    "    ax_reward.plot(xs, smooth(ys.tolist(), window=6), color=C['trained'], linewidth=2.5, label='Smoothed reward')\n",
+    "    if std_col in reward_df.columns:\n",
+    "        std = reward_df[std_col].fillna(0).astype(float).to_numpy()\n",
+    "        ax_reward.fill_between(xs, ys - std, ys + std, color=C['reward'], alpha=0.12, label='Reward std')\n",
+    "    if len(xs) > 8:\n",
+    "        z = np.polyfit(xs, ys, 1)\n",
+    "        p = np.poly1d(z)\n",
+    "        ax_reward.plot(xs, p(xs), '--', color=C['text'], alpha=0.45, linewidth=1.5, label=f'Trend ({z[0]:+.5f}/step)')\n",
+    "    ax_reward.set_xlabel('Training Step')\n",
+    "    ax_reward.set_ylabel('Reward')\n",
+    "    ax_reward.legend(facecolor=C['grid'], labelcolor=C['text'], edgecolor=C['border'])\n",
+    "else:\n",
+    "    ax_reward.text(0.5, 0.5, 'No logged reward data available.', transform=ax_reward.transAxes,\n",
+    "                   ha='center', va='center', color=C['subtext'])\n",
+    "fig_reward.savefig(reward_curve_path, dpi=180, facecolor=fig_reward.get_facecolor(), bbox_inches='tight')\n",
+    "plt.close(fig_reward)\n",
+    "\n",
+    "# Reference-style simple plots from trainer.state.log_history.\n",
+    "log_history = trainer.state.log_history\n",
+    "simple_steps = []\n",
+    "simple_rewards = []\n",
+    "simple_losses = []\n",
+    "simple_loss_steps = []\n",
+    "\n",
+    "for entry in log_history:\n",
+    "    reward_key = \"reward\" if \"reward\" in entry else (\"rewards / reward_func / mean\" if \"rewards / reward_func / mean\" in entry else None)\n",
+    "    if reward_key is not None:\n",
+    "        simple_steps.append(entry.get(\"step\", len(simple_steps) + 1))\n",
+    "        simple_rewards.append(float(entry[reward_key]))\n",
+    "    if \"loss\" in entry:\n",
+    "        simple_loss_steps.append(entry.get(\"step\", len(simple_loss_steps) + 1))\n",
+    "        simple_losses.append(float(entry[\"loss\"]))\n",
+    "\n",
+    "# Plot 1: Reward over training\n",
+    "simple_reward_curve_path = \"plots/reward_curve.png\"\n",
+    "fig_simple_reward, ax_simple_reward = plt.subplots(1, 1, figsize=(10, 5))\n",
+    "if simple_rewards:\n",
+    "    ax_simple_reward.plot(simple_steps[:len(simple_rewards)], simple_rewards, color=\"#4285f4\", linewidth=2, label=\"GRPO Reward\")\n",
+    "    if len(simple_rewards) > 5:\n",
+    "        window = max(3, len(simple_rewards) // 10)\n",
+    "        smoothed = [\n",
+    "            sum(simple_rewards[max(0, i-window):i+1]) / len(simple_rewards[max(0, i-window):i+1])\n",
+    "            for i in range(len(simple_rewards))\n",
+    "        ]\n",
+    "        ax_simple_reward.plot(simple_steps[:len(smoothed)], smoothed, color=\"#ea4335\", linewidth=2, linestyle=\"--\", label=f\"Smoothed (window={window})\")\n",
+    "else:\n",
+    "    ax_simple_reward.text(0.5, 0.5, \"No reward logs found\", transform=ax_simple_reward.transAxes, ha=\"center\", va=\"center\")\n",
+    "ax_simple_reward.set_xlabel(\"Training Step\", fontsize=12)\n",
+    "ax_simple_reward.set_ylabel(\"Reward\", fontsize=12)\n",
+    "ax_simple_reward.set_title(\"GridMind-RL GRPO Training - Reward Curve\", fontsize=14, fontweight=\"bold\")\n",
+    "ax_simple_reward.legend()\n",
+    "ax_simple_reward.grid(True, alpha=0.3)\n",
+    "fig_simple_reward.tight_layout()\n",
+    "fig_simple_reward.savefig(simple_reward_curve_path, dpi=150)\n",
+    "plt.show()\n",
+    "print(f\"Saved: {simple_reward_curve_path}\")\n",
+    "\n",
+    "# Plot 2: Loss over training\n",
+    "simple_loss_curve_path = \"plots/loss_curve.png\"\n",
+    "if simple_losses:\n",
+    "    fig_simple_loss, ax_simple_loss = plt.subplots(1, 1, figsize=(10, 5))\n",
+    "    ax_simple_loss.plot(simple_loss_steps[:len(simple_losses)], simple_losses, color=\"#34a853\", linewidth=2)\n",
+    "    ax_simple_loss.set_xlabel(\"Training Step\", fontsize=12)\n",
+    "    ax_simple_loss.set_ylabel(\"Loss\", fontsize=12)\n",
+    "    ax_simple_loss.set_title(\"GridMind-RL GRPO Training - Loss Curve\", fontsize=14, fontweight=\"bold\")\n",
+    "    ax_simple_loss.grid(True, alpha=0.3)\n",
+    "    fig_simple_loss.tight_layout()\n",
+    "    fig_simple_loss.savefig(simple_loss_curve_path, dpi=150)\n",
+    "    plt.show()\n",
+    "    print(f\"Saved: {simple_loss_curve_path}\")\n",
+    "else:\n",
+    "    simple_loss_curve_path = None\n",
+    "    print(\"No loss logs found; skipped plots/loss_curve.png\")\n",
+    "\n",
     "# Separate before/after comparison graph for quick judge inspection.\n",
     "fig2, ax2 = plt.subplots(figsize=(11, 6))\n",
     "fig2.patch.set_facecolor(C['bg'])\n",
@@ -1048,6 +1257,10 @@
     "plt.close(fig2)\n",
     "\n",
     "print(f\"Saved dashboard graph to {dashboard_path}\")\n",
+    "print(f\"Saved training reward curve to {reward_curve_path}\")\n",
+    "print(f\"Saved simple reward curve to {simple_reward_curve_path}\")\n",
+    "if simple_loss_curve_path:\n",
+    "    print(f\"Saved simple loss curve to {simple_loss_curve_path}\")\n",
     "print(f\"Saved before/after graph to {comparison_path}\")\n",
     "\n",
     "results = {\n",
@@ -1066,8 +1279,12 @@
     "    \"training_rewards_log\": training_rewards[-20:] if training_rewards else [],\n",
     "    \"training_step_logs\": training_steps_log[-20:] if training_steps_log else [],\n",
     "    \"step_losses\": step_losses if 'step_losses' in globals() else [],\n",
+    "    \"training_metrics_table\": training_metrics_path,\n",
     "    \"graphs\": {\n",
     "        \"dashboard\": dashboard_path,\n",
+    "        \"training_reward_curve\": reward_curve_path,\n",
+    "        \"simple_reward_curve\": simple_reward_curve_path,\n",
+    "        \"simple_loss_curve\": simple_loss_curve_path,\n",
     "        \"before_after\": comparison_path,\n",
     "    },\n",
     "}\n",