fix(pm_ops_trainer): capture rewards via rollout_func wrapper, not inputs

The original approach checked inputs[] for the 'reward' key, but TRL
only puts dataset columns into inputs — rollout_func extra keys never
arrive there. Fix: wrap rollout_func in __init__ to capture the reward
list into self._rollout_reward_cache, then inject it in
_calculate_rewards before the standard reward_funcs path runs.
Also handles num_generations > 1 by repeating rewards to match
completions count.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

training/pm_ops_trainer.py

@@ -2,23 +2,20 @@
 
 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.
+TRL 1.2.0 + Unsloth PatchFastRL strips custom rollout_func output keys
+before calling reward_funcs. The 'reward' key never arrives in kwargs, so
+reward_func always returns 0.0 and GRPO sees constant reward → zero gradient.
 
-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.
+Root cause: TRL builds reward_kwargs from the *dataset* batch columns (only
+'prompt'). rollout_func extra keys are not stored in the dataset, so they
+never reach _calculate_rewards via inputs.
 
 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.
+Wrap rollout_func in __init__ to capture the 'reward' list into an instance
+cache (_rollout_reward_cache) each time rollout_func is called. Then override
+_calculate_rewards to inject those cached rewards directly — no kwargs, no
+inputs lookup needed.
 
 Usage
 -----
@@ -27,24 +24,21 @@ Usage
     trainer = PMOpsGRPOTrainer(
         model=model,
         processing_class=tokenizer,
-        reward_funcs=reward_func,   # kept as-is; only used as fallback
+        reward_funcs=reward_func,   # kept as fallback; not called on inject path
         train_dataset=dataset,
         args=grpo_config,
-        rollout_func=rollout_func,  # must put 'reward' key in output dict
+        rollout_func=rollout_func,  # must return '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
+        "prompt_ids":    list[list[int]],
         "completion_ids": list[list[int]],
         "logprobs":      list[list[float]],
-        "reward":        list[float],       # one combined float per episode
+        "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
@@ -52,15 +46,39 @@ from trl import GRPOTrainer
 
 
 class PMOpsGRPOTrainer(GRPOTrainer):
-    """Drop-in GRPOTrainer replacement with direct reward injection.
+    """Drop-in GRPOTrainer replacement with instance-level reward caching.
 
-    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.
+    Wraps rollout_func to capture the 'reward' list, then injects it in
+    _calculate_rewards before the standard reward_funcs path runs.
+    Falls back to standard GRPOTrainer behaviour if cache is empty.
     """
 
-    # Key written by rollout_func and read by _calculate_rewards
-    ROLLOUT_REWARD_KEY = "reward"
+    def __init__(self, *args, rollout_func=None, **kwargs):
+        # Cache for rewards captured from each rollout_func call
+        self._rollout_reward_cache: list[float] = []
+
+        if rollout_func is not None:
+            original_rollout = rollout_func
+            # Use a list cell to forward-reference self before super().__init__
+            _self_ref: list = []
+
+            def _capturing_rollout(prompts, trainer=None):
+                trainer_obj = trainer or (_self_ref[0] if _self_ref else None)
+                result = original_rollout(prompts, trainer=trainer_obj)
+                # Capture rewards from this rollout batch
+                _self_ref[0]._rollout_reward_cache = list(result.get("reward", []))
+                return result
+
+            rollout_func = _capturing_rollout
+
+        super().__init__(*args, rollout_func=rollout_func, **kwargs)
+
+        # Now self is fully initialised — populate the forward ref
+        if hasattr(self, "_rollout_reward_cache"):
+            try:
+                _self_ref.append(self)  # type: ignore[name-defined]  # noqa: F821
+            except NameError:
+                pass  # rollout_func was None, no closure to fill
 
     def _calculate_rewards(
         self,
@@ -69,38 +87,35 @@ class PMOpsGRPOTrainer(GRPOTrainer):
         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.
+        """Inject pre-computed rewards when cache is populated.
+
+        Falls back to standard TRL path (calls reward_funcs) if cache empty.
+        Handles num_generations > 1 by repeating rewards to match completions.
         """
-        # Fast-path: pre-computed reward key present in every example
-        if inputs and all(
-            self.ROLLOUT_REWARD_KEY in ex for ex in inputs
-        ):
+        cache = self._rollout_reward_cache
+        n = len(completions)
+
+        if cache:
+            if len(cache) == n:
+                rewards_list = cache
+            else:
+                # num_generations > 1: TRL may call reward_func with n > len(cache)
+                # Repeat rewards in round-robin to fill all completions.
+                rewards_list = [cache[i % len(cache)] for i in range(n)]
+
             device = self.accelerator.device
             rewards = torch.tensor(
-                [float(ex[self.ROLLOUT_REWARD_KEY]) for ex in inputs],
+                [float(r) for r in rewards_list],
                 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})
-
+            self._rollout_reward_cache = []  # consume cache
             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
         )