fix(pm_ops_trainer): use direct closure capture instead of _self_ref

The _self_ref forward-reference trick was the bug — _self_ref was populated after super().__init__() but the cache assignment still used it unnecessarily. In Python, self is captured directly in nested function closures without any forward reference, so _capturing_rollout can assign self._rollout_reward_cache straight away. Added diagnostic prints to confirm capture/inject on each step.

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

training/pm_ops_trainer.py

@@ -43,7 +43,7 @@ rollout_func(prompts, trainer=None) must return a dict containing at minimum:
 """
 
 import torch
-from trl import GRPOTrainer
+from trl.trainer.grpo_trainer import GRPOTrainer
 
 
 class PMOpsGRPOTrainer(GRPOTrainer):
@@ -58,24 +58,45 @@ class PMOpsGRPOTrainer(GRPOTrainer):
     def __init__(self, *args, rollout_func=None, **kwargs):
         # Reward cache populated by the rollout wrapper, consumed by _calculate_rewards
         self._rollout_reward_cache: list[float] = []
+        self._rollout_capture_calls = 0
+        self._warned_rollout_bypass = False
 
+        wrapped_rollout_func = None
         if rollout_func is not None:
             original_rollout = rollout_func
 
-            # self is captured directly by closure — no forward reference needed
-            def _capturing_rollout(prompts, trainer=None):
-                result = original_rollout(prompts, trainer=trainer)
-                rewards = list(result.get("reward", []))
+            # self is captured directly by closure; accept flexible call signatures
+            # because cloud runtimes may invoke rollout_func with positional/keyword variations.
+            def _capturing_rollout(prompts, trainer=None, *rollout_args, **rollout_kwargs):
+                rollout_kwargs.setdefault("trainer", trainer)
+                result = original_rollout(prompts, *rollout_args, **rollout_kwargs)
+
+                raw_rewards = result.get("reward", result.get("rewards", []))
+                if raw_rewards is None:
+                    rewards = []
+                elif isinstance(raw_rewards, torch.Tensor):
+                    rewards = [float(r) for r in raw_rewards.detach().cpu().flatten().tolist()]
+                elif isinstance(raw_rewards, (int, float)):
+                    rewards = [float(raw_rewards)]
+                else:
+                    rewards = [float(r) for r in list(raw_rewards)]
+
                 self._rollout_reward_cache = rewards
+                self._rollout_capture_calls += 1
                 print(f"[PMOpsGRPOTrainer] captured {len(rewards)} rewards "
                       f"(mean={sum(rewards)/len(rewards):.3f})" if rewards else
                       "[PMOpsGRPOTrainer] WARNING: rollout returned no rewards")
                 return result
 
-            rollout_func = _capturing_rollout
+            wrapped_rollout_func = _capturing_rollout
+            rollout_func = wrapped_rollout_func
 
         super().__init__(*args, rollout_func=rollout_func, **kwargs)
 
+        # Keep wrapper bound explicitly in case an upstream patch reassigns rollout_func.
+        if wrapped_rollout_func is not None:
+            self.rollout_func = wrapped_rollout_func
+
     def _calculate_rewards(
         self,
         inputs,
@@ -87,6 +108,18 @@ class PMOpsGRPOTrainer(GRPOTrainer):
         cache = self._rollout_reward_cache
         n = len(completions)
 
+        # Some TRL variants forward rollout extra_fields into `inputs` directly.
+        if not cache:
+            input_rewards: list[float] = []
+            for row in inputs:
+                if isinstance(row, dict) and "reward" in row:
+                    input_rewards.append(float(row["reward"]))
+                else:
+                    input_rewards = []
+                    break
+            if input_rewards:
+                cache = input_rewards
+
         if cache:
             # Handle num_generations > 1: TRL may call with n > len(cache)
             if len(cache) == n:
@@ -106,6 +139,14 @@ class PMOpsGRPOTrainer(GRPOTrainer):
             print(f"[PMOpsGRPOTrainer] injected {n} rewards, mean={rewards.mean().item():.3f}")
             return rewards
 
+        if self.rollout_func is not None and self._rollout_capture_calls == 0 and not self._warned_rollout_bypass:
+            print(
+                "[PMOpsGRPOTrainer] WARNING: rollout_func was never called before reward calculation. "
+                "This cloud runtime is likely bypassing rollout_func (TRL/Unsloth mismatch), so "
+                "reward_funcs only receive prompts/completion_ids/trainer_state and no 'reward' key."
+            )
+            self._warned_rollout_bypass = True
+
         # Fallback: standard TRL reward_funcs path
         print("[PMOpsGRPOTrainer] cache empty — falling back to reward_funcs")
         return super()._calculate_rewards(