fix(grpo): fresh reward design — runbook-compliance scoring

Root cause analysis
-------------------
env_score was always 0.0 for two independent reasons:

1. create_ticket(label='sec-issue', priority='P1') fails TicketingApp
validation for most org configs. The app returns ok=False, no ticket is
stored, and the grader returns 0.0 immediately (no new tickets found).

2. Dataset seeds produced triage briefs but PMOpsEnvironment.reset(seed)
independently ran rng.choice(_TASK_TYPES) — silently generating
release_notes or dep_update tasks for the same seed. The triage grader
was never called; the wrong grader always returned 0.0.

Fixes
-----
dataset.py: generate_triage_dataset now pre-simulates the env's first two
RNG calls (difficulty + task_type) and only includes seeds where the env
WILL run a triage episode, using the exact same difficulty and org params.

rewards.py: replaced all previous reward formulas with compute_rollout_reward,
a single function based on RUNBOOK COMPLIANCE:
read_runbook +0.10 — process: read runbook first
valid_label +0.20/-0.10 — label from label_taxonomy?
valid_priority +0.15/-0.10 — priority from priority_levels?
valid_team +0.20/-0.10 — team from team_map?
right_channel +0.25/-0.10 — posted to oncall channel?
env_bonus +0.10 * env_score — bonus when grader agrees

This creates genuine reward VARIANCE: the same hardcoded model template
(sec-issue, P1, infra, #infra-team) scores differently for different org
configs — correct for some seeds, wrong for others — giving GRPO the
advantage signal it needs.

rollout.py: updated rollout_once to track ticket_label, ticket_priority,
assigned_team, valid_labels/priorities/teams/channels from the runbook
response, and call compute_rollout_reward. Log now shows ✓/✗ per component.

train_v3.ipynb: updated run_grpo_episode and dataset cell to match.

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

training/dataset.py

@@ -1,7 +1,10 @@
 """Generate fixed-seed training dataset for PM-Ops triage task.
 
-Each row contains a seed embedded in the prompt string so the rollout
-function can pass it to env.reset(seed=...) for reproducible episodes.
+Each seed is validated against the env's own RNG so the task type, difficulty,
+org_config, and scenario the env generates at reset(seed=S) EXACTLY matches the
+brief embedded in the prompt. Previously, the dataset generated triage briefs but
+the env silently ran a different task type (release_notes, dep_update, etc.) for
+the same seed — causing env_score=0 for every episode.
 """
 import json
 import os
@@ -15,6 +18,22 @@ from server.world.scenario_gen import generate_scenario
 
 SEED_PREFIX = "SEED:"
 
+# Must match pm_ops_environment.py constants exactly
+_ENV_DIFFICULTY_POOL = ["easy", "medium", "medium", "hard"]
+_ENV_TASK_TYPES = ["triage", "incident_routing", "release_notes", "dep_update"]
+
+
+def _env_params(seed: int) -> tuple[str, str]:
+    """Predict difficulty + task_type the env will choose for this seed.
+
+    Replicates the first two RNG calls in PMOpsEnvironment.reset() so we can
+    filter seeds to those that produce the task type we want.
+    """
+    rng = random.Random(seed)
+    difficulty = rng.choice(_ENV_DIFFICULTY_POOL)
+    task_type  = rng.choice(_ENV_TASK_TYPES)
+    return difficulty, task_type
+
 
 def _valid_triage(scenario: dict) -> bool:
     exp = scenario.get("expected", {})
@@ -22,14 +41,28 @@ def _valid_triage(scenario: dict) -> bool:
 
 
 def generate_triage_dataset(n_episodes: int = 150, base_seed: int = 42) -> list[dict]:
-    """Return list of dataset rows with embedded seeds."""
+    """Return dataset rows where the env WILL run a triage episode for the embedded seed.
+
+    We pre-simulate the env's RNG to only include seeds where:
+      1. env.reset(seed) picks task_type="triage"
+      2. The resulting org + scenario are valid (have expected channel + team)
+      3. The difficulty, org_config, and brief exactly match what the env will use
+
+    This eliminates the mismatch where the dataset had triage briefs but the env
+    graded as release_notes → guaranteed env_score=0.
+    """
     rng = random.Random(base_seed)
     rows = []
 
     while len(rows) < n_episodes:
         seed = rng.randint(0, 2**31)
-        difficulty = rng.choice(["easy", "medium", "medium", "hard"])
 
+        # Only use seeds the env will run as triage
+        difficulty, task_type = _env_params(seed)
+        if task_type != "triage":
+            continue
+
+        # Generate org + scenario using the SAME difficulty + seed the env will use
         org, scenario = None, None
         for attempt in range(10):
             org = generate_org_config(seed + attempt, difficulty)
@@ -41,7 +74,6 @@ def generate_triage_dataset(n_episodes: int = 150, base_seed: int = 42) -> list[
             continue
 
         rows.append({
-            # Seed embedded so rollout can parse it and pass to reset()
             "prompt": f"{SEED_PREFIX}{seed} | {scenario['brief']}",
             "seed": seed,
             "difficulty": difficulty,

training/rewards.py

@@ -1,149 +1,100 @@
-"""Reward functions for PM-Ops GRPO training.
-
-Weights are applied inside each function so GRPOTrainer sums them to a
-total reward. All weights sum to 1.0:
-
-  final_score        0.45  — correctness: label + priority + team + channel
-  no_wrong_channels  0.15  — anti-hack: penalise channel-spray behaviour
-  valid_json         0.15  — format discipline: fraction of valid JSON outputs
-  read_runbook       0.15  — process: did agent read runbook before acting?
-  efficiency         0.10  — speed: steps saved when task completed correctly
-
-Reward hacking protection (point 8 of hackathon guide):
-  The main PM-Ops hack is posting to every channel to guarantee hitting the
-  right one. reward_no_wrong_channels makes this explicit and gradient-visible
-  rather than buried inside reward_final_score.
-
-To change weights, edit the WEIGHT_* constants below.
-Equal-weight Option B: set all five to 0.20.
-
-Architecture note — two reward modes:
-  1. Split mode (ALL_REWARD_FUNCS): five separate functions, each reading one
-     kwargs key from the rollout dict. TRL sums them. Requires TRL to pass
-     rollout extra keys as kwargs to reward_funcs (works in TRL 0.17–0.24 +
-     Unsloth patch).
-  2. Combined mode (COMBINED_REWARD_FUNC): single function that reads all five
-     kwargs keys and returns the weighted sum. Use this if split mode shows
-     all-zero rewards (indicates TRL is not passing kwargs, e.g. version
-     mismatch). Switch by replacing reward_funcs=ALL_REWARD_FUNCS with
-     reward_funcs=COMBINED_REWARD_FUNC in GRPOTrainer.
+"""Runbook-compliance reward for PM-Ops GRPO training.
+
+Design philosophy (v4)
+----------------------
+Previous reward formulas (json_ratio, ok_ratio, diversity_bonus) were all constant
+once the model learned to output valid JSON — giving zero GRPO advantage.
+
+Root problem: env_score was always 0 because:
+  1. The model used hardcoded label='sec-issue' / priority='P1', which fail
+     TicketingApp validation for most org configs → no ticket created → grader
+     returns 0.0 immediately.
+  2. Dataset seeds produced triage briefs but the env ran a DIFFERENT task type
+     (release_notes, dep_update) for the same seed → guaranteed mismatch.
+
+Fix: reward is computed from the agent's ACTIONS compared to RUNBOOK DATA.
+The org_config (returned by meta.read_runbook) varies by seed — different orgs
+have different valid labels, priorities, teams, and oncall channels. The model's
+fixed template (sec-issue, P1, infra) scores well for some orgs and badly for
+others, creating the reward VARIANCE that GRPO needs.
+
+Components (sum = 1.0 when all correct):
+  read_runbook   0.10  — process: did agent read runbook first?
+  valid_label    0.20  — used a label from label_taxonomy? (+0.20 / -0.10)
+  valid_priority 0.15  — used a priority from priority_levels? (+0.15 / -0.10)
+  valid_team     0.20  — assigned to a team from team_map? (+0.20 / -0.10)
+  right_channel  0.25  — posted to an oncall channel? (+0.25 / -0.10 per wrong)
+  env_bonus      0.10  — env grader bonus when everything lines up correctly
+
+This is computed IN the rollout (not by TRL reward_funcs) because it needs
+access to the per-episode runbook data.
 """
 
-import warnings
-
-WEIGHT_FINAL_SCORE = 0.45
-WEIGHT_NO_WRONG_CHANNELS = 0.15
-WEIGHT_VALID_JSON = 0.15
-WEIGHT_READ_RUNBOOK = 0.15
-WEIGHT_EFFICIENCY = 0.10
-
-assert abs(
-    WEIGHT_FINAL_SCORE + WEIGHT_NO_WRONG_CHANNELS + WEIGHT_VALID_JSON
-    + WEIGHT_READ_RUNBOOK + WEIGHT_EFFICIENCY - 1.0
-) < 1e-9, "Reward weights must sum to 1.0"
-
-_EXPECTED_KEYS = frozenset([
-    "final_score_reward", "no_wrong_channels_reward",
-    "valid_json_reward", "read_runbook_reward", "efficiency_reward",
-])
-_warned_empty = False
-
-
-def _extract(kwargs: dict, key: str, n: int) -> list[float]:
-    global _warned_empty
-    rewards = kwargs.get(key, [])
-    if not rewards:
-        if not _warned_empty:
-            present = set(kwargs.keys()) & _EXPECTED_KEYS
-            missing = _EXPECTED_KEYS - set(kwargs.keys())
-            warnings.warn(
-                f"[rewards] kwargs missing rollout reward keys — all rewards will be 0!\n"
-                f"  present: {present or 'none'}\n"
-                f"  missing: {missing}\n"
-                f"  all kwargs keys: {list(kwargs.keys())}\n"
-                f"  CAUSE: TRL version mismatch or rollout_func not returning these keys.\n"
-                f"  FIX: switch to reward_funcs=COMBINED_REWARD_FUNC (see rewards.py).",
-                stacklevel=3,
-            )
-            _warned_empty = True
-        return [0.0] * n
-    return [float(r) for r in rewards]
-
-
-def reward_final_score(completions, **kwargs) -> list[float]:
-    """Primary correctness signal from the env grader (0–1). Weight: 0.45"""
-    raw = _extract(kwargs, "final_score_reward", len(completions))
-    return [r * WEIGHT_FINAL_SCORE for r in raw]
-
-
-def reward_no_wrong_channels(completions, **kwargs) -> list[float]:
-    """Anti-hack: penalise posting to non-oncall channels. Weight: 0.15
-
-    Graduated penalty: 1.0 for zero wrong posts, -0.5 per wrong post.
-    Computed in rollout from oncall_channels extracted from the runbook.
-    Explicitly discourages the channel-spray failure mode.
-    """
-    raw = _extract(kwargs, "no_wrong_channels_reward", len(completions))
-    return [r * WEIGHT_NO_WRONG_CHANNELS for r in raw]
-
-
-def reward_valid_json(completions, **kwargs) -> list[float]:
-    """Fraction of steps with parseable JSON output (0–1). Weight: 0.15
-
-    Anti-staleness: penalises the 'output garbage every step' failure mode.
-    """
-    raw = _extract(kwargs, "valid_json_reward", len(completions))
-    return [r * WEIGHT_VALID_JSON for r in raw]
-
-
-def reward_read_runbook(completions, **kwargs) -> list[float]:
-    """Binary: did agent call meta.read_runbook before acting? Weight: 0.15
-
-    Teaches the agent to learn org conventions first instead of guessing.
-    """
-    raw = _extract(kwargs, "read_runbook_reward", len(completions))
-    return [r * WEIGHT_READ_RUNBOOK for r in raw]
-
-
-def reward_efficiency(completions, **kwargs) -> list[float]:
-    """Steps saved when task completed correctly (0–1). Weight: 0.10
-
-    Only non-zero when final_score >= 0.3 — does not reward fast failure.
-    """
-    raw = _extract(kwargs, "efficiency_reward", len(completions))
-    return [r * WEIGHT_EFFICIENCY for r in raw]
-
-
-def combined_reward(completions, **kwargs) -> list[float]:
-    """Single combined reward — fallback when TRL does not pass split kwargs.
 
-    Reads all five rollout reward keys and returns their weighted sum.
-    Use with: reward_funcs=[combined_reward]
+def compute_rollout_reward(
+    *,
+    read_runbook_done: bool,
+    valid_labels: set,        # from org_config.label_taxonomy.values()
+    valid_priorities: set,    # from org_config.priority_levels
+    valid_teams: set,         # from org_config.team_map.values()
+    oncall_channels: set,     # from org_config.oncall_channels.values()
+    ticket_label: str | None,       # label used in create_ticket (None if no ticket)
+    ticket_priority: str | None,    # priority used in create_ticket
+    assigned_team: str | None,      # team from assign_ticket (None if not called)
+    posted_channels: list[str],     # all channels from chat.post_message
+    env_score: float,               # final env grader score (0–1)
+    valid_json_count: int,          # steps with parseable JSON output
+) -> float:
+    """Compute a single reward scalar for one episode.
+
+    Returns a value in [-1.0, 1.0].
     """
-    n = len(completions)
-    final    = _extract(kwargs, "final_score_reward",       n)
-    no_wrong = _extract(kwargs, "no_wrong_channels_reward", n)
-    vj       = _extract(kwargs, "valid_json_reward",        n)
-    rb       = _extract(kwargs, "read_runbook_reward",      n)
-    eff      = _extract(kwargs, "efficiency_reward",        n)
-    return [
-        f * WEIGHT_FINAL_SCORE
-        + nw * WEIGHT_NO_WRONG_CHANNELS
-        + v  * WEIGHT_VALID_JSON
-        + r  * WEIGHT_READ_RUNBOOK
-        + e  * WEIGHT_EFFICIENCY
-        for f, nw, v, r, e in zip(final, no_wrong, vj, rb, eff)
-    ]
-
-
-# Split mode: TRL sums five separate signals (preferred — gives per-signal trackio curves)
-ALL_REWARD_FUNCS = [
-    reward_final_score,
-    reward_no_wrong_channels,
-    reward_valid_json,
-    reward_read_runbook,
-    reward_efficiency,
-]
-
-# Combined mode: single function — use when kwargs aren't flowing in split mode
-COMBINED_REWARD_FUNC = [combined_reward]
+    if valid_json_count == 0:
+        return -1.0
+
+    reward = 0.0
+
+    # 1. Runbook read (+0.10 process bonus)
+    if read_runbook_done:
+        reward += 0.10
+
+    # 2. Ticket label valid (only scored if a ticket was created)
+    if ticket_label is not None:
+        if valid_labels:
+            if ticket_label in valid_labels:
+                reward += 0.20
+            else:
+                reward -= 0.10  # wrong label — validation would have rejected it
+
+    # 3. Ticket priority valid
+    if ticket_priority is not None:
+        if valid_priorities:
+            if ticket_priority in valid_priorities:
+                reward += 0.15
+            else:
+                reward -= 0.10
+
+    # 4. Ticket assigned to a valid team
+    if assigned_team is not None:
+        if valid_teams:
+            if assigned_team in valid_teams:
+                reward += 0.20
+            else:
+                reward -= 0.10
+
+    # 5. Posted to the right oncall channel
+    if posted_channels:
+        if oncall_channels:
+            correct_posts = [ch for ch in posted_channels if ch in oncall_channels]
+            wrong_posts   = [ch for ch in posted_channels if ch not in oncall_channels]
+            if correct_posts:
+                reward += 0.25
+            reward -= 0.10 * len(wrong_posts)  # -0.10 per channel-spray post
+        else:
+            # Posted without reading runbook — can't verify, mild penalty
+            reward -= 0.05 * len(posted_channels)
+
+    # 6. Env grader bonus — full env score adds on top when everything is correct
+    reward += env_score * 0.10
+
+    return max(-1.0, min(1.0, reward))

training/rollout.py

@@ -262,17 +262,22 @@ def rollout_once(
     # Each entry: {"obs_text": str, "completion": str, "is_runbook": bool}
     turn_history: list[dict] = []
 
-    # Reward accumulators
+    # Rollout accumulators
     valid_action_count = 0
-    read_runbook_done = False
     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
-    posted_channels: list[str] = []     # every channel the agent posts to
+    # Runbook-compliance reward tracking
+    read_runbook_done  = False
+    valid_labels:     set[str] = set()   # org label_taxonomy values
+    valid_priorities: set[str] = set()   # org priority_levels
+    valid_teams:      set[str] = set()   # org team_map values
+    oncall_channels:  set[str] = set()   # org oncall_channels values
+    ticket_label:    str | None = None   # label used in create_ticket
+    ticket_priority: str | None = None   # priority used in create_ticket
+    assigned_team:   str | None = None   # team from assign_ticket
+    posted_channels: list[str] = []      # every channel posted to
 
     while not done and step < max_steps:
         # obs_text for THIS step — stored in history BEFORE stepping
@@ -309,10 +314,16 @@ def rollout_once(
         action_type: str = parsed.get("action_type", "meta.noop")
         args: dict = parsed.get("args", {})
 
-        if action_type == "meta.read_runbook" and is_valid_json and not read_runbook_done:
+        if action_type == "meta.read_runbook" and is_valid_json:
             read_runbook_done = True
 
-        # Track chat.post_message targets for anti-hack reward
+        if action_type == "ticketing.create_ticket" and is_valid_json and ticket_label is None:
+            ticket_label    = args.get("label")
+            ticket_priority = args.get("priority")
+
+        if action_type == "ticketing.assign_ticket" and is_valid_json and assigned_team is None:
+            assigned_team = args.get("team")
+
         if action_type == "chat.post_message" and is_valid_json:
             ch = args.get("channel", "")
             if ch:
@@ -331,86 +342,49 @@ 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
+        # Extract full org_config from runbook response (one step after the call)
         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":
-            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())
+        if action_type == "meta.read_runbook" and last_result.get("ok"):
+            data = last_result.get("data") or {}
+            if isinstance(data, dict):
+                org = data.get("org_config") or {}
+                valid_labels     = set(org.get("label_taxonomy", {}).values())
+                valid_priorities = set(org.get("priority_levels", []))
+                valid_teams      = set(org.get("team_map", {}).values())
+                oncall_channels  = set(org.get("oncall_channels", {}).values())
 
         done = bool(getattr(result, "done", obs_dict.get("done", False)))
         final_score = float(getattr(result, "reward", obs_dict.get("reward", 0.0)))
         step += 1
 
-    # --- Auxiliary reward signals ---
-
-    valid_json_ratio = valid_action_count / max(step, 1)
-    efficiency = max(0.0, 1.0 - step / max_steps) if final_score >= 0.3 else 0.0
-
-    # no_wrong_channels: 1.0 baseline, -0.5 per wrong post
-    # Only meaningful if we know the oncall channels (runbook was read)
-    if oncall_channels:
-        wrong_posts = sum(1 for ch in posted_channels if ch not in oncall_channels)
-        no_wrong_channels = max(0.0, 1.0 - 0.5 * wrong_posts)
-    elif posted_channels:
-        # Posted without reading runbook — can't verify, treat all as wrong
-        no_wrong_channels = 0.0
-    else:
-        # No posts at all — no spray happened (reward_final_score handles missing notif)
-        no_wrong_channels = 1.0
-
-    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.
-    if valid_action_count == 0:
-        combined = -1.0
-    elif final_score == 0.0:
-        # 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 = (
-            ok_ratio              * 0.20
-            + read_runbook_reward * 0.10
-            - 0.30                        # hard penalty for zero task completion
-        )
-    else:
-        combined = (
-            final_score           * 0.40
-            + ok_ratio            * 0.10
-            + no_wrong_channels   * 0.15
-            + valid_json_ratio    * 0.10
-            + read_runbook_reward * 0.15
-            + efficiency          * 0.10
-        )
+    from training.rewards import compute_rollout_reward
+    combined = compute_rollout_reward(
+        read_runbook_done  = read_runbook_done,
+        valid_labels       = valid_labels,
+        valid_priorities   = valid_priorities,
+        valid_teams        = valid_teams,
+        oncall_channels    = oncall_channels,
+        ticket_label       = ticket_label,
+        ticket_priority    = ticket_priority,
+        assigned_team      = assigned_team,
+        posted_channels    = posted_channels,
+        env_score          = final_score,
+        valid_json_count   = valid_action_count,
+    )
 
     print(
-        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}"
+        f"[rollout] steps={step} env={final_score:.3f} "
+        f"label={'✓' if ticket_label and ticket_label in valid_labels else '✗' if ticket_label else '-'} "
+        f"priority={'✓' if ticket_priority and ticket_priority in valid_priorities else '✗' if ticket_priority else '-'} "
+        f"team={'✓' if assigned_team and assigned_team in valid_teams else '✗' if assigned_team else '-'} "
+        f"channel={'✓' if any(ch in oncall_channels for ch in posted_channels) else '✗' if posted_channels else '-'} "
+        f"→ reward={combined:.3f}"
     )
     return {
-        "prompt_ids": prompt_ids,
+        "prompt_ids":     prompt_ids,
         "completion_ids": completion_ids,
-        "logprobs": logprobs,
-        "final_score_reward": final_score,
-        "no_wrong_channels_reward": no_wrong_channels,
-        "valid_json_reward": valid_json_ratio,
-        "read_runbook_reward": read_runbook_reward,
-        "efficiency_reward": efficiency,
+        "logprobs":       logprobs,
+        "reward":         combined,
     }
 
 
@@ -429,14 +403,10 @@ def make_rollout_func(sync_env, tokenizer, max_steps: int = 15):
     """
     def rollout_func(prompts: list[str], trainer=None) -> dict:
         out: dict[str, list] = {
-            "prompt_ids": [],
+            "prompt_ids":     [],
             "completion_ids": [],
-            "logprobs": [],
-            "final_score_reward": [],
-            "no_wrong_channels_reward": [],
-            "valid_json_reward": [],
-            "read_runbook_reward": [],
-            "efficiency_reward": [],
+            "logprobs":       [],
+            "reward":         [],
         }
         # Track how many times each unique prompt has appeared so we can pass
         # a gen_offset — ensures repeated prompts (num_generations > 1) hit

training/train_v3.ipynb

@@ -489,14 +489,7 @@
    "id": "cell-13",
    "metadata": {},
    "outputs": [],
-   "source": [
-    "from training.dataset import generate_triage_dataset\n",
-    "\n",
-    "rows         = generate_triage_dataset(n_episodes=N_GRPO_EPISODES, base_seed=42)\n",
-    "grpo_dataset = Dataset.from_list([{'prompt': r['prompt']} for r in rows])\n",
-    "print(f'GRPO dataset: {len(grpo_dataset)} triage episodes')\n",
-    "print(f'Difficulties: {set(r[\"difficulty\"] for r in rows)}')"
-   ]
+   "source": "from training.dataset import generate_triage_dataset\n\n# generate_triage_dataset now pre-simulates the env's RNG to only include seeds\n# where env.reset(seed) will actually run a TRIAGE episode — previously the env\n# silently ran release_notes/dep_update for the same seed, guaranteeing env_score=0.\nrows         = generate_triage_dataset(n_episodes=N_GRPO_EPISODES, base_seed=42)\ngrpo_dataset = Dataset.from_list([{'prompt': r['prompt']} for r in rows])\nprint(f'GRPO dataset: {len(grpo_dataset)} triage episodes')\nprint(f'Difficulties: {set(r[\"difficulty\"] for r in rows)}')\nprint(f'Sample prompt: {rows[0][\"prompt\"][:120]}')"
   },
   {
    "cell_type": "markdown",
@@ -512,7 +505,7 @@
    "id": "cell-14",
    "metadata": {},
    "outputs": [],
-   "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}')"
+   "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\nfrom training.rewards import compute_rollout_reward\n\ngrpo_env = GenericEnvClient(base_url=ENV_URL).sync()\ngrpo_env.connect()\nprint('GRPO training env connected')\n\nROLLOUT_TEMPERATURE = 1.1   # must be > 1.0 for rollout diversity\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 triage episode and return trajectory + reward.\n\n    Reward is runbook-compliance based (see training/rewards.py):\n    - Did the model read the runbook?\n    - Did it use a valid label/priority/team from the runbook?\n    - Did it notify a correct oncall channel?\n\n    This creates genuine reward variance across rollouts because org configs\n    vary by seed — the same hardcoded label/team/channel is correct for some\n    orgs and wrong for others, giving GRPO a real gradient signal.\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    env_score        = 0.0\n    step, done       = 0, False\n    _sample_logged   = False\n\n    # Runbook-compliance tracking\n    read_runbook_done  = False\n    valid_labels:     set = set()\n    valid_priorities: set = set()\n    valid_teams:      set = set()\n    oncall_channels:  set = set()\n    ticket_label:    str | None = None\n    ticket_priority: str | None = None\n    assigned_team:   str | None = None\n    posted_channels: list = []\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_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' and parsed is not None:\n            read_runbook_done = True\n\n        if action_type == 'ticketing.create_ticket' and parsed is not None and ticket_label is None:\n            ticket_label    = args.get('label')\n            ticket_priority = args.get('priority')\n\n        if action_type == 'ticketing.assign_ticket' and parsed is not None and assigned_team is None:\n            assigned_team = args.get('team')\n\n        if action_type == 'chat.post_message' and parsed is not None:\n            ch = args.get('channel', '')\n            if ch:\n                posted_channels.append(ch)\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        # Extract org_config from runbook response (available the step AFTER read_runbook)\n        last_res = obs_dict.get('last_action_result') or {}\n        if action_type == 'meta.read_runbook' and last_res.get('ok'):\n            data = last_res.get('data') or {}\n            if isinstance(data, dict):\n                org = data.get('org_config') or {}\n                valid_labels     = set(org.get('label_taxonomy', {}).values())\n                valid_priorities = set(org.get('priority_levels', []))\n                valid_teams      = set(org.get('team_map', {}).values())\n                oncall_channels  = set(org.get('oncall_channels', {}).values())\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    reward = compute_rollout_reward(\n        read_runbook_done = read_runbook_done,\n        valid_labels      = valid_labels,\n        valid_priorities  = valid_priorities,\n        valid_teams       = valid_teams,\n        oncall_channels   = oncall_channels,\n        ticket_label      = ticket_label,\n        ticket_priority   = ticket_priority,\n        assigned_team     = assigned_team,\n        posted_channels   = posted_channels,\n        env_score         = env_score,\n        valid_json_count  = valid_json_count,\n    )\n\n    lbl_ok = '✓' if ticket_label and ticket_label in valid_labels else ('✗' if ticket_label else '-')\n    pri_ok = '✓' if ticket_priority and ticket_priority in valid_priorities else ('✗' if ticket_priority else '-')\n    tm_ok  = '✓' if assigned_team and assigned_team in valid_teams else ('✗' if assigned_team else '-')\n    ch_ok  = '✓' if any(ch in oncall_channels for ch in posted_channels) else ('✗' if posted_channels else '-')\n    print(f'  [rollout] steps={step} env={env_score:.3f} '\n          f'label={lbl_ok} pri={pri_ok} team={tm_ok} ch={ch_ok} '\n          f'offset={gen_offset} -> reward={reward:.3f}')\n\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    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, no reward key. kwargs: {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",