Update README.md

README.md

@@ -1,259 +1,317 @@
+# PM-Ops 🛠️
+
+> **Can a small model learn to operate inside a company it has never seen before?**
+
+PM-Ops is a reinforcement learning benchmark and training environment where an LLM agent operates as a product manager inside a fully simulated software organization — navigating ticketing, codebases, and chat — using only the organization's own runbook as its guide.
+
 ---
-title: PM-Ops RL Environment
-emoji: 📋
-colorFrom: blue
-colorTo: purple
-sdk: docker
-app_port: 8000
-pinned: false
-tags:
-  - openenv
-  - reinforcement-learning
+
+## Links
+
+| Resource | Link |
+|---|---|
+| 🤗 **Live Environment** | [TheCrustaceans/Pm-ops — HuggingFace Space](https://huggingface.co/spaces/TheCrustaceans/Pm-ops) |
+| 📝 **Blog Post** | [BlogPost.mdx](https://huggingface.co/spaces/TheCrustaceans/Pm-ops/blob/main/BlogPost.mdx) |
+|  **training** | [Google Drive](https://drive.google.com/file/d/1JdwYukKrEaMTaOwc1W4Q8bBRjxwGZ2be/view?usp=drive_link) |
+|  **evaluation** | [Google Drive](https://drive.google.com/file/d/1JdwYukKrEaMTaOwc1W4Q8bBRjxwGZ2be/view?usp=sharing) |
+
 ---
 
-# PM-Ops: Procedural PM Operations RL Environment
+## Table of Contents
+
+- [Motivation](#motivation)
+- [Environment](#environment)
+  - [The Three Apps](#the-three-apps)
+  - [Action Space](#action-space)
+  - [Observation Schema](#observation-schema)
+  - [Org Generator](#org-generator)
+  - [Task Types](#task-types)
+  - [Difficulty Tiers](#difficulty-tiers)
+- [Reward Design](#reward-design)
+- [Training Pipeline](#training-pipeline)
+  - [SFT Warmup](#sft-warmup)
+  - [GRPO](#grpo)
+  - [Reward Shaping](#reward-shaping)
+- [Results](#results)
+- [Project Structure](#project-structure)
+- [Setup & Usage](#setup--usage)
+
+---
 
-A simulated software-org environment — Ticketing, Chat, Codebase — designed to train LLM agents that learn *organizational conventions* rather than memorize fixed answers.
+## Motivation
 
-**Hosted env:** `https://huggingface.co/spaces/TheCrustaceans/Pm-ops`  
-**Trained checkpoint:** `Saurav1/pm-ops-grpo-Qwen3-1.7B-triage-v4`  
-**Stack:** Qwen3-1.7B · Unsloth · TRL 0.22.2 · SFT → GRPO
+Frontier models achieve near-perfect scores on standard benchmarks. Yet when deployed inside a real organization, they routinely fail tasks that a junior employee handles on day one.
+
+The reason: **organizational context is not on the internet.**
+
+A model knows that database failures are serious. It does not know that *your* company outsources DB infrastructure to a vendor and that the correct response is an email, not an internal incident ticket. No benchmark measures this gap. PM-Ops does.
+
+**Why system prompts are not the answer:**
+
+| Problem | Detail |
+|---|---|
+| Attention decay | As conversation grows, the model attends to early context instructions with diminishing weight. Org conventions buried in a long system prompt get effectively ignored. |
+| Static snapshots | Channels get renamed. Teams get reorganized. A system prompt written last quarter is already wrong. |
+| No gradient | The model is told what to do, not trained on the consequences of ignoring it. There is no learning signal. |
+
+PM-Ops addresses this by making the agent *experience* the consequences of skipping the runbook across thousands of varied organizational configurations during training.
 
 ---
 
-## The Core Idea
+## Environment
 
-Every episode samples a fresh `OrgConfig` — different label names, priority schemes, team-to-service mappings, and oncall channels. An agent that hardcodes `label="bug"` or `priority="P1"` will fail most episodes. A trained agent learns to **read the runbook first**, then use exactly what it says.
+PM-Ops runs as a WebSocket-based environment server (OpenEnv-compatible) hosted on HuggingFace Spaces. Each episode presents the agent with a freshly generated organization — different team names, channel names, label taxonomies, and priority levels every time.
+
+### The Three Apps
 
 ```
-Episode A                        Episode B
-──────────────────────────────   ──────────────────────────────
-labels:    issue, perf, docs     labels:    defect, vuln, guide
-priority:  critical/high/low     priority:  P0/P1/P2/P3
-team:      payments → data        team:      payments → backend
-channel:   #data-alerts           channel:   #backend-oncall
+┌─────────────────────────────────────────────────────────┐
+│                    PM-Ops Environment                    │
+│                                                         │
+│  ┌──────────────┐  ┌──────────────┐  ┌──────────────┐  │
+│  │  Ticketing   │  │   Codebase   │  │     Chat     │  │
+│  │  (Jira-like) │  │ (GitHub-like)│  │ (Slack-like) │  │
+│  │              │  │              │  │              │  │
+│  │ tickets      │  │ repositories │  │ channels     │  │
+│  │ projects     │  │ commits      │  │ threads      │  │
+│  │ teams        │  │ pull requests│  │ DMs          │  │
+│  │ labels       │  │ file authors │  │ user profiles│  │
+│  │ priorities   │  │ changed files│  │              │  │
+│  └──────────────┘  └──────────────┘  └──────────────┘  │
+│                                                         │
+│  Ground truth verified against DB — not LLM-judged      │
+└─────────────────────────────────────────────────────────┘
 ```
 
-The same hardcoded action scores 1.0 in Episode A and 0.0 in Episode B. This is the learning signal.
+**Ticketing App:** The source of truth for org state. Agent creates tickets, assigns teams, sets labels and priorities. Correctness is verified by checking the database, not asking another LLM.
 
----
+**Codebase App:** Pure detective work. No code is written. When a bug report arrives, the agent traces commit history, identifies the responsible change, and finds the author.
 
-## Tasks
+**Chat App:** Episode-specific channels. The channel `#oncall-payments` in episode 1 may be `#urgent-billing` in episode 2. The agent must read the runbook to know which one is live — not guess from training data.
 
-| Task | What the agent must do | Max Steps |
-|---|---|---|
-| **triage** | File a bug ticket with the correct org label + priority, assign to the owning team, page the oncall channel | 25 |
-| **incident_routing** | Identify the commit that broke a service, page the right oncall channel, optionally ticket it | 25 |
-| **release_notes** | Compile resolved tickets into release notes in the org's style (terse / verbose / bulleted) and post them | 40 |
-| **dep_update** | Coordinate a dependency update across all teams that own affected services | 40 |
+### Action Space
 
----
+```json
+{ "action_type": "meta.read_runbook",       "args": {} }
+{ "action_type": "meta.finish",             "args": {} }
+{ "action_type": "meta.noop",               "args": {} }
+{ "action_type": "ticketing.create_ticket", "args": { "summary": "...", "label": "...", "priority": "...", "assignee": "..." } }
+{ "action_type": "ticketing.assign_ticket", "args": { "ticket_id": "...", "team": "..." } }
+{ "action_type": "chat.list_channels",      "args": {} }
+{ "action_type": "chat.post_message",       "args": { "channel": "...", "text": "..." } }
+{ "action_type": "codebase.get_commits",    "args": { "repo": "..." } }
+```
 
-## Scoring
+All actions are emitted as chain-of-thought reasoning followed by a ` ```json ` block. Three-pass extraction handles malformed outputs: code block → raw JSON → regex fallback.
 
-Reward is **delayed** to episode end via a deterministic Python verifier — no LLM-as-judge.
+### Observation Schema
 
-### Triage (example)
-```
-+0.25  ticket created
-+0.20  correct label (must match label_taxonomy)
-+0.20  correct priority (must match priority_levels)
-+0.20  assigned to correct team (via team_map)
-+0.10  message posted to correct oncall channel
-+0.05  message body ≥ 20 chars
-─0.10  per duplicate ticket
-─0.05  per message posted to a noise channel
+Each step returns:
+
+```python
+{
+    "task_brief":             str,   # the incident/task description
+    "last_action_result":     dict,  # {"ok": bool, "data": ..., "error": ...}
+    "step":                   int,
+    "steps_remaining":        int,
+    "reward":                 float, # always 0.0 until meta.finish
+    "done":                   bool,
+    "token_budget_remaining": int,
+}
 ```
 
-### GRPO Training Reward (runbook-compliance)
+**Delayed reward:** `reward` is always `0.0` until the agent calls `meta.finish`. This forces the agent to commit to a plan and execute it — step-by-step reward hacking is not possible.
 
-During RL training we add a **dense shaped reward** on top of the env score:
+### Org Generator
 
-| Signal | Reward |
-|---|---|
-| Called `meta.read_runbook` | +0.10 |
-| Ticket label in `label_taxonomy` | +0.20 / −0.10 |
-| Ticket priority in `priority_levels` | +0.15 / −0.10 |
-| Team in `team_map` | +0.20 / −0.10 |
-| Posted to `oncall_channels` | +0.25 / −0.10 per wrong channel |
-| Env grader score | ×0.10 bonus |
+LLM-generated scenarios sound plausible but fail formal verification. PM-Ops uses a **deterministic, parameterized, seedable org generator** instead.
+
+Each org is generated from a seed and produces:
 
-Reward naturally varies per seed because every org has different valid values — the model can't memorize answers.
+- `label_taxonomy` — org-specific bug/feature labels
+- `priority_levels` — org-specific severity scale
+- `team_map` — service → owning team
+- `oncall_channels` — team → notification channel
+- `required_ticket_fields` — what fields must be set for a valid ticket
+
+The same seed always produces the same org. Evaluation is reproducible. Training sees a different org every episode.
+
+### Task Types
+
+| Task | Description | Key Challenge |
+|---|---|---|
+| **Triage** | Bug report arrives. Create ticket with correct label, priority, assignee. Notify correct channel. | Using org taxonomy, not generic labels. |
+| **Incident Routing** | Production incident. Identify affected services, find owning team via codebase + team map, escalate. | Cross-system reasoning: chat → codebase → ticketing. |
+| **Release Notes** | Compile and post release notes for resolved tickets in the org's specific format. | Format compliance, not just content. |
+| **Dependency Update** | Breaking library change hits multiple services. Notify each team through their own oncall channel, create per-team tickets. | Multi-target coordination without collapsing to single-service logic. |
+
+### Difficulty Tiers
+
+| Tier | Services | Labels | Priorities | Notes |
+|---|---|---|---|---|
+| **Easy** | 2 | 3 | 2 | No ambiguity. Runbook is complete. |
+| **Medium** | 3 | 4 | 4 | Partial runbook. Agent must infer across systems. Multi-owner configs. |
+| **Hard** | ≥5 | 5+ | 5+ | Noise channels designed to distract. Outdated/missing docs. Simulated panicking users in general chat — correct behavior is to ignore them. |
+
+The hard tier specifically tests resistance to NLP pressure. A model that responds to "THE SITE IS DOWN PLEASE HELP" in `#general` instead of reading the runbook and acting on verified info fails the episode.
 
 ---
 
-## Action Space
+## Reward Design
 
-The agent emits one JSON block per step:
+Rewards are **delayed** — revealed only at `meta.finish`.
 
-```json
-{"action_type": "ticketing.create_ticket", "args": {"summary": "...", "label": "issue", "priority": "high"}}
+### Scoring Breakdown
+
+```
+Ticket created correctly     +0.25
+Correct label                +0.20
+Correct priority             +0.20
+Correct team assignment      +0.20
+Correct channel notification +0.15
+─────────────────────────────────
+Maximum per episode           1.00
 ```
 
-| Namespace | Actions |
-|---|---|
-| `meta` | `read_runbook` · `finish` · `noop` |
-| `ticketing` | `create_ticket` · `update_ticket` · `get_ticket` · `list_tickets` · `assign_ticket` · `comment_ticket` · `transition_ticket` |
-| `chat` | `post_message` · `read_channel` · `list_channels` · `search` |
-| `codebase` | `list_commits` · `get_commit` · `list_prs` |
+### Penalties
 
-The optimal triage sequence takes **5 steps**:
 ```
-meta.read_runbook → ticketing.create_ticket → ticketing.assign_ticket → chat.post_message → meta.finish
+Wrong channel post           -0.05 each
+                             (>2 wrong posts → net negative; defeats channel-spray)
+Duplicate ticket             -0.10 each
+Zero valid actions (inaction)-1.00
 ```
 
----
-
-## Procedural World Generation
+The `-0.05` per wrong channel is calibrated: posting to every channel to guarantee hitting the right one becomes net negative after 2 wrong posts. The agent must read the runbook to know the correct channel.
 
-Each episode is fully seeded — same seed always produces the same org + scenario.
+### Combined Training Reward
 
 ```python
-# org_generator.py
-rng = random.Random(seed)
-difficulty = rng.choice(["easy", "medium", "medium", "hard"])
-label_taxonomy  = {k: rng.choice(variants[k]) for k in sampled_keys}
-priority_levels = rng.choice([["P0","P1","P2","P3"], ["critical","high","medium","low"], ...])
-team_map        = {service: team for service, team in zip(services, teams)}
-oncall_channels = {svc: f"#{team_map[svc]}{rng.choice(suffixes)}" for svc in services}
+if valid_action_count == 0:
+    combined = -1.0                              # never output valid JSON
+elif final_score == 0.0:
+    combined = valid_json_ratio * 0.15           # tried but failed
+            + read_runbook_reward * 0.10
+            - 0.30                               # hard penalty for zero completion
+else:
+    combined = final_score         * 0.45        # task correctness
+            + no_wrong_channels    * 0.15        # anti channel-spray
+            + valid_json_ratio     * 0.15        # format discipline
+            + read_runbook_reward  * 0.15        # process compliance
+            + efficiency           * 0.10        # steps saved
 ```
 
-**Difficulty** controls the number of services (2 / 3 / 5), labels (3 / 4 / 5), noise channels (1 / 2 / 3), and required ticket fields.
-
 ---
 
 ## Training Pipeline
 
-### Phase 1 — SFT Warmup (~15 min on A100)
+### SFT Warmup
 
-Run the deterministic `baseline_agent` for N episodes and collect every `(observation, action)` pair as a supervised example. Without this, the model outputs freeform text → all GRPO advantages are zero.
+Before GRPO, a short supervised fine-tuning phase runs on baseline agent demonstrations. Without this, the model outputs prose instead of structured JSON actions and the GRPO gradient is zero — the model needs to learn the output format before it can learn the task.
 
-```
-baseline_agent: read_runbook → create_ticket → assign_ticket → list_channels → post_message → finish
-SFT target: prompt + ```json\n{"action_type": ..., "args": {...}}\n``` + EOS
-```
+### GRPO
 
-After 2 epochs, the model reliably emits `\`\`\`json ... \`\`\`` blocks.
+We use **Group Relative Policy Optimization** to train a 7B parameter model directly against PM-Ops episode rewards.
 
-### Phase 2 — GRPO (~90 min on A100)
+Each training step is a full PM-Ops episode:
+1. Model plays through up to 15 turns
+2. Environment scores the final state
+3. Gradient updates weights based on relative performance across the generation group
 
-GRPO learns by comparing rewards across N generations of the same prompt. Each generation gets a different env seed (`gen_slot` offset), so they explore distinct org configs and produce varied rewards — giving GRPO a real gradient signal.
+No intermediate reward signal is given. The model must learn to plan across multiple steps.
 
-```
-prompt (same seed) → gen_slot=0 → org_A → reward=0.45
-                  → gen_slot=1 → org_B → reward=0.80
-                  → gen_slot=2 → org_C → reward=0.10
-advantage = reward - mean(group)  →  non-zero gradient
-```
+### Reward Shaping
+
+| Component | Weight | Purpose |
+|---|---|---|
+| `final_score` | 0.45 | Primary correctness from env grader |
+| `no_wrong_channels` | 0.15 | Anti-hack: penalise channel spray |
+| `valid_json_ratio` | 0.15 | Format discipline |
+| `read_runbook` | 0.15 | Process compliance |
+| `efficiency` | 0.10 | Steps saved (only when task succeeds) |
 
-**Critical: the dataset is pre-filtered.** `dataset.py` simulates the env's RNG to only include seeds where `env.reset(seed)` will run a *triage* task — eliminating the 75% of seeds that would produce release_notes or dep_update episodes and always score 0.
+`PMOpsGRPOTrainer` subclasses TRL's `GRPOTrainer` and overrides `_calculate_rewards()` to inject pre-computed rewards from the rollout directly, bypassing TRL's broken kwargs flow for multi-turn rollouts.
 
 ---
 
-## Repository Layout
+## Results
+
+Evaluated on 13 matched episodes, comparing the heuristic baseline agent against the GRPO-trained model:
+
+| Model | Avg Score (13 eps) |
+|---|---|
+| Heuristic Baseline | 0.269 |
+| GRPO Trained (early) | 0.362 |
+| **Improvement** | **+34.6%** |
+
+Training reward trend across 14 steps shows a positive slope of **+0.0024/step** with KL divergence remaining stable and controlled throughout.
+
+![Eval Results](eval_results.png)
+![Training Metrics](training_metrics.png)
+
+---
+
+## Project Structure
 
 ```
+pm_ops/
 ├── server/
-│   ├── pm_ops_environment.py    # OpenEnv Environment — reset/step/grade
-│   ├── apps/
-│   │   ├── ticketing.py         # Jira-like app (create, assign, transition tickets)
-│   │   ├── chat.py              # Slack-like app (post, read, search channels)
-│   │   └── codebase.py          # GitHub-like app (commits, PRs)
-│   ├── tasks/
-│   │   ├── triage_task.py       # Grader: label + priority + team + channel
-│   │   ├── incident_routing_task.py
-│   │   ├── release_notes_task.py
-│   │   └── dep_update_task.py
-│   └── world/
-│       ├── org_generator.py     # Procedural OrgConfig from seed
-│       └── scenario_gen.py      # Task brief + expected solution from seed
-│
-├── training/
-│   ├── train_v4.ipynb           # ← run this (current clean notebook)
-│   ├── dataset.py               # Pre-filtered triage dataset (seed-matched to env)
-│   ├── rollout.py               # Multi-turn episode runner + JSON extractor
-│   ├── rewards.py               # compute_rollout_reward (runbook-compliance)
-│   ├── prompts.py               # System prompt + observation formatter
-│   └── pm_ops_trainer.py        # PMOpsGRPOTrainer — reward cache injection
-│
-├── inference.py                 # Deterministic baseline_agent (heuristic)
-├── models.py                    # Pydantic types: PMOpsAction, PMOpsObservation
-├── client.py                    # OpenEnv client helpers
-└── Dockerfile                   # HF Space deployment
+│   ├── pm_ops_environment.py   # OpenEnv Environment subclass
+│   ├── org_generator.py        # Deterministic, seedable org factory
+│   ├── grader.py               # Ground-truth reward computation
+│   └── app.py                  # FastAPI server entry point
+├── inference.py                # Heuristic baseline agent + eval runner
+└── pyproject.toml
+
+training/
+├── rollout.py                  # Multi-turn rollout: build_messages, extract_json_action
+├── rewards.py                  # Reward functions + weight constants
+├── pm_ops_trainer.py           # PMOpsGRPOTrainer (_calculate_rewards override)
+├── prompts.py                  # SYSTEM_PROMPT, format_observation
+├── dataset.py                  # Dataset loading + seed parsing
+└── train_v3.ipynb              # SFT warmup + GRPO training notebook
 ```
 
 ---
 
-## Quick Start
+## Setup & Usage
 
-### Run the env locally
+### Run the environment server
 
 ```bash
+cd pm_ops
 pip install -e .
-uvicorn server.app:app --host 0.0.0.0 --port 8000
+server  # starts FastAPI on :7860
 ```
 
 ### Run the baseline agent
 
 ```bash
-python inference.py
-# [END] success=True steps=6 score=0.75 ...
+export API_BASE_URL=https://thecrustaceans-pm-ops.hf.space
+python pm_ops/inference.py
 ```
 
-### Connect a custom agent
+### Connect via OpenEnv client
 
 ```python
 from openenv.core import GenericEnvClient
 
-with GenericEnvClient(base_url="http://localhost:8000").sync() as env:
-    result   = env.reset(seed=42)
-    obs      = result.observation
-    result   = env.step({"action_type": "meta.read_runbook", "args": {}})
-    # ... agent loop ...
-    result   = env.step({"action_type": "meta.finish", "args": {}})
-    print(f"score = {result.reward:.3f}")
-```
-
-### Train (cloud — A100 or T4)
-
-Open `training/train_v4.ipynb` on your cloud runtime, run top to bottom. The notebook is self-contained: installs deps, clones the HF Space, runs SFT then GRPO, saves merged 16-bit weights to HF Hub.
-
-Key hyperparameters auto-adapt to GPU:
-
-| | A100 (40 GB) | T4 (16 GB) |
-|---|---|---|
-| `NUM_GEN` | 6 | 2 |
-| `GRAD_ACCUM` | 32 | 8 |
-| `N_SFT_EPISODES` | 120 | 50 |
-| `N_GRPO_EPISODES` | 150 | 30 |
-| `MAX_COMP_LEN` | 192 | 192 |
-
----
-
-## Key Design Decisions
+env = GenericEnvClient(base_url="https://thecrustaceans-pm-ops.hf.space").sync()
 
-**Why SFT before GRPO?**  
-GRPO computes `advantage = reward − mean(group)`. If the model outputs garbage on every step, all rewards are ~0 and all advantages are 0 → zero gradient. SFT teaches the JSON format first so GRPO has something to improve.
+with env:
+    result = env.reset()
+    obs = result.observation
 
-**Why runbook-compliance reward instead of env score?**  
-The env grader requires a precise 5-step sequence with *exact* field values (label, priority, team, channel). Early in training the model rarely completes the full sequence, so env score is 0 for nearly every episode — zero variance, zero GRPO signal. The compliance reward fires on individual correct actions and creates variance even from partial completions.
-
-**Why filter training seeds to triage-only?**  
-`PMOpsEnvironment.reset(seed)` draws task type uniformly from 4 options. Without filtering, 75% of training prompts describe a triage task but the env runs release_notes or dep_update → guaranteed env score = 0 regardless of what the model does. The dataset pre-simulates the env's RNG to only include seeds where `task_type == "triage"`.
+    result = env.step({
+        "action_type": "meta.read_runbook",
+        "args": {}
+    })
+```
 
-**Why `gen_slot` offsets?**  
-GRPO needs reward *variance across generations of the same prompt*. Without the offset, all N generations of the same prompt reset the env with the same seed → same org config → same correct answers → potentially same reward → zero advantage. `gen_slot=k` adds `k` to the seed so generation k explores a slightly different org.
+### Training
 
-**Why `PMOpsGRPOTrainer`?**  
-TRL's GRPOTrainer builds `reward_kwargs` from dataset columns only. Our multi-turn rollout reward (computed inside `rollout_func`) never reaches `_calculate_rewards` through the normal kwargs path. `PMOpsGRPOTrainer` wraps `rollout_func` to cache rewards and injects them directly as a tensor, bypassing TRL's broken plumbing.
+Open `training/train_v3.ipynb` in a Colab instance with a GPU. The notebook handles SFT warmup, GRPO setup, and checkpointing.
 
 ---
 
-## Sponsor Alignment
-
-| Sponsor | Relevance |
-|---|---|
-| **Meta RFC 004** | Delayed trajectory reward is the native reward mode |
-| **Patronus** | Schema drift — tests whether agents track per-org convention changes |
-| **Scale AI** | Enterprise workflow simulation with procedurally generated organizations |
-| **AI Labs** | Direct PM-in-the-loop positioning for agentic systems |
+**Live Environment:** [https://huggingface.co/spaces/TheCrustaceans/Pm-ops](https://huggingface.co/spaces/TheCrustaceans/Pm-ops)