fix(grpo): eliminate zero-advantage collapse from stale rewards

Three root causes all produced reward=const → advantage=0 → zero gradient:

1. generate_rollout_completions uses near-greedy decoding, making every
rollout for the same prompt produce identical outputs. Replaced with
_generate_no_vllm at temperature=1.1 (top_p=0.95, top_k=50).

2. Reward formula env_score*0.85 + json_ratio*0.15 collapses to 0.15
whenever env_score=0 and json_ratio=1.0 (always true after SFT). Added
ok_ratio (fraction of steps env accepted) as the variance source when
the task isn't yet solved; ok_ratio differs across rollouts even with
identical JSON output.

3. All num_generations copies of the same prompt hit the same env seed →
same starting state → same reward. Added gen_offset (tracked per prompt
in grpo_rollout_func) so each generation resets the env with seed+offset,
exploring a distinct org config and task variant.

Also fixed _calculate_rewards tile-mod (cache[i % k] gave every generation
of the same prompt an identical reward even if rollouts differed). Now pads
with mean instead of repeating, and logs reward/injected_std to confirm
variance is non-zero.

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

training/pm_ops_trainer.py

@@ -130,12 +130,13 @@ class PMOpsGRPOTrainer(GRPOTrainer):
                 cache = input_rewards
 
         # Cloud fallback: some patched runtimes skip the normal rollout capture path
-        # before calling _calculate_rewards. Actively invoke rollout_func once here
-        # to populate reward cache from the exact prompt batch.
+        # before calling _calculate_rewards. Actively invoke rollout_func here.
+        # Pass ALL n prompts (including repeated ones for num_generations > 1) so
+        # the rollout_func can generate distinct rewards per generation — NOT tile-mod.
         if not cache and self.rollout_func is not None and prompts:
             try:
-                print("[PMOpsGRPOTrainer] cache empty — probing rollout_func for rewards")
-                out = self.rollout_func(prompts, trainer=self)
+                print(f"[PMOpsGRPOTrainer] cache empty — probing rollout_func for {n} rewards")
+                out = self.rollout_func(list(prompts), trainer=self)
                 cache = self._rollout_reward_cache
                 if not cache and isinstance(out, dict):
                     cache = _coerce_rewards(out.get("reward", out.get("rewards", [])))
@@ -144,11 +145,17 @@ class PMOpsGRPOTrainer(GRPOTrainer):
                 print(f"[PMOpsGRPOTrainer] rollout probe failed: {exc!r}")
 
         if cache:
-            # Handle num_generations > 1: TRL may call with n > len(cache)
             if len(cache) == n:
                 rewards_list = cache
+            elif len(cache) > n:
+                rewards_list = cache[:n]
             else:
-                rewards_list = [cache[i % len(cache)] for i in range(n)]
+                # Still short — extend with mean rather than tile-mod so we don't
+                # duplicate rewards for the same prompt (tile-mod → zero advantage).
+                mean_r = sum(cache) / len(cache)
+                rewards_list = list(cache) + [mean_r] * (n - len(cache))
+                print(f"[PMOpsGRPOTrainer] WARNING: cache has {len(cache)} rewards for n={n}; "
+                      f"padding with mean={mean_r:.3f}. Consider matching num_generations.")
 
             device = self.accelerator.device
             rewards = torch.tensor(
@@ -157,9 +164,12 @@ class PMOpsGRPOTrainer(GRPOTrainer):
                 device=device,
             ).unsqueeze(1)  # [batch_size, 1]
 
-            self.log({"reward/injected_mean": rewards.mean().item()})
+            std = rewards.std().item() if n > 1 else 0.0
+            self.log({"reward/injected_mean": rewards.mean().item(),
+                      "reward/injected_std": std})
             self._rollout_reward_cache = []  # consume cache
-            print(f"[PMOpsGRPOTrainer] injected {n} rewards, mean={rewards.mean().item():.3f}")
+            print(f"[PMOpsGRPOTrainer] injected {n} rewards "
+                  f"mean={rewards.mean().item():.3f} std={std:.3f}")
             return rewards
 
         if self.rollout_func is not None and self._rollout_capture_calls == 0 and not self._warned_rollout_bypass:

training/rollout.py

@@ -42,7 +42,8 @@ def _get_model_for_generation(trainer):
     return trainer.model
 
 
-def _generate_no_vllm(trainer, prompt_text: str, tokenizer, max_new_tokens: int = 512) -> dict:
+def _generate_no_vllm(trainer, prompt_text: str, tokenizer,
+                      max_new_tokens: int = 512, temperature: float = 1.1) -> dict:
     """Generate one completion using HF model.generate() without vLLM.
 
     Returns the same dict shape as generate_rollout_completions so the rest
@@ -68,7 +69,9 @@ def _generate_no_vllm(trainer, prompt_text: str, tokenizer, max_new_tokens: int
             **enc,
             max_new_tokens=max_new_tokens,
             do_sample=True,
-            temperature=0.7,
+            temperature=temperature,
+            top_p=0.95,
+            top_k=50,
             pad_token_id=tokenizer.pad_token_id or tokenizer.eos_token_id,
             output_scores=True,
             return_dict_in_generate=True,
@@ -233,15 +236,18 @@ def rollout_once(
     tokenizer,
     dataset_prompt: str,
     max_steps: int = 15,
+    gen_offset: int = 0,
 ) -> dict:
     """Play one full PM-Ops episode. Returns trajectory + reward signals.
 
-    max_steps: training cap per task type. Triage can be solved in 5 steps;
-    capping at 15 avoids burning compute on aimless late-episode steps.
-    The env's hard limit (MAX_STEPS=40) still applies server-side.
+    gen_offset: added to seed so each GRPO generation explores a different env
+    episode even when receiving the same prompt (same base seed).
     """
     seed = parse_seed_from_prompt(dataset_prompt)
-    result = sync_env.reset(seed=seed) if seed is not None else sync_env.reset()
+    if seed is not None:
+        result = sync_env.reset(seed=seed + gen_offset)
+    else:
+        result = sync_env.reset()
 
     obs = result.observation if hasattr(result, "observation") else result
     obs_dict = _obs_to_dict(obs)
@@ -262,6 +268,7 @@ def rollout_once(
     final_score = 0.0
     step = 0
     done = False
+    ok_count = 0  # number of steps where env accepted the action (last_action_result.ok)
 
     # Channel tracking for reward_no_wrong_channels
     oncall_channels: set[str] = set()   # populated after reading runbook
@@ -324,11 +331,15 @@ def rollout_once(
         new_obs = result.observation if hasattr(result, "observation") else result
         obs_dict = _obs_to_dict(new_obs)
 
+        # Track per-step env acceptance for intermediate reward signal
+        last_result = obs_dict.get("last_action_result") or {}
+        if last_result.get("ok"):
+            ok_count += 1
+
         # Extract oncall channels from runbook response (available one step later)
         if action_type == "meta.read_runbook":
-            last = obs_dict.get("last_action_result") or {}
-            if last.get("ok"):
-                data = last.get("data") or {}
+            if last_result.get("ok"):
+                data = last_result.get("data") or {}
                 if isinstance(data, dict):
                     org = data.get("org_config") or {}
                     oncall_channels = set(org.get("oncall_channels", {}).values())
@@ -356,30 +367,37 @@ def rollout_once(
 
     read_runbook_reward = 1.0 if read_runbook_done else 0.0
 
+    # ok_ratio: fraction of steps where env accepted the action — provides per-step
+    # intermediate signal that varies across rollouts even when final_score=0.
+    ok_ratio = ok_count / max(step, 1)
+
     # --- Reward gating ---
     # If model never output valid JSON, it never actually tried anything.
     # Strip all process rewards and apply a harsh penalty.
-    # Process rewards (runbook, no_wrong, efficiency) only matter if model acted.
     if valid_action_count == 0:
         combined = -1.0
     elif final_score == 0.0:
-        # Model tried (valid JSON) but task failed — small process credit, no final bonus
+        # Model tried (valid JSON) but task failed.
+        # ok_ratio gives gradient signal that distinguishes "wrong but env-valid" actions
+        # from "env-rejected" ones — this is the variance source GRPO needs when
+        # final_score is always 0 (model not yet solving tasks).
         combined = (
-            valid_json_ratio    * 0.15
+            ok_ratio              * 0.20
             + read_runbook_reward * 0.10
             - 0.30                        # hard penalty for zero task completion
         )
     else:
         combined = (
-            final_score           * 0.45
+            final_score           * 0.40
+            + ok_ratio            * 0.10
             + no_wrong_channels   * 0.15
-            + valid_json_ratio    * 0.15
+            + valid_json_ratio    * 0.10
             + read_runbook_reward * 0.15
             + efficiency          * 0.10
         )
 
     print(
-        f"[rollout] steps={step} final={final_score:.3f} "
+        f"[rollout] steps={step} final={final_score:.3f} ok={ok_ratio:.2f} "
         f"json={valid_json_ratio:.2f} runbook={read_runbook_reward:.0f} "
         f"no_wrong={no_wrong_channels:.2f} eff={efficiency:.2f} "
         f"valid_acts={valid_action_count} → combined={combined:.3f}"
@@ -420,13 +438,20 @@ def make_rollout_func(sync_env, tokenizer, max_steps: int = 15):
             "read_runbook_reward": [],
             "efficiency_reward": [],
         }
+        # Track how many times each unique prompt has appeared so we can pass
+        # a gen_offset — ensures repeated prompts (num_generations > 1) hit
+        # different env seeds and produce different rollouts.
+        prompt_seen: dict[str, int] = {}
         for prompt_text in prompts:
+            gen_offset = prompt_seen.get(prompt_text, 0)
+            prompt_seen[prompt_text] = gen_offset + 1
             episode = rollout_once(
                 trainer=trainer,
                 sync_env=sync_env,
                 tokenizer=tokenizer,
                 dataset_prompt=prompt_text,
                 max_steps=max_steps,
+                gen_offset=gen_offset,
             )
             for k in out:
                 out[k].append(episode[k])

training/train_v3.ipynb

@@ -512,106 +512,7 @@
    "id": "cell-14",
    "metadata": {},
    "outputs": [],
-   "source": [
-    "from trl.experimental.openenv import generate_rollout_completions\n",
-    "from training.rollout import (\n",
-    "    _obs_to_dict, _current_obs_text, build_messages,\n",
-    "    extract_json_action, step_aware_fallback,\n",
-    ")\n",
-    "from training.dataset import parse_seed_from_prompt\n",
-    "\n",
-    "grpo_env = GenericEnvClient(base_url=ENV_URL).sync()\n",
-    "grpo_env.connect()\n",
-    "print('GRPO training env connected')\n",
-    "\n",
-    "\n",
-    "def run_grpo_episode(trainer, env, tok, dataset_prompt, max_steps=TRAIN_MAX_STEPS):\n",
-    "    \"\"\"Run one full PM-ops episode and return flat trajectory + reward.\"\"\"\n",
-    "    seed     = parse_seed_from_prompt(dataset_prompt)\n",
-    "    result   = env.reset(seed=seed) if seed is not None else env.reset()\n",
-    "    obs_dict = _obs_to_dict(result.observation if hasattr(result, 'observation') else result)\n",
-    "    task_brief = obs_dict.get('task_brief') or dataset_prompt\n",
-    "\n",
-    "    prompt_ids, completion_ids, logprobs = [], [], []\n",
-    "    turn_history      = []\n",
-    "    valid_json_count  = 0\n",
-    "    env_score         = 0.0\n",
-    "    step, done        = 0, False\n",
-    "    _sample_logged    = False\n",
-    "\n",
-    "    while not done and step < max_steps:\n",
-    "        obs_text    = _current_obs_text(obs_dict, step, task_brief)\n",
-    "        msgs        = build_messages(turn_history, obs_text)\n",
-    "        prompt_text = tok.apply_chat_template(\n",
-    "            msgs, add_generation_prompt=True, tokenize=False, enable_thinking=False\n",
-    "        )\n",
-    "\n",
-    "        rollout_out = generate_rollout_completions(trainer, [prompt_text])[0]\n",
-    "        prompt_ids.extend(rollout_out['prompt_ids'])\n",
-    "        completion_ids.extend(rollout_out['completion_ids'])\n",
-    "        logprobs.extend(rollout_out['logprobs'])\n",
-    "\n",
-    "        completion_text = rollout_out.get('text') or tok.decode(\n",
-    "            rollout_out['completion_ids'], skip_special_tokens=True\n",
-    "        )\n",
-    "\n",
-    "        # Log one sample per episode so we can visually track format quality\n",
-    "        if not _sample_logged:\n",
-    "            print(f'  [sample] {repr(completion_text[:180])}')\n",
-    "            _sample_logged = True\n",
-    "\n",
-    "        parsed = extract_json_action(completion_text)\n",
-    "        if parsed is not None:\n",
-    "            valid_json_count += 1\n",
-    "        else:\n",
-    "            parsed = step_aware_fallback(step, max_steps)\n",
-    "\n",
-    "        action_type = parsed.get('action_type', 'meta.noop')\n",
-    "        args        = parsed.get('args', {})\n",
-    "\n",
-    "        turn_history.append({\n",
-    "            'obs_text'  : obs_text,\n",
-    "            'completion': completion_text,\n",
-    "            'is_runbook': (action_type == 'meta.read_runbook' and parsed is not None),\n",
-    "        })\n",
-    "\n",
-    "        result   = env.step({'action_type': action_type, 'args': args})\n",
-    "        obs_dict = _obs_to_dict(result.observation if hasattr(result, 'observation') else result)\n",
-    "        done      = bool(getattr(result, 'done', obs_dict.get('done', False)))\n",
-    "        env_score = float(getattr(result, 'reward', obs_dict.get('reward', 0.0)))\n",
-    "        step     += 1\n",
-    "\n",
-    "    json_ratio = valid_json_count / max(step, 1)\n",
-    "    reward     = env_score * 0.85 + json_ratio * 0.15\n",
-    "    print(f'  [rollout] steps={step} env={env_score:.3f} json={json_ratio:.2f} -> reward={reward:.3f}')\n",
-    "    return {\n",
-    "        'prompt_ids'    : prompt_ids,\n",
-    "        'completion_ids': completion_ids,\n",
-    "        'logprobs'      : logprobs,\n",
-    "        'reward'        : reward,\n",
-    "    }\n",
-    "\n",
-    "\n",
-    "def grpo_rollout_func(prompts, trainer=None):\n",
-    "    out = {'prompt_ids': [], 'completion_ids': [], 'logprobs': [], 'reward': []}\n",
-    "    for prompt in prompts:\n",
-    "        ep = run_grpo_episode(trainer, grpo_env, tokenizer, prompt, TRAIN_MAX_STEPS)\n",
-    "        for k in out:\n",
-    "            out[k].append(ep[k])\n",
-    "    return out\n",
-    "\n",
-    "\n",
-    "def grpo_reward_func(completions, **kwargs):\n",
-    "    \"\"\"Passthrough — reward is pre-computed in grpo_rollout_func.\"\"\"\n",
-    "    rewards = kwargs.get('reward', [])\n",
-    "    if not rewards:\n",
-    "        print(f'[ERROR] reward_func: no reward key. kwargs keys: {list(kwargs.keys())}')\n",
-    "        return [0.0] * len(completions)\n",
-    "    return [float(r) for r in rewards]\n",
-    "\n",
-    "\n",
-    "print(f'GRPO rollout ready  max_steps={TRAIN_MAX_STEPS}')"
-   ]
+   "source": "from training.rollout import (\n    _obs_to_dict, _current_obs_text, build_messages,\n    extract_json_action, step_aware_fallback,\n    _generate_no_vllm,\n)\nfrom training.dataset import parse_seed_from_prompt\n\ngrpo_env = GenericEnvClient(base_url=ENV_URL).sync()\ngrpo_env.connect()\nprint('GRPO training env connected')\n\n# Must be > 1.0 so rollouts of the same prompt diverge and GRPO gets non-zero advantage.\nROLLOUT_TEMPERATURE = 1.1\n\n\ndef run_grpo_episode(trainer, env, tok, dataset_prompt, max_steps=TRAIN_MAX_STEPS,\n                     gen_offset=0):\n    \"\"\"Run one PM-ops episode. gen_offset shifts the env seed so each GRPO\n    generation for the same prompt explores a different env episode.\n    \"\"\"\n    seed = parse_seed_from_prompt(dataset_prompt)\n    if seed is not None:\n        result = env.reset(seed=seed + gen_offset)\n    else:\n        result = env.reset()\n\n    obs_dict   = _obs_to_dict(result.observation if hasattr(result, 'observation') else result)\n    task_brief = obs_dict.get('task_brief') or dataset_prompt\n\n    prompt_ids, completion_ids, logprobs = [], [], []\n    turn_history      = []\n    valid_json_count  = 0\n    ok_count          = 0   # steps where env accepted the action (last_action_result.ok)\n    read_runbook_done = False\n    env_score         = 0.0\n    step, done        = 0, False\n    _sample_logged    = False\n\n    while not done and step < max_steps:\n        obs_text    = _current_obs_text(obs_dict, step, task_brief)\n        msgs        = build_messages(turn_history, obs_text)\n        prompt_text = tok.apply_chat_template(\n            msgs, add_generation_prompt=True, tokenize=False, enable_thinking=False\n        )\n\n        # Always use _generate_no_vllm at ROLLOUT_TEMPERATURE.\n        # generate_rollout_completions defaults to greedy/low-temp, which collapses\n        # all rollouts for the same prompt to identical outputs → zero GRPO advantage.\n        rollout_out = _generate_no_vllm(\n            trainer, prompt_text, tok,\n            max_new_tokens=MAX_COMP_LEN,\n            temperature=ROLLOUT_TEMPERATURE,\n        )\n        prompt_ids.extend(rollout_out['prompt_ids'])\n        completion_ids.extend(rollout_out['completion_ids'])\n        logprobs.extend(rollout_out['logprobs'])\n        completion_text = rollout_out['text']\n\n        if not _sample_logged:\n            print(f'  [sample] {repr(completion_text[:180])}')\n            _sample_logged = True\n\n        parsed = extract_json_action(completion_text)\n        if parsed is not None:\n            valid_json_count += 1\n        else:\n            parsed = step_aware_fallback(step, max_steps)\n\n        action_type = parsed.get('action_type', 'meta.noop')\n        args        = parsed.get('args', {})\n\n        if action_type == 'meta.read_runbook':\n            read_runbook_done = True\n\n        turn_history.append({\n            'obs_text'  : obs_text,\n            'completion': completion_text,\n            'is_runbook': (action_type == 'meta.read_runbook' and parsed is not None),\n        })\n\n        try:\n            result = env.step({'action_type': action_type, 'args': args})\n        except RuntimeError as exc:\n            if 'VALIDATION_ERROR' in str(exc):\n                result = env.step({'action_type': 'meta.noop', 'args': {}})\n            else:\n                raise\n\n        obs_dict  = _obs_to_dict(result.observation if hasattr(result, 'observation') else result)\n\n        # Per-step intermediate signal: did env accept this action?\n        last_res = obs_dict.get('last_action_result') or {}\n        if last_res.get('ok'):\n            ok_count += 1\n\n        done      = bool(getattr(result, 'done', obs_dict.get('done', False)))\n        env_score = float(getattr(result, 'reward', obs_dict.get('reward', 0.0)))\n        step     += 1\n\n    json_ratio = valid_json_count / max(step, 1)\n    ok_ratio   = ok_count / max(step, 1)\n    rb_reward  = 1.0 if read_runbook_done else 0.0\n\n    # Reward formula designed to produce VARIANCE across rollouts:\n    # - ok_ratio varies per episode (different actions accepted/rejected by env)\n    # - env_score varies when the model solves the task with different quality\n    # Using json_ratio alone (always ~1.0 after SFT) gives a constant reward → zero gradient.\n    if valid_json_count == 0:\n        reward = -1.0\n    elif env_score == 0.0:\n        # Task not completed — ok_ratio differentiates rollouts that took env-valid\n        # actions from ones that didn't, giving GRPO a gradient even before task completion.\n        reward = ok_ratio * 0.20 + rb_reward * 0.10 - 0.30\n    else:\n        reward = env_score * 0.60 + ok_ratio * 0.10 + json_ratio * 0.10 + rb_reward * 0.20\n\n    reward = max(-1.0, min(1.0, reward))\n    print(f'  [rollout] steps={step} env={env_score:.3f} ok={ok_ratio:.2f} '\n          f'json={json_ratio:.2f} rb={rb_reward:.0f} offset={gen_offset} -> reward={reward:.3f}')\n    return {\n        'prompt_ids'    : prompt_ids,\n        'completion_ids': completion_ids,\n        'logprobs'      : logprobs,\n        'reward'        : reward,\n    }\n\n\ndef grpo_rollout_func(prompts, trainer=None):\n    out = {'prompt_ids': [], 'completion_ids': [], 'logprobs': [], 'reward': []}\n    # Count how many times each prompt has appeared so far in this batch.\n    # TRL sends [p, p, ...] (same prompt num_gen times) → gen_offset breaks symmetry.\n    prompt_seen: dict = {}\n    for prompt in prompts:\n        gen_offset          = prompt_seen.get(prompt, 0)\n        prompt_seen[prompt] = gen_offset + 1\n        ep = run_grpo_episode(trainer, grpo_env, tokenizer, prompt,\n                              TRAIN_MAX_STEPS, gen_offset=gen_offset)\n        for k in out:\n            out[k].append(ep[k])\n    return out\n\n\n_warned_missing_reward = False\ndef grpo_reward_func(completions, **kwargs):\n    \"\"\"Passthrough — reward is pre-computed in grpo_rollout_func.\"\"\"\n    global _warned_missing_reward\n    rewards = kwargs.get('reward', [])\n    if not rewards:\n        if not _warned_missing_reward:\n            print(f\"[WARN] reward_func fallback with no reward key. kwargs keys: {list(kwargs.keys())}\")\n            _warned_missing_reward = True\n        return [0.0] * len(completions)\n    return [float(r) for r in rewards]\n\n\nprint(f'GRPO rollout ready  max_steps={TRAIN_MAX_STEPS}  temperature={ROLLOUT_TEMPERATURE}')"
   },
   {
    "cell_type": "markdown",
@@ -917,4 +818,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 5
-}
+}