feat(training): PMOpsGRPOTrainer — override _calculate_rewards to inject rollout rewards directly

Bypasses TRL 1.2.0 kwargs-drop bug: instead of reward_funcs receiving
pre-computed rewards via broken **kwargs, _calculate_rewards reads the
'reward' key from the input batch and returns it as a [batch_size, 1]
tensor. Falls back to standard GRPOTrainer if key is absent.

- training/pm_ops_trainer.py: new PMOpsGRPOTrainer subclass
- training/train_v2.ipynb: swap GRPOTrainer -> PMOpsGRPOTrainer in cell 24

training/pm_ops_trainer.py

@@ -0,0 +1,106 @@
+"""PMOpsGRPOTrainer — GRPOTrainer subclass with direct reward injection.
+
+Problem being solved
+--------------------
+TRL 1.2.0's GRPOTrainer._calculate_rewards() calls reward_funcs with
+**reward_kwargs built from per-example keys in the input batch. This works
+correctly when reward_funcs produce rewards from text completions alone.
+
+It breaks for multi-turn env rollouts: our rollout_func computes rewards
+inside the episode loop (where the env state is available), stores them as
+a 'reward' key in the rollout output dict, and expects reward_func to read
+them back via kwargs. TRL silently drops these keys when building
+reward_kwargs from the batch, so reward_func always receives an empty dict
+and returns 0.0 for every episode.
+
+Fix
+---
+Override _calculate_rewards to check for a pre-computed 'reward' key in the
+input batch. If present, return it directly as a [batch_size, 1] tensor —
+no reward_funcs called, no kwargs needed. If absent, fall back to the
+standard TRL behaviour so the class works as a drop-in replacement.
+
+Usage
+-----
+    from training.pm_ops_trainer import PMOpsGRPOTrainer
+
+    trainer = PMOpsGRPOTrainer(
+        model=model,
+        processing_class=tokenizer,
+        reward_funcs=reward_func,   # kept as-is; only used as fallback
+        train_dataset=dataset,
+        args=grpo_config,
+        rollout_func=rollout_func,  # must put 'reward' key in output dict
+    )
+
+rollout_func contract
+---------------------
+rollout_func(prompts, trainer=None) must return a dict containing at minimum:
+    {
+        "prompt_ids":    list[list[int]],   # one per episode
+        "completion_ids": list[list[int]],
+        "logprobs":      list[list[float]],
+        "reward":        list[float],       # one combined float per episode
+    }
+
+The 'reward' value is the weighted combination of all sub-signals and is
+injected directly as the GRPO advantage signal.
+"""
+
+import torch
+from trl import GRPOTrainer
+
+
+class PMOpsGRPOTrainer(GRPOTrainer):
+    """Drop-in GRPOTrainer replacement with direct reward injection.
+
+    Overrides _calculate_rewards to read pre-computed rewards from the
+    rollout batch instead of calling reward_funcs with broken **kwargs.
+    Falls back to standard GRPOTrainer behaviour if 'reward' key is absent.
+    """
+
+    # Key written by rollout_func and read by _calculate_rewards
+    ROLLOUT_REWARD_KEY = "reward"
+
+    def _calculate_rewards(
+        self,
+        inputs,
+        prompts,
+        completions,
+        completion_ids_list,
+    ):
+        """Override: inject pre-computed rewards when available.
+
+        Args:
+            inputs:              list[dict] — one dict per example in the batch.
+                                 rollout_func output keys land here per example.
+            prompts:             list[str]  — prompt texts (unused in this path)
+            completions:         list[str]  — completion texts (unused)
+            completion_ids_list: list[list[int]] — token IDs (unused)
+
+        Returns:
+            Tensor of shape [batch_size, 1] when pre-computed rewards are found.
+            Falls back to super()._calculate_rewards(...) otherwise.
+        """
+        # Fast-path: pre-computed reward key present in every example
+        if inputs and all(
+            self.ROLLOUT_REWARD_KEY in ex for ex in inputs
+        ):
+            device = self.accelerator.device
+            rewards = torch.tensor(
+                [float(ex[self.ROLLOUT_REWARD_KEY]) for ex in inputs],
+                dtype=torch.float32,
+                device=device,
+            ).unsqueeze(1)  # [batch_size, 1]
+
+            # Log the mean reward so it shows up in training curves
+            mean_r = rewards.mean().item()
+            self.log({"reward/injected_mean": mean_r})
+
+            return rewards
+
+        # Fallback: standard TRL reward_funcs path
+        # Triggered when reward key is absent (e.g. non-rollout evaluation)
+        return super()._calculate_rewards(
+            inputs, prompts, completions, completion_ids_list
+        )

training/train_v2.ipynb

@@ -21,7 +21,9 @@
   {
    "cell_type": "markdown",
    "metadata": {},
-   "source": ["## 0. Install Dependencies"]
+   "source": [
+    "## 0. Install Dependencies"
+   ]
   },
   {
    "cell_type": "code",
@@ -47,7 +49,9 @@
   {
    "cell_type": "markdown",
    "metadata": {},
-   "source": ["## 1. Version Check + GPU Detect"]
+   "source": [
+    "## 1. Version Check + GPU Detect"
+   ]
   },
   {
    "cell_type": "code",
@@ -81,7 +85,9 @@
   {
    "cell_type": "markdown",
    "metadata": {},
-   "source": ["## 2. Clone PM-Ops Repo"]
+   "source": [
+    "## 2. Clone PM-Ops Repo"
+   ]
   },
   {
    "cell_type": "code",
@@ -111,7 +117,9 @@
   {
    "cell_type": "markdown",
    "metadata": {},
-   "source": ["## 3. HuggingFace Login"]
+   "source": [
+    "## 3. HuggingFace Login"
+   ]
   },
   {
    "cell_type": "code",
@@ -161,7 +169,9 @@
   {
    "cell_type": "markdown",
    "metadata": {},
-   "source": ["## 5. Verify Environment"]
+   "source": [
+    "## 5. Verify Environment"
+   ]
   },
   {
    "cell_type": "code",
@@ -239,7 +249,9 @@
   {
    "cell_type": "markdown",
    "metadata": {},
-   "source": ["## 7. Generate Training Dataset"]
+   "source": [
+    "## 7. Generate Training Dataset"
+   ]
   },
   {
    "cell_type": "code",
@@ -375,7 +387,9 @@
   {
    "cell_type": "markdown",
    "metadata": {},
-   "source": ["## 10. Configure GRPO Training"]
+   "source": [
+    "## 10. Configure GRPO Training"
+   ]
   },
   {
    "cell_type": "code",
@@ -423,7 +437,9 @@
   {
    "cell_type": "markdown",
    "metadata": {},
-   "source": ["## 11. Create Trainer"]
+   "source": [
+    "## 11. Create Trainer"
+   ]
   },
   {
    "cell_type": "code",
@@ -431,17 +447,17 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from trl import GRPOTrainer\n",
+    "from training.pm_ops_trainer import PMOpsGRPOTrainer\n",
     "\n",
-    "trainer = GRPOTrainer(\n",
+    "trainer = PMOpsGRPOTrainer(\n",
     "    model=model,\n",
     "    processing_class=tokenizer,\n",
-    "    reward_funcs=reward_func,\n",
+    "    reward_funcs=reward_func,   # fallback only — injected rewards take priority\n",
     "    train_dataset=dataset,\n",
     "    args=grpo_config,\n",
     "    rollout_func=rollout_func,\n",
     ")\n",
-    "print('GRPOTrainer ready')"
+    "print(\"PMOpsGRPOTrainer ready — direct reward injection active\")"
    ]
   },
   {
@@ -485,7 +501,9 @@
   {
    "cell_type": "markdown",
    "metadata": {},
-   "source": ["## 13. Save + Push"]
+   "source": [
+    "## 13. Save + Push"
+   ]
   },
   {
    "cell_type": "code",
@@ -509,7 +527,9 @@
   {
    "cell_type": "markdown",
    "metadata": {},
-   "source": ["## 14. Merge LoRA → bf16 (Optional — for full-weight inference)"]
+   "source": [
+    "## 14. Merge LoRA → bf16 (Optional — for full-weight inference)"
+   ]
   },
   {
    "cell_type": "code",
@@ -533,7 +553,9 @@
   {
    "cell_type": "markdown",
    "metadata": {},
-   "source": ["## 15. Evaluate: Baseline vs Trained"]
+   "source": [
+    "## 15. Evaluate: Baseline vs Trained"
+   ]
   },
   {
    "cell_type": "code",
@@ -621,7 +643,9 @@
   {
    "cell_type": "markdown",
    "metadata": {},
-   "source": ["## 16. Plot"]
+   "source": [
+    "## 16. Plot"
+   ]
   },
   {
    "cell_type": "code",
@@ -660,7 +684,9 @@
   {
    "cell_type": "markdown",
    "metadata": {},
-   "source": ["## 17. Teardown"]
+   "source": [
+    "## 17. Teardown"
+   ]
   },
   {
    "cell_type": "code",
@@ -674,9 +700,16 @@
   }
  ],
  "metadata": {
-  "kernelspec": {"display_name": "Python 3", "language": "python", "name": "python3"},
-  "language_info": {"name": "python", "version": "3.12.0"}
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3.12.0"
+  }
  },
  "nbformat": 4,
  "nbformat_minor": 4
-}
+}