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

@@ -12,10 +12,11 @@ never reach _calculate_rewards via inputs.
 
 Fix
 ---
-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.
+Wrap rollout_func inside __init__ using a closure that captures self directly
+(standard Python closure semantics — no forward reference tricks needed).
+Each time rollout_func is called during training, the wrapper stores the
+'reward' list into self._rollout_reward_cache. _calculate_rewards then injects
+those cached values as a tensor, bypassing the broken kwargs path entirely.
 
 Usage
 -----
@@ -27,7 +28,7 @@ Usage
         reward_funcs=reward_func,   # kept as fallback; not called on inject path
         train_dataset=dataset,
         args=grpo_config,
-        rollout_func=rollout_func,  # must return 'reward' key in output dict
+        rollout_func=rollout_func,  # must return 'reward': list[float] in output
     )
 
 rollout_func contract
@@ -46,40 +47,35 @@ from trl import GRPOTrainer
 
 
 class PMOpsGRPOTrainer(GRPOTrainer):
-    """Drop-in GRPOTrainer replacement with instance-level reward caching.
+    """Drop-in GRPOTrainer with closure-based reward caching.
 
-    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.
+    Wraps rollout_func at construction time to intercept the 'reward' list
+    on every call. _calculate_rewards injects these rewards directly as a
+    tensor — no kwargs, no inputs lookup.
+    Falls back to standard GRPOTrainer behaviour when cache is empty.
     """
 
     def __init__(self, *args, rollout_func=None, **kwargs):
-        # Cache for rewards captured from each rollout_func call
+        # Reward cache populated by the rollout wrapper, consumed by _calculate_rewards
         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 = []
 
+            # self is captured directly by closure — no forward reference needed
             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", []))
+                result = original_rollout(prompts, trainer=trainer)
+                rewards = list(result.get("reward", []))
+                self._rollout_reward_cache = rewards
+                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
 
         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,
         inputs,
@@ -87,20 +83,15 @@ class PMOpsGRPOTrainer(GRPOTrainer):
         completions,
         completion_ids_list,
     ):
-        """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.
-        """
+        """Inject cached rewards when available; fall back to reward_funcs otherwise."""
         cache = self._rollout_reward_cache
         n = len(completions)
 
         if cache:
+            # Handle num_generations > 1: TRL may call with n > len(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
@@ -110,12 +101,13 @@ class PMOpsGRPOTrainer(GRPOTrainer):
                 device=device,
             ).unsqueeze(1)  # [batch_size, 1]
 
-            mean_r = rewards.mean().item()
-            self.log({"reward/injected_mean": mean_r})
+            self.log({"reward/injected_mean": rewards.mean().item()})
             self._rollout_reward_cache = []  # consume cache
+            print(f"[PMOpsGRPOTrainer] injected {n} rewards, mean={rewards.mean().item():.3f}")
             return rewards
 
         # Fallback: standard TRL reward_funcs path
+        print("[PMOpsGRPOTrainer] cache empty — falling back to reward_funcs")
         return super()._calculate_rewards(
             inputs, prompts, completions, completion_ids_list
         )