feat: implement Unsloth GRPO training pipeline with environment-backed reward functions and balanced dataset generation

scripts/gridmind_grpo_colab.ipynb

@@ -272,88 +272,113 @@
    "outputs": [],
    "source": [
     "import json as _json\n",
-    "import requests as _requests\n",
-    "import random as _random\n",
     "import math as _math\n",
+    "import random as _random\n",
+    "import re as _re\n",
+    "import requests as _requests\n",
+    "import numpy as _np\n",
     "\n",
+    "training_rewards = []\n",
     "call_count = [0]\n",
+    "group_count = [0]\n",
+    "NUM_GENERATIONS_FOR_REWARD = 4\n",
+    "\n",
+    "_REQUIRED_ACTION_KEYS = {\"hvac_power_level\", \"thermal_charge_rate\", \"batch_job_slot\", \"load_shed_fraction\", \"building_id\"}\n",
+    "\n",
+    "def _extract_action(text):\n",
+    "    match = _re.search(r\"\\{.*?\\}\", text, _re.DOTALL)\n",
+    "    if not match:\n",
+    "        raise ValueError(\"completion did not contain a JSON object\")\n",
+    "    action = _json.loads(match.group())\n",
+    "    missing = _REQUIRED_ACTION_KEYS - set(action)\n",
+    "    if missing:\n",
+    "        raise ValueError(f\"missing action fields: {sorted(missing)}\")\n",
+    "    return {\n",
+    "        \"hvac_power_level\": float(max(0, min(1, action.get(\"hvac_power_level\", 0.5)))),\n",
+    "        \"thermal_charge_rate\": float(max(-1, min(1, action.get(\"thermal_charge_rate\", 0.0)))),\n",
+    "        \"batch_job_slot\": int(max(0, min(4, action.get(\"batch_job_slot\", 0)))),\n",
+    "        \"load_shed_fraction\": float(max(0, min(0.5, action.get(\"load_shed_fraction\", 0.0)))),\n",
+    "        \"building_id\": int(max(0, min(2, action.get(\"building_id\", 0)))),\n",
+    "    }\n",
     "\n",
     "def gridmind_reward_fn(completions, **kwargs):\n",
     "    \"\"\"\n",
-    "    Reward function for GridMind-RL GRPO training.\n",
-    "    - Parses JSON action from LLM output\n",
-    "    - Executes against environment\n",
-    "    - Returns normalized reward signal\n",
+    "    Environment-backed GRPO reward.\n",
+    "    Generations from the same prompt are evaluated on the same task/seed, so\n",
+    "    advantages reflect real action quality instead of random episode noise.\n",
     "    \"\"\"\n",
     "    rewards = []\n",
-    "    task_id = _random.choice([1, 2, 3, 4])\n",
+    "    batch_start = group_count[0]\n",
     "\n",
-    "    try:\n",
-    "        _requests.post(f\"{ENV_URL}/reset\", json={\"task_id\": task_id}, timeout=10)\n",
-    "    except:\n",
-    "        return [-0.1] * len(completions)\n",
+    "    for i, completion in enumerate(completions):\n",
+    "        call_count[0] += 1\n",
+    "        group_id = batch_start + (i // NUM_GENERATIONS_FOR_REWARD)\n",
+    "        text = completion[0][\"content\"] if isinstance(completion, list) else completion\n",
     "\n",
-    "    for completion in completions:\n",
     "        try:\n",
-    "            text = str(completion[0]) if isinstance(completion, list) and completion else str(completion)\n",
-    "            text = text.strip()\n",
-    "\n",
-    "            # Extract JSON from completion\n",
-    "            start = text.rfind('{')\n",
-    "            end = text.rfind('}') + 1\n",
-    "            if start < 0 or end <= start:\n",
-    "                rewards.append(-0.3)\n",
-    "                try:\n",
-    "                    _requests.post(f\"{ENV_URL}/reset\", json={\"task_id\": task_id}, timeout=6)\n",
-    "                except:\n",
-    "                    pass\n",
-    "                continue\n",
+    "            action = _extract_action(text)\n",
+    "        except _json.JSONDecodeError:\n",
+    "            reward = -0.8\n",
+    "            rewards.append(reward)\n",
+    "            training_rewards.append(reward)\n",
+    "            continue\n",
+    "        except ValueError:\n",
+    "            reward = -1.0\n",
+    "            rewards.append(reward)\n",
+    "            training_rewards.append(reward)\n",
+    "            continue\n",
+    "\n",
+    "        task_id = (group_id % 4) + 1\n",
+    "        seed = 10_000 + group_id\n",
     "\n",
-    "            try:\n",
-    "                action = _json.loads(text[start:end])\n",
-    "            except _json.JSONDecodeError:\n",
-    "                rewards.append(-0.2)\n",
-    "                try:\n",
-    "                    _requests.post(f\"{ENV_URL}/reset\", json={\"task_id\": task_id}, timeout=6)\n",
-    "                except:\n",
-    "                    pass\n",
-    "                continue\n",
-    "\n",
-    "            # Validate and clamp action fields\n",
-    "            cleaned = {\n",
-    "                \"hvac_power_level\": max(0.0, min(1.0, float(action.get(\"hvac_power_level\", 0.5)))),\n",
-    "                \"thermal_charge_rate\": max(-1.0, min(1.0, float(action.get(\"thermal_charge_rate\", 0.0)))),\n",
-    "                \"batch_job_slot\": max(0, min(4, int(action.get(\"batch_job_slot\", 0)))),\n",
-    "                \"load_shed_fraction\": max(0.0, min(0.5, float(action.get(\"load_shed_fraction\", 0.0)))),\n",
-    "                \"building_id\": int(action.get(\"building_id\", 0)),\n",
-    "            }\n",
+    "        try:\n",
+    "            reset_resp = _requests.post(\n",
+    "                f\"{ENV_URL}/reset\",\n",
+    "                json={\"task_id\": task_id, \"seed\": seed, \"num_buildings\": 1},\n",
+    "                timeout=15,\n",
+    "            )\n",
+    "            reset_resp.raise_for_status()\n",
+    "        except Exception:\n",
+    "            reward = -0.5\n",
+    "            rewards.append(reward)\n",
+    "            training_rewards.append(reward)\n",
+    "            continue\n",
     "\n",
-    "            try:\n",
-    "                step_r = _requests.post(f\"{ENV_URL}/step\", json=cleaned, timeout=8)\n",
-    "                data = step_r.json()\n",
+    "        total_env_reward = 0.0\n",
+    "        completed_steps = 0\n",
+    "        try:\n",
+    "            for _ in range(8):\n",
+    "                step_resp = _requests.post(f\"{ENV_URL}/step\", json=action, timeout=15)\n",
+    "                step_resp.raise_for_status()\n",
+    "                data = step_resp.json()\n",
     "                if isinstance(data, list):\n",
     "                    data = data[0]\n",
-    "                env_reward = float(data.get(\"reward\", 0.0))\n",
-    "                reward_signal = _math.tanh(env_reward * 1.5) * 0.5\n",
-    "                rewards.append(reward_signal)\n",
-    "            except:\n",
-    "                rewards.append(-0.15)\n",
-    "\n",
-    "            try:\n",
-    "                _requests.post(f\"{ENV_URL}/reset\", json={\"task_id\": task_id}, timeout=6)\n",
-    "            except:\n",
-    "                pass\n",
-    "\n",
+    "                if \"data\" in data and isinstance(data[\"data\"], dict):\n",
+    "                    data = data[\"data\"]\n",
+    "                total_env_reward += float(data.get(\"reward\", 0.0) or 0.0)\n",
+    "                completed_steps += 1\n",
+    "                if data.get(\"done\", False):\n",
+    "                    break\n",
+    "\n",
+    "            avg_step_reward = total_env_reward / max(completed_steps, 1)\n",
+    "            normalized_step_reward = max(-1.0, min(1.0, avg_step_reward / 10.0))\n",
+    "            grade_resp = _requests.get(f\"{ENV_URL}/grade\", timeout=15)\n",
+    "            if grade_resp.status_code == 200:\n",
+    "                normalized_grade = max(0.0, min(1.0, float(grade_resp.json().get(\"score\", 0.0))))\n",
+    "                reward = 0.7 * normalized_grade + 0.3 * normalized_step_reward\n",
+    "            else:\n",
+    "                reward = normalized_step_reward\n",
     "        except Exception:\n",
-    "            rewards.append(-0.15)\n",
+    "            reward = -0.5\n",
     "\n",
-    "    call_count[0] += 1\n",
-    "    if call_count[0] % 5 == 0:\n",
-    "        print(f\"  Step {call_count[0]}: Avg reward = {sum(rewards)/len(rewards):+.3f}\")\n",
+    "        rewards.append(reward)\n",
+    "        training_rewards.append(reward)\n",
+    "\n",
+    "    group_count[0] += _math.ceil(len(completions) / NUM_GENERATIONS_FOR_REWARD)\n",
     "\n",
     "    return rewards\n",
     "\n",
-    "print(\"Reward function ready\")"
+    "print(\"Environment-backed reward function ready\")\n"
    ]
   },
   {
@@ -373,6 +398,7 @@
    "source": [
     "from trl import GRPOTrainer, GRPOConfig\n",
     "from peft import LoraConfig, prepare_model_for_kbit_training\n",
+    "from transformers import PrinterCallback, TrainerCallback\n",
     "import inspect\n",
     "import os\n",
     "\n",
@@ -384,30 +410,106 @@
     "peft_config = LoraConfig(\n",
     "    r=16,\n",
     "    lora_alpha=32,\n",
-    "    target_modules=[\"q_proj\", \"v_proj\", \"k_proj\", \"o_proj\", \"gate_proj\", \"up_proj\", \"down_proj\"],\n",
+    "    target_modules=[\"q_proj\", \"k_proj\", \"v_proj\", \"o_proj\", \"gate_proj\", \"up_proj\", \"down_proj\"],\n",
     "    lora_dropout=0.05,\n",
     "    bias=\"none\",\n",
     "    task_type=\"CAUSAL_LM\",\n",
     ")\n",
     "\n",
-    "# Configure GRPO training\n",
-    "# Note: Disable AMP (fp16/bf16) when using quantized models to avoid gradient scaler issues\n",
+    "class MetricsTableCallback(TrainerCallback):\n",
+    "    columns = [\n",
+    "        (\"step\", \"Step\", 6),\n",
+    "        (\"loss\", \"Loss\", 10),\n",
+    "        (\"reward\", \"Reward\", 10),\n",
+    "        (\"reward_std\", \"RewardStd\", 10),\n",
+    "        (\"entropy\", \"Entropy\", 10),\n",
+    "        (\"learning_rate\", \"LR\", 11),\n",
+    "        (\"num_tokens\", \"Tokens\", 8),\n",
+    "        (\"step_time\", \"StepTime\", 10),\n",
+    "    ]\n",
+    "\n",
+    "    def __init__(self):\n",
+    "        self.header_printed = False\n",
+    "        self.rewards = []\n",
+    "\n",
+    "    def _format_value(self, key, value):\n",
+    "        if value is None:\n",
+    "            return \"-\"\n",
+    "        try:\n",
+    "            if key in {\"step\", \"num_tokens\"}:\n",
+    "                return str(int(float(value)))\n",
+    "            if key == \"learning_rate\":\n",
+    "                return f\"{float(value):.2e}\"\n",
+    "            return f\"{float(value):.4f}\"\n",
+    "        except (TypeError, ValueError):\n",
+    "            return str(value)\n",
+    "\n",
+    "    def _print_header(self):\n",
+    "        separator = \"+\" + \"+\".join(\"-\" * (width + 2) for _, _, width in self.columns) + \"+\"\n",
+    "        header = \"|\" + \"|\".join(f\" {title:<{width}} \" for _, title, width in self.columns) + \"|\"\n",
+    "        print(separator)\n",
+    "        print(header)\n",
+    "        print(separator)\n",
+    "        self.header_printed = True\n",
+    "\n",
+    "    def on_log(self, args, state, control, logs=None, **kwargs):\n",
+    "        if not logs or (\"loss\" not in logs and \"reward\" not in logs):\n",
+    "            return\n",
+    "        if not self.header_printed:\n",
+    "            self._print_header()\n",
+    "        row_values = []\n",
+    "        for key, _, width in self.columns:\n",
+    "            value = state.global_step if key == \"step\" else logs.get(key)\n",
+    "            row_values.append(f\" {self._format_value(key, value):>{width}} \")\n",
+    "        print(\"|\" + \"|\".join(row_values) + \"|\")\n",
+    "\n",
+    "        if \"reward\" in logs:\n",
+    "            try:\n",
+    "                self.rewards.append(float(logs[\"reward\"]))\n",
+    "            except (TypeError, ValueError):\n",
+    "                pass\n",
+    "\n",
+    "    def on_train_end(self, args, state, control, **kwargs):\n",
+    "        if not self.rewards:\n",
+    "            return\n",
+    "        first_window = self.rewards[: min(5, len(self.rewards))]\n",
+    "        last_window = self.rewards[-min(5, len(self.rewards)) :]\n",
+    "        first_avg = float(_np.mean(first_window))\n",
+    "        last_avg = float(_np.mean(last_window))\n",
+    "        overall_avg = float(_np.mean(self.rewards))\n",
+    "        best_reward = float(_np.max(self.rewards))\n",
+    "        print(\"+----------------------+------------+\")\n",
+    "        print(\"| Reward Summary       | Value      |\")\n",
+    "        print(\"+----------------------+------------+\")\n",
+    "        print(f\"| Logged rows          | {len(self.rewards):>10} |\")\n",
+    "        print(f\"| First rows avg       | {first_avg:>+10.4f} |\")\n",
+    "        print(f\"| Last rows avg        | {last_avg:>+10.4f} |\")\n",
+    "        print(f\"| Improvement          | {last_avg - first_avg:>+10.4f} |\")\n",
+    "        print(f\"| Overall avg          | {overall_avg:>+10.4f} |\")\n",
+    "        print(f\"| Best row reward      | {best_reward:>+10.4f} |\")\n",
+    "        print(\"+----------------------+------------+\")\n",
+    "\n",
+    "# GRPO config - stable for T4 / Colab\n",
+    "output_dir = \"gridmind-grpo-trained\"\n",
+    "os.makedirs(output_dir, exist_ok=True)\n",
+    "\n",
     "grpo_config_dict = {\n",
-    "    \"output_dir\": \"./gridmind-grpo-output\",\n",
+    "    \"output_dir\": output_dir,\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",
+    "    \"max_prompt_length\": 512,\n",
+    "    \"max_completion_length\": 80,\n",
     "    \"learning_rate\": 5e-5,\n",
+    "    \"lr_scheduler_type\": \"cosine\",\n",
+    "    \"warmup_ratio\": 0.1,\n",
     "    \"fp16\": False,\n",
     "    \"bf16\": False,\n",
     "    \"max_grad_norm\": 0.0,\n",
-    "    \"logging_steps\": 1,\n",
-    "    \"log_completions\": True,\n",
-    "    \"num_completions_to_print\": 1,\n",
+    "    \"logging_steps\": 5,\n",
+    "    \"log_completions\": False,\n",
     "    \"save_steps\": 60,\n",
     "    \"report_to\": \"none\",\n",
     "    \"disable_tqdm\": True,\n",
@@ -434,14 +536,16 @@
     "    train_dataset=dataset,\n",
     "    reward_funcs=gridmind_reward_fn,\n",
     "    peft_config=peft_config,\n",
+    "    callbacks=[MetricsTableCallback()],\n",
     ")\n",
+    "trainer.remove_callback(PrinterCallback)\n",
     "\n",
     "print(\"\\nStarting GRPO training (estimated 25-35 min on T4)...\\n\")\n",
     "train_result = trainer.train()\n",
     "\n",
-    "print(f\"\\n✔ Training complete!\")\n",
+    "print(f\"\\nTraining complete!\")\n",
     "print(f\"  Total steps: {train_result.global_step}\")\n",
-    "print(f\"  Final loss: {train_result.training_loss:.6f}\")"
+    "print(f\"  Final loss: {train_result.training_loss:.6f}\")\n"
    ]
   },
   {

scripts/train_unsloth.py

@@ -33,7 +33,7 @@ import matplotlib.gridspec as gridspec
 from datasets import Dataset
 from trl import GRPOTrainer, GRPOConfig
 from unsloth import FastLanguageModel
-from transformers import TrainerCallback
+from transformers import PrinterCallback, TrainerCallback
 
 os.makedirs("results", exist_ok=True)
 
@@ -65,69 +65,39 @@ def make_prompt(i):
     }]
 
 
-def reward_valid_json(completions, **kwargs):
-    rewards = []
-    for completion in completions:
-        text = completion[0]["content"] if isinstance(completion, list) else completion
-        try:
-            match = re.search(r'\{.*?\}', text, re.DOTALL)
-            if match:
-                json.loads(match.group())
-                rewards.append(0.3)
-            else:
-                rewards.append(0.0)
-        except Exception:
-            rewards.append(0.0)
-    return rewards
-
-
-def reward_has_required_keys(completions, **kwargs):
-    required = {"hvac_power_level", "thermal_charge_rate", "batch_job_slot", "load_shed_fraction"}
-    rewards = []
-    for completion in completions:
-        text = completion[0]["content"] if isinstance(completion, list) else completion
-        try:
-            match = re.search(r'\{.*?\}', text, re.DOTALL)
-            if match:
-                action = json.loads(match.group())
-                if required.issubset(action.keys()):
-                    rewards.append(0.3)
-                else:
-                    rewards.append(0.1)
-            else:
-                rewards.append(0.0)
-        except Exception:
-            rewards.append(0.0)
-    return rewards
-
-
 ENV_URL = "https://prajwal782007-gridmind.hf.space"
 
 
 class GridMindRewardFn:
-    """Fixed reward function with environment reset per completion call."""
+    """Environment-backed reward function with comparable rollouts per GRPO group."""
     
-    def __init__(self, env_url, num_steps=8):
+    def __init__(self, env_url, num_steps=8, num_generations=4):
         self.env_url = env_url
         self.num_steps = num_steps
+        self.num_generations = max(1, num_generations)
         self.call_count = [0]
         self.reward_variance_log = []
         self.training_rewards = []
+        self.group_count = 0
     
     def __call__(self, completions, **kwargs):
         rewards = []
         batch_rewards = []
         
+        batch_start = self.group_count
         for i, completion in enumerate(completions):
             self.call_count[0] += 1
+            group_id = batch_start + (i // self.num_generations)
             
             text = completion[0]["content"] if isinstance(completion, list) else completion
             
             try:
                 match = re.search(r'\{.*?\}', text, re.DOTALL)
                 if not match:
-                    rewards.append(-1.0)
-                    batch_rewards.append(-1.0)
+                    final_reward = -1.0
+                    rewards.append(final_reward)
+                    batch_rewards.append(final_reward)
+                    self.training_rewards.append(final_reward)
                     continue
                 
                 action = json.loads(match.group())
@@ -140,8 +110,10 @@ class GridMindRewardFn:
                     "building_id": 0
                 }
                 
-                seed = 1000 + self.call_count[0]
-                task_id = (self.call_count[0] % 3) + 1
+                # Evaluate all generations for the same prompt on the same scenario.
+                # This keeps GRPO advantages tied to action quality instead of seed noise.
+                seed = 10_000 + group_id
+                task_id = (group_id % 4) + 1
                 
                 reset_resp = requests.post(
                     f"{self.env_url}/reset",
@@ -149,11 +121,14 @@ class GridMindRewardFn:
                     timeout=30
                 )
                 if reset_resp.status_code != 200:
-                    rewards.append(-0.5)
-                    batch_rewards.append(-0.5)
+                    final_reward = -0.5
+                    rewards.append(final_reward)
+                    batch_rewards.append(final_reward)
+                    self.training_rewards.append(final_reward)
                     continue
                 
                 total_reward = 0.0
+                completed_steps = 0
                 for _ in range(self.num_steps):
                     step_resp = requests.post(
                         f"{self.env_url}/step",
@@ -166,36 +141,39 @@ class GridMindRewardFn:
                     if isinstance(step_data, list):
                         step_data = step_data[0]
                     total_reward += float(step_data.get("reward", 0))
+                    completed_steps += 1
                 
-                avg_reward = total_reward / self.num_steps if self.num_steps > 0 else 0
+                avg_step_reward = total_reward / max(completed_steps, 1)
+                normalized_step_reward = max(-1.0, min(1.0, avg_step_reward / 10.0))
                 
                 grade_resp = requests.get(f"{self.env_url}/grade", timeout=30)
                 if grade_resp.status_code == 200:
                     episode_score = float(grade_resp.json().get("score", 0.5))
-                    normalized = max(0.0, min(1.0, (episode_score - 0.4) / 0.32))
-                    final_reward = normalized
+                    normalized_grade = max(0.0, min(1.0, episode_score))
+                    final_reward = 0.7 * normalized_grade + 0.3 * normalized_step_reward
                 else:
-                    final_reward = max(-1.0, min(1.0, avg_reward / 10.0))
+                    final_reward = normalized_step_reward
                 
                 rewards.append(final_reward)
                 batch_rewards.append(final_reward)
                 self.training_rewards.append(final_reward)
                 
             except json.JSONDecodeError:
-                rewards.append(-0.8)
-                batch_rewards.append(-0.8)
+                final_reward = -0.8
+                rewards.append(final_reward)
+                batch_rewards.append(final_reward)
+                self.training_rewards.append(final_reward)
             except Exception as e:
                 print(f"Reward error: {e}", file=sys.stderr)
-                rewards.append(-0.5)
-                batch_rewards.append(-0.5)
+                final_reward = -0.5
+                rewards.append(final_reward)
+                batch_rewards.append(final_reward)
+                self.training_rewards.append(final_reward)
         
-        if len(batch_rewards) > 1 and self.call_count[0] % 10 == 0:
-            try:
-                variance = np.var(batch_rewards)
-                print(f"  [Step {self.call_count[0]}] Reward variance: {variance:.4f} | Avg: {np.mean(batch_rewards):.3f}")
-                self.reward_variance_log.append(variance)
-            except:
-                pass
+        self.group_count += math.ceil(len(completions) / self.num_generations)
+        
+        if len(batch_rewards) > 1:
+            self.reward_variance_log.append(float(np.var(batch_rewards)))
         
         return rewards
 
@@ -627,6 +605,85 @@ class CSVLogCallback(TrainerCallback):
             pd.DataFrame(self.log_history).to_csv(self.output_path, index=False)
 
 
+class MetricsTableCallback(TrainerCallback):
+    """Print compact GRPO metrics without dumping prompts or completions."""
+
+    columns = [
+        ("step", "Step", 6),
+        ("loss", "Loss", 10),
+        ("reward", "Reward", 10),
+        ("reward_std", "RewardStd", 10),
+        ("entropy", "Entropy", 10),
+        ("learning_rate", "LR", 11),
+        ("num_tokens", "Tokens", 8),
+        ("step_time", "StepTime", 10),
+    ]
+
+    def __init__(self):
+        self.header_printed = False
+        self.rewards = []
+
+    def _format_value(self, key, value):
+        if value is None:
+            return "-"
+        try:
+            if key in {"step", "num_tokens"}:
+                return str(int(float(value)))
+            if key == "learning_rate":
+                return f"{float(value):.2e}"
+            return f"{float(value):.4f}"
+        except (TypeError, ValueError):
+            return str(value)
+
+    def _print_header(self):
+        separator = "+" + "+".join("-" * (width + 2) for _, _, width in self.columns) + "+"
+        header = "|" + "|".join(f" {title:<{width}} " for _, title, width in self.columns) + "|"
+        print(separator)
+        print(header)
+        print(separator)
+        self.header_printed = True
+
+    def on_log(self, args, state, control, logs=None, **kwargs):
+        if not logs or ("loss" not in logs and "reward" not in logs):
+            return
+        if not self.header_printed:
+            self._print_header()
+
+        row_values = []
+        for key, _, width in self.columns:
+            value = state.global_step if key == "step" else logs.get(key)
+            row_values.append(f" {self._format_value(key, value):>{width}} ")
+        print("|" + "|".join(row_values) + "|")
+
+        if "reward" in logs:
+            try:
+                self.rewards.append(float(logs["reward"]))
+            except (TypeError, ValueError):
+                pass
+
+    def on_train_end(self, args, state, control, **kwargs):
+        if not self.rewards:
+            return
+
+        first_window = self.rewards[: min(5, len(self.rewards))]
+        last_window = self.rewards[-min(5, len(self.rewards)) :]
+        first_avg = float(np.mean(first_window))
+        last_avg = float(np.mean(last_window))
+        overall_avg = float(np.mean(self.rewards))
+        best_reward = float(np.max(self.rewards))
+
+        print("+----------------------+------------+")
+        print("| Reward Summary       | Value      |")
+        print("+----------------------+------------+")
+        print(f"| Logged rows          | {len(self.rewards):>10} |")
+        print(f"| First rows avg       | {first_avg:>+10.4f} |")
+        print(f"| Last rows avg        | {last_avg:>+10.4f} |")
+        print(f"| Improvement          | {last_avg - first_avg:>+10.4f} |")
+        print(f"| Overall avg          | {overall_avg:>+10.4f} |")
+        print(f"| Best row reward      | {best_reward:>+10.4f} |")
+        print("+----------------------+------------+")
+
+
 def main():
     parser = argparse.ArgumentParser(description="Train GridMind-RL agent with Unsloth GRPO")
     parser.add_argument("--env-url", type=str, default="http://localhost:7860", help="OpenEnv server URL")
@@ -683,9 +740,8 @@ def main():
         "learning_rate": 5e-6,  # FIXED: was 5e-5, too high
         "lr_scheduler_type": "cosine",
         "warmup_ratio": 0.1,
-        "logging_steps": 1,  # Log every step to produce dense table
-        "log_completions": True,  # Enable completion metrics in table
-        "num_completions_to_print": 1,  # Print 1 completion per step
+        "logging_steps": 5,  # Keep 60-step output clean: rows at 5, 10, ..., 60
+        "log_completions": False,
         "save_steps": 100,
         "fp16": False,  # Disable AMP with quantized models (avoid grad scaler issues)
         "bf16": False,
@@ -708,20 +764,21 @@ def main():
         print(f"Skipping unsupported GRPOConfig args: {skipped_training_args}")
     training_args = GRPOConfig(**training_arg_kwargs)
     
-    reward_fn = GridMindRewardFn(args.env_url, num_steps=8)
+    reward_fn = GridMindRewardFn(
+        args.env_url,
+        num_steps=8,
+        num_generations=requested_training_args["num_generations"],
+    )
     
     trainer = GRPOTrainer(
         model=model,
         processing_class=tokenizer,
         args=training_args,
         train_dataset=dataset,
-        reward_funcs=[
-            reward_valid_json,
-            reward_has_required_keys,
-            reward_fn,
-        ],
-        callbacks=[CSVLogCallback(args.output_csv)]
+        reward_funcs=[reward_fn],
+        callbacks=[CSVLogCallback(args.output_csv), MetricsTableCallback()]
     )
+    trainer.remove_callback(PrinterCallback)
     
     print("🚀 Starting GRPO training...")
     trainer.train()