Fix zero rewards: layered reward function with keyword matching + skip identical-reward steps

GRPO_Training_Pipeline.ipynb

@@ -80,7 +80,7 @@
   },
   {
    "cell_type": "code",
-   "source": "###############################################################\n# REAL GRPO TRAINING\n# model generates actions -> env rewards -> advantages -> backprop\n###############################################################\n\nimport torch, requests, json, time, gc\nimport torch.nn.functional as F\nimport matplotlib\nmatplotlib.use(\"Agg\")\nimport matplotlib.pyplot as plt\nfrom unsloth import FastLanguageModel\n\nENV = \"https://cobedigger-dataengenv.hf.space\"\n\nG           = 4      # completions per prompt\nTRAIN_STEPS = 25     # gradient steps\nLR          = 2e-5\nMAX_NEW     = 200\nMAX_PROMPT  = 400\nMAX_FULL    = MAX_PROMPT + MAX_NEW  # 600\n\nSYSTEM = (\n    \"You are an ML engineer debugging a broken pipeline. \"\n    \"Respond with ONLY a JSON object. No explanation, no markdown.\\n\"\n    \"Actions: inspect_data | run_script | edit_script | query_actor | submit\\n\"\n    \"edit_script payload must have 'old' and 'new' keys.\\n\\n\"\n    \"Stage 1: rename age_years to age AND add df.dropna()\\n\"\n    \"Stage 2: add StandardScaler fitted on X_train AFTER train_test_split\\n\"\n    \"Stage 3: move scaler.fit to AFTER train_test_split\\n\"\n    \"Stage 4: add class_weight='balanced' to LogisticRegression\"\n)\n\nSTAGE_PROMPTS = {\n    1: \"Stage 1 | KeyError: 'age_years'\\nColumn is 'age' not 'age_years'. Also 12 NaN rows. What is your action?\",\n    2: \"Stage 2 | MLPClassifier loss is NaN after 5 iters\\nFeatures not normalised. What is your action?\",\n    3: \"Stage 3 | Accuracy=0.98 — data leakage suspected\\nscaler.fit_transform called before train_test_split. What is your action?\",\n    4: \"Stage 4 | fairness_gap=0.18 > threshold 0.05\\nMinority class poorly predicted. What is your action?\",\n}\n\ndef score_action(action_str, stage):\n    try:\n        s = action_str.find(\"{\"); e = action_str.rfind(\"}\")\n        if s == -1 or e == -1: return 0.0\n        action = json.loads(action_str[s:e+1])\n        if \"action_type\" not in action: return 0.0\n        action.setdefault(\"payload\", {})\n        if isinstance(action[\"payload\"], str):\n            action[\"payload\"] = {\"script\": action[\"payload\"]}\n    except Exception:\n        return 0.0\n\n    atype = action[\"action_type\"]\n    if atype == \"inspect_data\": return 0.15\n    if atype == \"query_actor\" and stage == 4: return 0.20\n    if atype == \"run_script\": return 0.10\n    if atype == \"submit\": return 0.05\n    if atype == \"edit_script\":\n        try:\n            requests.post(f\"{ENV}/reset\", timeout=30)\n            r = requests.post(f\"{ENV}/step\", json=action, timeout=30).json()\n            score = float(r.get(\"reward\", {}).get(\"score\", 0.0))\n            return min(1.0, score + 0.1)\n        except Exception:\n            return 0.05\n    return 0.0\n\ndef make_prompt(stage):\n    msgs = [{\"role\": \"system\", \"content\": SYSTEM},\n            {\"role\": \"user\",   \"content\": STAGE_PROMPTS[stage]}]\n    return tokenizer.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True)\n\ndef reinforce_loss(input_ids, attention_mask, advantages):\n    \"\"\"Pure REINFORCE: maximise advantage-weighted log-prob of completions.\"\"\"\n    logits = model(input_ids=input_ids, attention_mask=attention_mask).logits\n    shift_logits = logits[..., :-1, :].contiguous()\n    shift_labels = input_ids[..., 1:].contiguous()\n\n    lp = F.log_softmax(shift_logits, dim=-1)\n    tok_lp = lp.gather(-1, shift_labels.unsqueeze(-1)).squeeze(-1)\n\n    mask   = attention_mask[..., 1:].float()\n    seq_lp = (tok_lp * mask).sum(-1) / mask.sum(-1).clamp(min=1)\n\n    loss = -(advantages * seq_lp).mean()\n    return loss\n\n# ── Setup ────────────────────────────────────────────────────\nprint(\"Setting up GRPO training...\")\ntorch.cuda.empty_cache(); gc.collect()\nFastLanguageModel.for_training(model)\n\noptimizer = torch.optim.AdamW(\n    [p for p in model.parameters() if p.requires_grad], lr=LR\n)\n\ngrpo_rewards, grpo_losses = [], []\nprint(f\"\\nStarting GRPO: {TRAIN_STEPS} steps, G={G} completions/step\")\nprint(\"=\" * 60)\n\nfor step in range(TRAIN_STEPS):\n    stage  = [1, 2, 3, 4][step % 4]\n    prompt = make_prompt(stage)\n\n    enc = tokenizer(\n        prompt, return_tensors=\"pt\",\n        truncation=True, max_length=MAX_PROMPT\n    ).to(\"cuda\")\n    prompt_len = enc[\"input_ids\"].shape[1]\n\n    # 1. Generate G completions in inference mode\n    FastLanguageModel.for_inference(model)\n    model.eval()\n    with torch.no_grad():\n        outs = model.generate(\n            **enc,\n            max_new_tokens=MAX_NEW,\n            do_sample=True,\n            temperature=0.9,\n            num_return_sequences=G,\n            pad_token_id=tokenizer.eos_token_id,\n            repetition_penalty=1.1,\n        )\n    completions = [\n        tokenizer.decode(outs[i][prompt_len:], skip_special_tokens=True)\n        for i in range(G)\n    ]\n\n    # 2. Score each via live env\n    rewards = [score_action(c, stage) for c in completions]\n    time.sleep(0.3)\n\n    rewards_t  = torch.tensor(rewards, dtype=torch.float32)\n    mean_r     = rewards_t.mean()\n    std_r      = rewards_t.std().clamp(min=1e-6)\n    advantages = ((rewards_t - mean_r) / std_r).to(\"cuda\")\n\n    # 3. Tokenise prompt + completion pairs\n    full_ids_list, full_mask_list = [], []\n    for c in completions:\n        enc_full = tokenizer(\n            prompt + c, return_tensors=\"pt\",\n            truncation=True, padding=\"max_length\", max_length=MAX_FULL,\n        )\n        full_ids_list.append(enc_full[\"input_ids\"])\n        full_mask_list.append(enc_full[\"attention_mask\"])\n    full_ids  = torch.cat(full_ids_list,  dim=0).to(\"cuda\")\n    full_mask = torch.cat(full_mask_list, dim=0).to(\"cuda\")\n\n    # 4. REINFORCE loss + backprop (training mode)\n    FastLanguageModel.for_training(model)\n    model.train()\n    loss = reinforce_loss(full_ids, full_mask, advantages)\n    loss.backward()\n    torch.nn.utils.clip_grad_norm_(\n        [p for p in model.parameters() if p.requires_grad], 1.0\n    )\n    optimizer.step()\n    optimizer.zero_grad()\n    torch.cuda.empty_cache()\n\n    avg_r = rewards_t.mean().item()\n    grpo_rewards.append(avg_r)\n    grpo_losses.append(loss.item())\n    print(\n        f\"  Step {step+1:02d}/{TRAIN_STEPS} | Stage {stage} | \"\n        f\"Rewards={[f'{r:.2f}' for r in rewards]} | \"\n        f\"Avg={avg_r:.3f} | Loss={loss.item():.4f}\"\n    )\n\nprint(\"\\nGRPO training complete!\")\nprint(f\"   First 5 avg reward : {sum(grpo_rewards[:5])/5:.3f}\")\nprint(f\"   Last  5 avg reward : {sum(grpo_rewards[-5:])/5:.3f}\")\n\n# ── Plot ─────────────────────────────────────────────────────\nfig, (ax1, ax2) = plt.subplots(1, 2, figsize=(13, 4))\nax1.plot(grpo_rewards, color=\"#22c55e\", linewidth=2, marker=\"o\", markersize=4)\nax1.set_title(\"GRPO Avg Reward per Step\", fontweight=\"bold\")\nax1.set_xlabel(\"Training Step\"); ax1.set_ylabel(\"Avg Group Reward\")\nax1.set_ylim(-0.05, 1.1); ax1.grid(True, alpha=0.3)\nax2.plot(grpo_losses, color=\"#f59e0b\", linewidth=2, marker=\"o\", markersize=4)\nax2.set_title(\"GRPO Policy Loss per Step\", fontweight=\"bold\")\nax2.set_xlabel(\"Training Step\"); ax2.set_ylabel(\"Loss\")\nax2.grid(True, alpha=0.3)\nplt.tight_layout()\nplt.savefig(\"grpo_training_curve.png\", dpi=150, bbox_inches=\"tight\")\nplt.show()\n\nwith open(\"grpo_rewards.json\", \"w\") as f:\n    json.dump({\"rewards\": grpo_rewards, \"losses\": grpo_losses}, f)\nprint(\"Saved grpo_training_curve.png and grpo_rewards.json\")\n",
+   "source": "###############################################################\n# REAL GRPO TRAINING\n# model generates actions -> env rewards -> advantages -> backprop\n###############################################################\n\nimport torch, requests, json, time, gc, re\nimport torch.nn.functional as F\nimport matplotlib\nmatplotlib.use(\"Agg\")\nimport matplotlib.pyplot as plt\nfrom unsloth import FastLanguageModel\n\nENV = \"https://cobedigger-dataengenv.hf.space\"\n\nG           = 4\nTRAIN_STEPS = 25\nLR          = 2e-5\nMAX_NEW     = 200\nMAX_PROMPT  = 400\nMAX_FULL    = MAX_PROMPT + MAX_NEW  # 600\n\nSYSTEM = (\n    \"You are an ML engineer debugging a broken pipeline. \"\n    \"Respond with ONLY a JSON object. No explanation, no markdown.\\n\"\n    \"Actions: inspect_data | run_script | edit_script | query_actor | submit\\n\"\n    \"edit_script payload must have 'old' and 'new' keys.\\n\\n\"\n    \"Stage 1: rename age_years to age AND add df.dropna()\\n\"\n    \"Stage 2: add StandardScaler fitted on X_train AFTER train_test_split\\n\"\n    \"Stage 3: move scaler.fit to AFTER train_test_split\\n\"\n    \"Stage 4: add class_weight='balanced' to LogisticRegression\"\n)\n\nSTAGE_PROMPTS = {\n    1: \"Stage 1 | KeyError: 'age_years'\\nColumn is 'age' not 'age_years'. Also 12 NaN rows. What is your action?\",\n    2: \"Stage 2 | MLPClassifier loss is NaN after 5 iters\\nFeatures not normalised. What is your action?\",\n    3: \"Stage 3 | Accuracy=0.98 — data leakage suspected\\nscaler.fit_transform called before train_test_split. What is your action?\",\n    4: \"Stage 4 | fairness_gap=0.18 > threshold 0.05\\nMinority class poorly predicted. What is your action?\",\n}\n\n# Keywords that indicate the model understands what fix is needed\nSTAGE_KEYWORDS = {\n    1: [\"age_years\", \"age\", \"dropna\", \"rename\"],\n    2: [\"StandardScaler\", \"scaler\", \"fit_transform\", \"normalise\", \"normalize\"],\n    3: [\"train_test_split\", \"after\", \"leakage\", \"fit_transform\"],\n    4: [\"class_weight\", \"balanced\", \"fairness\", \"minority\"],\n}\n\ndef score_action(action_str, stage):\n    \"\"\"\n    Layered reward:\n      0.05 — valid JSON produced at all\n      0.10 — valid JSON with a known action_type\n      0.20 — action_type is contextually sensible for this stage\n      0.40 — edit_script that mentions stage-relevant keywords\n      0.70+ — edit_script that actually moves grader score\n    \"\"\"\n    # Layer 1: valid JSON?\n    try:\n        s = action_str.find(\"{\"); e = action_str.rfind(\"}\")\n        if s == -1 or e == -1:\n            return 0.0\n        action = json.loads(action_str[s:e+1])\n        if \"action_type\" not in action:\n            return 0.0\n        action.setdefault(\"payload\", {})\n        if isinstance(action[\"payload\"], str):\n            action[\"payload\"] = {\"script\": action[\"payload\"]}\n    except Exception:\n        return 0.0\n\n    reward = 0.05  # valid JSON baseline\n    atype  = action[\"action_type\"]\n    VALID  = {\"inspect_data\", \"run_script\", \"edit_script\", \"query_actor\", \"submit\"}\n\n    if atype not in VALID:\n        return reward\n\n    reward = 0.10  # known action_type\n\n    # Layer 2: contextually good action for this stage\n    good_actions = {\n        1: {\"inspect_data\", \"edit_script\", \"run_script\"},\n        2: {\"inspect_data\", \"edit_script\", \"run_script\"},\n        3: {\"run_script\",   \"edit_script\"},\n        4: {\"query_actor\",  \"edit_script\", \"run_script\"},\n    }\n    if atype in good_actions.get(stage, set()):\n        reward = 0.20\n\n    # Layer 3: edit_script with relevant keywords in payload\n    if atype == \"edit_script\":\n        payload_str = json.dumps(action.get(\"payload\", {})).lower()\n        keywords    = STAGE_KEYWORDS.get(stage, [])\n        hits        = sum(1 for kw in keywords if kw.lower() in payload_str)\n        if hits >= 1:\n            reward = 0.35\n        if hits >= 2:\n            reward = 0.50\n\n        # Layer 4: actually execute against env and get grader signal\n        try:\n            requests.post(f\"{ENV}/reset\", timeout=20)\n            r     = requests.post(f\"{ENV}/step\", json=action, timeout=20).json()\n            score = float(r.get(\"reward\", {}).get(\"score\", 0.0))\n            if score > 0:\n                reward = min(1.0, 0.50 + score)\n        except Exception:\n            pass  # keep keyword-based reward\n\n    return reward\n\ndef make_prompt(stage):\n    msgs = [{\"role\": \"system\", \"content\": SYSTEM},\n            {\"role\": \"user\",   \"content\": STAGE_PROMPTS[stage]}]\n    return tokenizer.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True)\n\ndef reinforce_loss(input_ids, attention_mask, advantages):\n    logits     = model(input_ids=input_ids, attention_mask=attention_mask).logits\n    shift_log  = logits[..., :-1, :].contiguous()\n    shift_lab  = input_ids[..., 1:].contiguous()\n    lp         = F.log_softmax(shift_log, dim=-1)\n    tok_lp     = lp.gather(-1, shift_lab.unsqueeze(-1)).squeeze(-1)\n    mask       = attention_mask[..., 1:].float()\n    seq_lp     = (tok_lp * mask).sum(-1) / mask.sum(-1).clamp(min=1)\n    return -(advantages * seq_lp).mean()\n\n# ── Setup ────────────────────────────────────────────────────\nprint(\"Setting up GRPO training...\")\ntorch.cuda.empty_cache(); gc.collect()\nFastLanguageModel.for_training(model)\n\noptimizer = torch.optim.AdamW(\n    [p for p in model.parameters() if p.requires_grad], lr=LR\n)\n\ngrpo_rewards, grpo_losses = [], []\nprint(f\"\\nStarting GRPO: {TRAIN_STEPS} steps, G={G} completions/step\")\nprint(\"=\" * 60)\n\nfor step in range(TRAIN_STEPS):\n    stage  = [1, 2, 3, 4][step % 4]\n    prompt = make_prompt(stage)\n\n    enc = tokenizer(\n        prompt, return_tensors=\"pt\",\n        truncation=True, max_length=MAX_PROMPT\n    ).to(\"cuda\")\n    prompt_len = enc[\"input_ids\"].shape[1]\n\n    # 1. Generate G completions\n    FastLanguageModel.for_inference(model)\n    model.eval()\n    with torch.no_grad():\n        outs = model.generate(\n            **enc,\n            max_new_tokens=MAX_NEW,\n            do_sample=True,\n            temperature=0.9,\n            num_return_sequences=G,\n            pad_token_id=tokenizer.eos_token_id,\n            repetition_penalty=1.1,\n        )\n    completions = [\n        tokenizer.decode(outs[i][prompt_len:], skip_special_tokens=True)\n        for i in range(G)\n    ]\n\n    # 2. Score — print first completion for visibility\n    print(f\"    Sample completion: {completions[0][:120]!r}\")\n    rewards = [score_action(c, stage) for c in completions]\n    time.sleep(0.2)\n\n    rewards_t = torch.tensor(rewards, dtype=torch.float32)\n    std_r     = rewards_t.std().clamp(min=1e-6)\n    mean_r    = rewards_t.mean()\n\n    # If all rewards identical → no gradient signal, skip update\n    if std_r < 1e-4:\n        print(\n            f\"  Step {step+1:02d}/{TRAIN_STEPS} | Stage {stage} | \"\n            f\"Rewards={[f'{r:.2f}' for r in rewards]} | SKIP (no variance)\"\n        )\n        grpo_rewards.append(mean_r.item())\n        grpo_losses.append(0.0)\n        torch.cuda.empty_cache()\n        continue\n\n    advantages = ((rewards_t - mean_r) / std_r).to(\"cuda\")\n\n    # 3. Tokenise prompt + completion\n    full_ids_list, full_mask_list = [], []\n    for c in completions:\n        enc_full = tokenizer(\n            prompt + c, return_tensors=\"pt\",\n            truncation=True, padding=\"max_length\", max_length=MAX_FULL,\n        )\n        full_ids_list.append(enc_full[\"input_ids\"])\n        full_mask_list.append(enc_full[\"attention_mask\"])\n    full_ids  = torch.cat(full_ids_list,  dim=0).to(\"cuda\")\n    full_mask = torch.cat(full_mask_list, dim=0).to(\"cuda\")\n\n    # 4. REINFORCE loss + backprop\n    FastLanguageModel.for_training(model)\n    model.train()\n    loss = reinforce_loss(full_ids, full_mask, advantages)\n    loss.backward()\n    torch.nn.utils.clip_grad_norm_(\n        [p for p in model.parameters() if p.requires_grad], 1.0\n    )\n    optimizer.step()\n    optimizer.zero_grad()\n    torch.cuda.empty_cache()\n\n    avg_r = mean_r.item()\n    grpo_rewards.append(avg_r)\n    grpo_losses.append(loss.item())\n    print(\n        f\"  Step {step+1:02d}/{TRAIN_STEPS} | Stage {stage} | \"\n        f\"Rewards={[f'{r:.2f}' for r in rewards]} | \"\n        f\"Avg={avg_r:.3f} | Loss={loss.item():.4f}\"\n    )\n\nprint(\"\\nGRPO training complete!\")\nprint(f\"   First 5 avg reward : {sum(grpo_rewards[:5])/5:.3f}\")\nprint(f\"   Last  5 avg reward : {sum(grpo_rewards[-5:])/5:.3f}\")\n\n# ── Plot ─────────────────────────────────────────────────────\nfig, (ax1, ax2) = plt.subplots(1, 2, figsize=(13, 4))\nax1.plot(grpo_rewards, color=\"#22c55e\", linewidth=2, marker=\"o\", markersize=4)\nax1.set_title(\"GRPO Avg Reward per Step\", fontweight=\"bold\")\nax1.set_xlabel(\"Training Step\"); ax1.set_ylabel(\"Avg Group Reward\")\nax1.set_ylim(-0.05, 1.1); ax1.grid(True, alpha=0.3)\nax2.plot(grpo_losses, color=\"#f59e0b\", linewidth=2, marker=\"o\", markersize=4)\nax2.set_title(\"GRPO Policy Loss per Step\", fontweight=\"bold\")\nax2.set_xlabel(\"Training Step\"); ax2.set_ylabel(\"Loss\")\nax2.grid(True, alpha=0.3)\nplt.tight_layout()\nplt.savefig(\"grpo_training_curve.png\", dpi=150, bbox_inches=\"tight\")\nplt.show()\n\nwith open(\"grpo_rewards.json\", \"w\") as f:\n    json.dump({\"rewards\": grpo_rewards, \"losses\": grpo_losses}, f)\nprint(\"Saved grpo_training_curve.png and grpo_rewards.json\")\n",
    "metadata": {
     "trusted": true
    },