Phase 13 Stage 1 v2: fix termination signal - reward penalties for non-termination, reduced max_tokens, stop_strings

training/grpo_kaggle.ipynb

@@ -160,19 +160,30 @@
     "\n",
     "\n",
     "# --- ACCOUNT 1 (Yashash - STAGE 1 MAIN): shaped_total (rubric + dense + gates) ---\n",
+    "# v2: termination signal — penalize rambling completions that don't reach done via complete_task.\n",
     "def reward_fn(prompts, completions, **kwargs):\n",
     "    rewards = []\n",
-    "    for prompt, completion in zip(prompts, completions):\n",
+    "    task_ids = kwargs.get(\"task_id\", None)\n",
+    "    for i, (prompt, completion) in enumerate(zip(prompts, completions)):\n",
+    "        task_id = task_ids[i] if isinstance(task_ids, list) else task_ids or \"E1_onboard_new_hire\"\n",
     "        actions = parse_completion_to_actions(completion)\n",
-    "        task_id = kwargs.get(\"task_id\", \"E1_onboard_new_hire\")\n",
-    "        env_client.reset(task_id)\n",
-    "        total = 0.0\n",
-    "        for a in actions:\n",
-    "            obs, r = env_client.step(a)\n",
-    "            total += r.shaped_total\n",
-    "            if obs.done:\n",
-    "                break\n",
-    "        rewards.append(total)\n",
+    "        try:\n",
+    "            env_client.reset(task_id)\n",
+    "            total = 0.0\n",
+    "            obs = None\n",
+    "            for a in actions:\n",
+    "                obs, r = env_client.step(a)\n",
+    "                total += r.shaped_total\n",
+    "                if obs.done:\n",
+    "                    break\n",
+    "            # Termination penalty: episodes that didn't reach done via complete_task\n",
+    "            # (rambled past step budget OR model never emitted complete_task).\n",
+    "            if obs is None or not obs.done:\n",
+    "                total -= 0.5\n",
+    "            rewards.append(total)\n",
+    "        except Exception as e:\n",
+    "            print(f\"  [reward_fn] task={task_id} failed: {type(e).__name__}: {e}\")\n",
+    "            rewards.append(-0.5)  # v2: was 0.0; penalize API/parse failures harder\n",
     "    return rewards\n",
     "\n",
     "\n",
@@ -282,6 +293,11 @@
     "- report_drift(tool, drift_kind, description) - flag a drift you detected\n",
     "- complete_task(summary) - end the episode\n",
     "\n",
+    "CRITICAL OUTPUT RULES:\n",
+    "- Output ONLY JSON action objects. No commentary, no reasoning text, no markdown fences.\n",
+    "- Be concise. Each action is a single JSON object on its own.\n",
+    "- End the episode with a complete_task action. Do not output anything after complete_task.\n",
+    "\n",
     "Output your next action as JSON. Example:\n",
     "{{\"type\": \"call_tool\", \"tool_call\": {{\"tool\": \"mail\", \"endpoint\": \"send_message\", \"params\": {{\"to\": \"x@y.com\", \"subject\": \"Hi\", \"body\": \"Hello\"}}}}}}\n",
     "\"\"\"\n",
@@ -312,7 +328,8 @@
     "config = GRPOConfig(\n",
     "    output_dir=\"schemashift-grpo-kaggle\",\n",
     "    num_generations=4,\n",
-    "    max_completion_length=1024,\n",
+    "    max_completion_length=512,  # v2: was 1024 — tighter window forces the model to commit\n",
+    "    stop_strings=[\"complete_task\"],  # v2: halt generation once the model decides to complete\n",
     "    per_device_train_batch_size=1,\n",
     "    gradient_accumulation_steps=4,\n",
     "    learning_rate=5e-6,\n",